Furen Xiao
a516408c4e
- Implemented 09_run_iterative.py for MONAI to orchestrate iterative pseudo-labeling rounds, including training, predicting, and evaluation. - Created 01_prepare_dataset.py, 02_plan_preprocess.py, 06_pseudo_label.py, 07_train.py, 08_eval.py, and 09_run_iterative.py for nnU-Net to manage dataset preparation, planning, training, evaluation, and iterative pseudo-labeling. - Enhanced error handling and logging throughout the scripts for better traceability. - Added functionality to handle negative pseudo-labels and TTA options during inference. - Included visualization of Dice scores across rounds for both MONAI and nnU-Net.
32 lines
No EOL
1.3 KiB
Python
32 lines
No EOL
1.3 KiB
Python
"""Train the nnU-Net model for one pseudo-labeling round.
|
|
|
|
Epochs / initial lr / warm-start checkpoint are passed to NTUHLPLTrainer via
|
|
NNU_PL_* env vars (see scripts_nnu/trainers/ntuh_pl_trainer.py).
|
|
|
|
Usage:
|
|
nnUNetv2_train wrapper: python scripts_nnu/03_nnu_train.py \
|
|
--gpus 3 --epochs 250 --lr 1e-2 [--warmstart runs_nnu/round0/best_nnu.pth]
|
|
"""
|
|
import argparse
|
|
import os
|
|
import sys
|
|
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
|
|
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
|
|
from nnu_common import run, nnu_env, train_cmd, best_ckpt
|
|
|
|
|
|
def main():
|
|
ap = argparse.ArgumentParser()
|
|
ap.add_argument("--gpus", type=int, default=3)
|
|
ap.add_argument("--epochs", type=int, default=250)
|
|
ap.add_argument("--lr", type=float, default=1e-2)
|
|
ap.add_argument("--warmstart", default=None, help="checkpoint for full-weight warm start")
|
|
ap.add_argument("--log", default=None)
|
|
args = ap.parse_args()
|
|
log = args.log or os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "logs", "nnu_train.log")
|
|
run(train_cmd(args.gpus), log, env=nnu_env(epoch=args.epochs, lr=args.lr, warmstart=args.warmstart))
|
|
print(f"[nnu:train] done; best checkpoint at {best_ckpt()}", flush=True)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main() |