from data_loading.data_module import DataModule from models.nn_unet import NNUnet from utils.gpu_affinity import set_affinity from utils.logger import LoggingCallback from utils.utils import get_main_args, is_main_process, log, make_empty_dir, set_cuda_devices, verify_ckpt_path if __name__ == "__main__": args = get_main_args() if args.affinity != "disabled": affinity = set_affinity(os.getenv("LOCAL_RANK", "0"), args.affinity) set_cuda_devices(args) seed_everything(args.seed) data_module = DataModule(args) data_module.prepare_data() data_module.setup() ckpt_path = verify_ckpt_path(args) callbacks = None model_ckpt = None if args.benchmark: model = NNUnet(args) batch_size = args.batch_size if args.exec_mode == "train" else args.val_batch_size log_dir = os.path.join(args.results, args.logname if args.logname is not None else "perf.json") callbacks = [ LoggingCallback( log_dir=log_dir, global_batch_size=batch_size * args.gpus, mode=args.exec_mode,
choices=["pre", "post"], help="Type of task to run; pre - localization, post - damage assesment", ) arg("--seed", type=int, default=1) parser = Model.add_model_specific_args(parser) args = parser.parse_args() if args.interpolate: args.deep_supervision = False args.dec_interp = False set_cuda_devices(args.gpus) affinity = set_affinity(os.getenv("LOCAL_RANK", "0"), "socket_unique_interleaved") seed_everything(args.seed) data_module = DataModule(args) callbacks = None checkpoint = args.ckpt if args.ckpt is not None and os.path.exists( args.ckpt) else None if args.exec_mode == "train": model = Model(args) model_ckpt = ModelCheckpoint(monitor="f1_score", mode="max", save_last=True) callbacks = [ EarlyStopping(monitor="f1_score", patience=args.patience, verbose=True, mode="max") ]