p.requires_grad = False model = DataParallel(model).cuda() if NUM_GPUS > 1 else model.cuda() optimizer = Optimizer.from_params( [x for x in model.named_parameters() if x[1].requires_grad], params['trainer']['optimizer']) lr_scheduler_params = params['trainer'].pop("learning_rate_scheduler", None) scheduler = LearningRateScheduler.from_params( optimizer, lr_scheduler_params) if lr_scheduler_params else None if os.path.exists(args.folder): print("Found folder! restoring", flush=True) start_epoch, val_metric_per_epoch = restore_checkpoint( model, optimizer, serialization_dir=args.folder, learning_rate_scheduler=scheduler) else: print("Making directories") os.makedirs(args.folder, exist_ok=True) start_epoch, val_metric_per_epoch = 0, [] shutil.copy2(args.params, args.folder) param_shapes = print_para(model) num_batches = 0 for epoch_num in range(start_epoch, params['trainer']['num_epochs'] + start_epoch): train_results = [] norms = [] model.train()
def restore_checkpoint(self, serialization_dir, epoch_to_load): # Restore from a training dir return restore_checkpoint(self.model, self.optimizer, serialization_dir, epoch_to_load)