ryan/hypertower
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

added better memory caching, multithreaded processing, cleanup scripts dir

Browse Source
This commit is contained in:
rpotter6298
2026-03-13 08:08:36 +01:00
parent 8282461a23
commit 7ea85d5426
39 changed files with 1850 additions and 1998 deletions
Download Patch File Download Diff File Expand all files Collapse all files
classes/v2/loader_factory.py
+2
View File
@@ -206,6 +206,7 @@ def make_loader(
*,
image_transform,
image_preprocessor=None,
image_cache=None,
batch_size: int,
shuffle: bool,
num_workers: int,
@@ -216,6 +217,7 @@ def make_loader(
slots,
image_transform=image_transform,
image_preprocessor=image_preprocessor,
image_cache=image_cache,
)
return DataLoader(
ds,
classes/v2/slot_dataset.py
+61
View File
@@ -1,5 +1,6 @@
from __future__ import annotations
from concurrent.futures import ThreadPoolExecutor, as_completed
from typing import Any, Callable, Optional
from pathlib import Path
@@ -45,12 +46,14 @@ class SlotDataset(Dataset):
image_transform: Optional[Callable[[Image.Image], torch.Tensor]] = None,
matrix_transform: Optional[Callable[[Any], torch.Tensor]] = None,
image_preprocessor: Optional[Callable[..., Image.Image]] = None,
image_cache: Optional[dict[str, np.ndarray]] = None,
) -> None:
self.samples = samples
self.slot_descriptors = slot_descriptors
self.image_transform = image_transform or transforms.ToTensor()
self.matrix_transform = matrix_transform or self._default_matrix_transform
self.image_preprocessor = image_preprocessor
self.image_cache = image_cache
def __len__(self) -> int:
return len(self.samples)
@@ -74,14 +77,72 @@ class SlotDataset(Dataset):
raise ValueError("Missing required image slot")
return None
path = Path(value)
cache_key = str(value)
if self.image_cache is not None:
cached = self.image_cache.get(cache_key)
if cached is not None:
return self.image_transform(Image.fromarray(cached, mode="RGB"))
img = Image.open(path).convert("RGB")
if self.image_preprocessor is not None:
try:
img = self.image_preprocessor(img, path)
except TypeError:
img = self.image_preprocessor(img)
if self.image_cache is not None:
self.image_cache[cache_key] = np.asarray(img, dtype=np.uint8)
return self.image_transform(img)
def prebuild_image_cache(self, cache_workers: int = 0) -> None:
"""Pre-populate image_cache for all samples in this dataset."""
if self.image_cache is None:
return
paths = list({
str(record[key])
for record in self.samples
for key, desc in self.slot_descriptors.items()
if desc.kind == "image" and record.get(key) is not None
})
to_warm = [p for p in paths if p not in self.image_cache]
if not to_warm:
return
print(
f"[image_cache] warming {len(to_warm)} images "
f"({len(paths) - len(to_warm)} already cached)",
flush=True,
)
def _warm_one(path_str: str) -> None:
if path_str in self.image_cache:
return
p = Path(path_str)
img = Image.open(p).convert("RGB")
if self.image_preprocessor is not None:
try:
img = self.image_preprocessor(img, p)
except TypeError:
img = self.image_preprocessor(img)
self.image_cache[path_str] = np.asarray(img, dtype=np.uint8)
try:
from tqdm import tqdm
except ImportError:
tqdm = None
if cache_workers <= 1:
it = tqdm(to_warm, desc="Warm image cache", unit="img") if tqdm else to_warm
for path_str in it:
_warm_one(path_str)
else:
with ThreadPoolExecutor(max_workers=cache_workers) as ex:
futures = {ex.submit(_warm_one, p): p for p in to_warm}
it = tqdm(as_completed(futures), total=len(futures), desc="Warm image cache", unit="img") if tqdm else as_completed(futures)
for fut in it:
fut.result()
def _load_matrix(self, value: Any, *, required: bool) -> Optional[torch.Tensor]:
if value is None:
if required:
classes/v2/v2_hypertower.py
+28 -5
View File
@@ -222,7 +222,13 @@ class V2HyperTower:
ap.add_argument("--augment", action="store_true")
ap.add_argument("--balanced-sampling", action="store_true",
help="Use WeightedRandomSampler during training to equalise class frequency (default: off).")
ap.add_argument("--num-workers", type=int, default=0)
ap.add_argument("--num-workers", type=int, default=4)
ap.add_argument("--in-memory-cache", action="store_true", default=True,
help="Cache preprocessed images in RAM (default: on).")
ap.add_argument("--no-in-memory-cache", action="store_false", dest="in_memory_cache",
help="Disable in-memory image cache.")
ap.add_argument("--cache-workers", type=int, default=4,
help="Threads for prebuilding in-memory image cache (default: 4).")
ap.add_argument("--device", choices=["auto", "cpu", "cuda"], default="auto")
ap.add_argument("--seed", type=int, default=1234)
ap.add_argument("--run-name", default=None)
@@ -452,6 +458,8 @@ class V2HyperTower:
n_classes=num_classes,
)
image_cache: dict | None = {} if getattr(args, "in_memory_cache", False) else None
for fold in range(n_folds):
seed_everything(args.seed + fold * 100)
fold_dir = tm_dir / f"fold{fold}"
@@ -469,6 +477,7 @@ class V2HyperTower:
fold_dir=fold_dir,
tower_mode=tower_mode,
pred_store=pred_store,
image_cache=image_cache,
)
fold_results.append(result)
if artifacts.y_true_ensemble is not None:
@@ -632,6 +641,7 @@ class V2HyperTower:
fold_dir: Path,
tower_mode: str,
pred_store: "PredictionStore | None" = None,
image_cache: "dict | None" = None,
) -> tuple[FoldResult, FoldArtifacts]:
args = self.args
device = self.device
@@ -743,7 +753,8 @@ class V2HyperTower:
slots_eye = profile_eye.slot_descriptors()
slots_patient = profile_patient.slot_descriptors()
loader_kw = dict(batch_size=args.batch_size, num_workers=args.num_workers)
loader_kw = dict(batch_size=args.batch_size, num_workers=args.num_workers,
image_cache=image_cache)
# ---- loaders ---------------------------------------------------
use_balanced = bool(getattr(args, "balanced_sampling", False))
@@ -831,6 +842,16 @@ class V2HyperTower:
if pred_store is not None:
pred_store.set_split(fold, [str(s["id_1"]) for s in holdout_bilat], "holdout")
# ---- prebuild in-memory image cache (fold 0 only; shared dict fills for later folds) ----
if image_cache is not None:
cache_workers = int(getattr(args, "cache_workers", 4))
_loaders_to_warm = [
train_single_loader, train_bilat_loader, val_loader, holdout_loader,
]
for _ldr in _loaders_to_warm:
if _ldr is not None:
_ldr.dataset.prebuild_image_cache(cache_workers=cache_workers)
opt_single = torch.optim.Adam(single.parameters(), lr=args.lr) if run_single else None
opt_bilateral = torch.optim.Adam(bilateral.parameters(), lr=args.lr) if run_bilat else None
@@ -942,6 +963,7 @@ class V2HyperTower:
# ---- epoch loop ------------------------------------------------
_prev_phase_single = "inactive" # used to detect md_warmup → next phase transition
for epoch in range(total_epochs):
_epoch_t0 = time.time()
if not run_single:
phase_single, main_epoch_single, single_active = "inactive", 0, False
elif epoch < single_warmup_md:
@@ -1329,6 +1351,7 @@ class V2HyperTower:
print() # seal the progress bar line
if args.log_every > 0 and (epoch + 1) % args.log_every == 0:
_epoch_secs = time.time() - _epoch_t0
hld_auc = target_holdout_single_auc if run_single else bi_auc_h
hld_suffix = f" hld_auc={hld_auc:.4f}" if holdout_loader is not None else ""
@@ -1356,7 +1379,7 @@ class V2HyperTower:
if run_single:
if tower_mode == "single":
msg = (
f" ep {epoch+1:>3}/{total_epochs} "
f" ep {epoch+1:>3}/{total_epochs} ({_epoch_secs:.1f}s) "
f"[single:{phase_single} {single_phase_epoch}/{single_phase_total}] "
f"fused(acc={cl_acc:.4f},auc={cl_auc:.4f}) "
f"img(acc={cl_acc_img:.4f},auc={cl_auc_img:.4f}) "
@@ -1366,7 +1389,7 @@ class V2HyperTower:
)
else:
msg = (
f" ep {epoch+1:>3}/{total_epochs} "
f" ep {epoch+1:>3}/{total_epochs} ({_epoch_secs:.1f}s) "
f"[single:{phase_single} {single_phase_epoch}/{single_phase_total}] "
f"fused(acc={en_acc:.4f},auc={en_auc:.4f}) "
f"img(acc={en_acc_img:.4f},auc={en_auc_img:.4f}) "
@@ -1376,7 +1399,7 @@ class V2HyperTower:
)
else:
msg = (
f" ep {epoch+1:>3}/{total_epochs} "
f" ep {epoch+1:>3}/{total_epochs} ({_epoch_secs:.1f}s) "
f"[bilat:{phase_bilat} {bilat_phase_epoch}/{bilat_phase_total}] "
f"fused(acc={bi_acc:.4f},auc={bi_auc:.4f}) "
f"img(acc={bi_acc_img:.4f},auc={bi_auc_img:.4f}) "