05b2b1c02a8ce5f471877243ebebe5df876ccc37,models/unet/train_unet.py,,main,#Any#,194

Before Change

def main(args):
    if args.mode == "train":
        load_version = 0 if args.resume else None
        logger = TestTubeLogger(save_dir=args.exp_dir, name=args.exp, version=load_version)
        trainer = create_trainer(args, logger)
        model = UnetMRIModel(args)
        trainer.fit(model)
    else:  // args.mode == "test"
        assert args.checkpoint is not None
        model = UnetMRIModel.load_from_checkpoint(str(args.checkpoint))
        model.hparams.sample_rate = 1.

After Change

        trainer.test(model)

def main(args=None):
    parser = Args()
    parser.add_argument("--mode", choices=["train", "test", "challenge"], default="train")
    parser.add_argument("--num-epochs", type=int, default=50, help="Number of training epochs")
    parser.add_argument("--gpus", type=int, default=1)
    parser.add_argument("--nodes", type=int, default=1)
    parser.add_argument("--exp-dir", type=pathlib.Path, default="experiments",
                        help="Path where model and results should be saved")
    parser.add_argument("--exp", type=str, help="Name of the experiment")
    parser.add_argument("--checkpoint", type=pathlib.Path,
                        help="Path to pre-trained model. Use with --mode {test,challenge}")
    parser = UnetMRIModel.add_model_specific_args(parser)
    if args is not None:
        parser.set_defaults(**args)

    args, _ = parser.parse_known_args()
    random.seed(args.seed)
    np.random.seed(args.seed)
    torch.manual_seed(args.seed)
    run(args)

Instances