max_epochs=args.epochs,
sync_batchnorm=args.sync_batchnorm,
gradient_clip_val=args.gradient_clip_val,
callbacks=callbacks,
num_sanity_val_steps=0,
default_root_dir=args.results,
resume_from_checkpoint=ckpt_path,
accelerator="ddp" if args.gpus > 1 else None,
limit_train_batches=1.0 if args.train_batches == 0 else args.train_batches,
limit_val_batches=1.0 if args.test_batches == 0 else args.test_batches,
limit_test_batches=1.0 if args.test_batches == 0 else args.test_batches,
)
if args.benchmark:
if args.exec_mode == "train":
trainer.fit(model, train_dataloader=data_module.train_dataloader())
else:
# warmup
trainer.test(model, test_dataloaders=data_module.test_dataloader())
# benchmark run
trainer.current_epoch = 1
trainer.test(model, test_dataloaders=data_module.test_dataloader())
elif args.exec_mode == "train":
trainer.fit(model, data_module)
if is_main_process():
logname = args.logname if args.logname is not None else "train_log.json"
log(logname, torch.tensor(model.best_mean_dice), results=args.results)
elif args.exec_mode == "evaluate":
model.args = args
trainer.test(model, test_dataloaders=data_module.val_dataloader())
if is_main_process():
default_root_dir=args.results,
resume_from_checkpoint=ckpt_path,
accelerator="ddp" if args.gpus > 1 else None,
limit_train_batches=1.0
if args.train_batches == 0 else args.train_batches,
limit_val_batches=1.0 if args.test_batches == 0 else args.test_batches,
limit_test_batches=1.0
if args.test_batches == 0 else args.test_batches,
)
if args.benchmark:
if args.exec_mode == "train":
if args.profile:
with torch.autograd.profiler.emit_nvtx():
trainer.fit(
model, train_dataloader=data_module.train_dataloader())
else:
trainer.fit(model,
train_dataloader=data_module.train_dataloader())
else:
# warmup
trainer.test(model, test_dataloaders=data_module.test_dataloader())
# benchmark run
trainer.current_epoch = 1
trainer.test(model, test_dataloaders=data_module.test_dataloader())
elif args.exec_mode == "train":
trainer.fit(model, data_module)
elif args.exec_mode == "evaluate":
model.args = args
trainer.test(model, test_dataloaders=data_module.val_dataloader())
elif args.exec_mode == "predict":