Furen Xiao
8c813db209
Update README.md to include detailed project overview, environment requirements, data sources, and a comprehensive directory layout. Add nnU-Net pipeline documentation and directory descriptions. Update .gitignore to exclude nnU-Net specific directories and add new scripts directory for nnU-Net pipeline. Add initial nnU-Net pipeline scripts.
46 lines
No EOL
1.9 KiB
Python
46 lines
No EOL
1.9 KiB
Python
import os
|
|
import torch
|
|
|
|
from nnunetv2.training.nnUNetTrainer.nnUNetTrainer import nnUNetTrainer
|
|
|
|
|
|
class NTUHLPLTrainer(nnUNetTrainer):
|
|
"""nnU-Net trainer for iterative pseudo-labeling on NTUH T1c tumor segmentation.
|
|
|
|
Env knobs (read in __init__, applied before initialize() builds optimizer/LR):
|
|
NNU_PL_EPOCHS total epochs of this run (default 250)
|
|
NNU_PL_LR initial lr, PolyLR decay (default 1e-2)
|
|
|
|
Warm start (round-to-round fine-tuning):
|
|
NNU_PL_WARMSTART path to an nnU-Net checkpoint. Its FULL network weights,
|
|
segmentation head included, are loaded in on_train_start(). The CLI
|
|
-pretrained_weights flag deliberately skips `.seg_layers.` keys, which is
|
|
wrong for iterative pseudo-labeling, hence this hook.
|
|
"""
|
|
|
|
def __init__(self, plans: dict, configuration: str, fold: int, dataset_json: dict,
|
|
device: torch.device = torch.device("cuda")):
|
|
super().__init__(plans, configuration, fold, dataset_json, device)
|
|
self.num_epochs = int(float(os.environ.get("NNU_PL_EPOCHS", "250")))
|
|
self.initial_lr = float(os.environ.get("NNU_PL_LR", "1e-2"))
|
|
self._pl_warmstart_done = False
|
|
|
|
def on_train_start(self):
|
|
super().on_train_start()
|
|
if self._pl_warmstart_done:
|
|
return
|
|
self._pl_warmstart_done = True
|
|
ws = os.environ.get("NNU_PL_WARMSTART")
|
|
if not ws:
|
|
return
|
|
mod = self.network
|
|
if hasattr(mod, "module"):
|
|
mod = mod.module
|
|
if hasattr(mod, "_orig_mod"):
|
|
mod = mod._orig_mod
|
|
ckpt = torch.load(ws, map_location=self.device, weights_only=False)
|
|
w = ckpt["network_weights"]
|
|
mod.load_state_dict(w, strict=True)
|
|
torch.cuda.empty_cache()
|
|
print(f"[NTUHLPLTrainer] warm-started full network weights from {ws}", flush=True)
|
|
self.print_to_log_file(f"warm-started full network weights from {ws}") |