def main(): if opt.is_continue: model = torch.load('unet.pth').to(device) else: model = Unet(19).to(device) dataset = Cityscapes(opt.root, resize=opt.resize, crop=opt.crop) dataloader = DataLoader(dataset, batch_size=opt.batch_size, shuffle=True, num_workers=1) criterion = BCELoss().to(device) optimizer = Adam(model.parameters(), lr=0.001) t_now = time.time() for epoch in range(opt.n_epochs): print('epoch {}'.format(epoch)) for i, batch in enumerate(dataloader): inputs, labels = batch inputs = inputs.to(device) labels = labels.to(device) optimizer.zero_grad() outputs = model(inputs) loss = criterion(outputs, labels) loss.backward() optimizer.step() if i % 100 == 0: print(loss) print('time:', time.time() - t_now) t_now = time.time() print(loss) torch.save(model, 'unet.pth')
def main(): args = parser.parse_args() step = 0 exp_name = f'{args.name}_{hp.max_lr}_{hp.cycle_length}' transforms = segtrans.JointCompose([segtrans.Resize(400), segtrans.RandomRotate(0, 90), segtrans.RandomCrop(256, 256), segtrans.ToTensor(), segtrans.Normalize(mean=hp.mean, std=hp.std)]) val_transforms = segtrans.JointCompose([segtrans.PadToFactor(), segtrans.ToTensor(), segtrans.Normalize(mean=hp.mean, std=hp.std)]) train_dataset = DSBDataset(f'{args.data}/train', transforms=transforms) val_dataset = DSBDataset(f'{args.data}/val', transforms=val_transforms) model = Unet() if args.checkpoint: checkpoint = torch.load(args.checkpoint) model.load_state_dict(checkpoint['state']) step = checkpoint['step'] exp_name = checkpoint['exp_name'] optimizer = Adam(model.parameters(), lr=hp.max_lr) if args.find_lr: scheduler = LRFinderScheduler(optimizer) else: scheduler = SGDRScheduler(optimizer, min_lr=hp.min_lr, max_lr=hp.max_lr, cycle_length=hp.cycle_length, current_step=step) model.cuda(device=args.device) train(model, optimizer, scheduler, train_dataset, val_dataset, n_epochs=args.epochs, batch_size=args.batch_size, exp_name=exp_name, device=args.device, step=step)
def train(args): # device device = torch.device("cuda:%d" % args.gpu if torch.cuda.is_available() else "cpu") # data trainset = SonyDataset(args.input_dir, args.gt_dir, args.ps) train_loader = DataLoader(trainset, batch_size=args.batch_size, shuffle=True, num_workers=12, pin_memory=True) logging.info("data loading okay") # model model = Unet().to(device) # loss function criterion = nn.L1Loss() # optimizer optimizer = optim.Adam(model.parameters(), lr=args.lr, weight_decay=args.wd) # lr scheduler scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=2000, gamma=0.1) # training running_loss = 0.0 for epoch in range(args.num_epoch): scheduler.step() for i, databatch in enumerate(train_loader): # get the inputs input_patch, gt_patch, train_id, ratio = databatch input_patch, gt_patch = input_patch.to(device), gt_patch.to(device) # zero the parameter gradients optimizer.zero_grad() # forward + backward + optimize outputs = model(input_patch) loss = criterion(outputs, gt_patch) loss.backward() optimizer.step() # print statistics running_loss += loss.item() if i % args.log_interval == (args.log_interval - 1): print('[%d, %5d] loss: %.3f %s' % (epoch, i, running_loss / args.log_interval, datetime.now())) running_loss = 0.0 if epoch % args.save_freq == 0: if not os.path.isdir( os.path.join(args.result_dir, '%04d' % epoch)): os.makedirs(os.path.join(args.result_dir, '%04d' % epoch)) gt_patch = gt_patch.cpu().detach().numpy() outputs = outputs.cpu().detach().numpy() train_id = train_id.numpy() ratio = ratio.numpy() temp = np.concatenate( (gt_patch[0, :, :, :], outputs[0, :, :, :]), axis=2) scipy.misc.toimage( temp * 255, high=255, low=0, cmin=0, cmax=255).save(args.result_dir + '%04d/%05d_00_train_%d.jpg' % (epoch, train_id[0], ratio[0])) # at the end of epoch if epoch % args.model_save_freq == 0: torch.save(model.state_dict(), args.checkpoint_dir + './model_%d.pl' % epoch)
def main(args): logging.basicConfig(level=logging.DEBUG if args.verbose else logging.INFO) device = torch.device( 'cuda' if args.use_cuda and torch.cuda.is_available() else 'cpu') # Create output folder if (args.output_folder is not None): if not os.path.exists(args.output_folder): os.makedirs(args.output_folder) logging.debug('Creating folder `{0}`'.format(args.output_folder)) output_folder = os.path.join(args.output_folder, time.strftime('%Y-%m-%d_%H%M%S')) os.makedirs(output_folder) logging.debug('Creating folder `{0}`'.format(output_folder)) args.datafolder = os.path.abspath(args.datafolder) args.model_path = os.path.abspath( os.path.join(output_folder, 'model.th')) # Save the configuration in a config.json file with open(os.path.join(output_folder, 'config.json'), 'w') as f: json.dump(vars(args), f, indent=2) logging.info('Saving configuration file in `{0}`'.format( os.path.abspath(os.path.join(output_folder, 'config.json')))) # Get datasets and load into meta learning format meta_train_dataset, meta_val_dataset, _ = get_datasets( args.dataset, args.datafolder, args.num_ways, args.num_shots, args.num_shots_test, augment=augment, fold=args.fold, download=download_data) meta_train_dataloader = BatchMetaDataLoader(meta_train_dataset, batch_size=args.batch_size, shuffle=True, num_workers=args.num_workers, pin_memory=True) meta_val_dataloader = BatchMetaDataLoader(meta_val_dataset, batch_size=args.batch_size, shuffle=True, num_workers=args.num_workers, pin_memory=True) # Define model model = Unet(device=device, feature_scale=args.feature_scale) model = model.to(device) print(f'Using device: {device}') # Define optimizer meta_optimizer = torch.optim.Adam(model.parameters(), lr=args.meta_lr) #, weight_decay=1e-5) #meta_optimizer = torch.optim.RMSprop(model.parameters(), lr=learning_rate, momentum = 0.99) # Define meta learner metalearner = ModelAgnosticMetaLearning( model, meta_optimizer, first_order=args.first_order, num_adaptation_steps=args.num_adaption_steps, step_size=args.step_size, learn_step_size=False, loss_function=loss_function, device=device) best_value = None # Training loop epoch_desc = 'Epoch {{0: <{0}d}}'.format(1 + int(math.log10(args.num_epochs))) train_losses = [] val_losses = [] train_ious = [] train_accuracies = [] val_accuracies = [] val_ious = [] start_time = time.time() for epoch in range(args.num_epochs): print('start epoch ', epoch + 1) print('start train---------------------------------------------------') train_loss, train_accuracy, train_iou = metalearner.train( meta_train_dataloader, max_batches=args.num_batches, verbose=args.verbose, desc='Training', leave=False) print(f'\n train accuracy: {train_accuracy}, train loss: {train_loss}') print('end train---------------------------------------------------') train_losses.append(train_loss) train_accuracies.append(train_accuracy) train_ious.append(train_iou) # Evaluate in given intervals if epoch % args.val_step_size == 0: print( 'start evaluate-------------------------------------------------' ) results = metalearner.evaluate(meta_val_dataloader, max_batches=args.num_batches, verbose=args.verbose, desc=epoch_desc.format(epoch + 1), is_test=False) val_acc = results['accuracy'] val_loss = results['mean_outer_loss'] val_losses.append(val_loss) val_accuracies.append(val_acc) val_ious.append(results['iou']) print( f'\n validation accuracy: {val_acc}, validation loss: {val_loss}' ) print( 'end evaluate-------------------------------------------------' ) # Save best model if 'accuracies_after' in results: if (best_value is None) or (best_value < results['accuracies_after']): best_value = results['accuracies_after'] save_model = True elif (best_value is None) or (best_value > results['mean_outer_loss']): best_value = results['mean_outer_loss'] save_model = True else: save_model = False if save_model and (args.output_folder is not None): with open(args.model_path, 'wb') as f: torch.save(model.state_dict(), f) print('end epoch ', epoch + 1) elapsed_time = time.time() - start_time print('Finished after ', time.strftime('%H:%M:%S', time.gmtime(elapsed_time))) r = {} r['train_losses'] = train_losses r['train_accuracies'] = train_accuracies r['train_ious'] = train_ious r['val_losses'] = val_losses r['val_accuracies'] = val_accuracies r['val_ious'] = val_ious r['time'] = time.strftime('%H:%M:%S', time.gmtime(elapsed_time)) with open(os.path.join(output_folder, 'train_results.json'), 'w') as g: json.dump(r, g) logging.info('Saving results dict in `{0}`'.format( os.path.abspath(os.path.join(output_folder, 'train_results.json')))) # Plot results plot_errors(args.num_epochs, train_losses, val_losses, val_step_size=args.val_step_size, output_folder=output_folder, save=True, bce_dice_focal=bce_dice_focal) plot_accuracy(args.num_epochs, train_accuracies, val_accuracies, val_step_size=args.val_step_size, output_folder=output_folder, save=True) plot_iou(args.num_epochs, train_ious, val_ious, val_step_size=args.val_step_size, output_folder=output_folder, save=True) if hasattr(meta_train_dataset, 'close'): meta_train_dataset.close() meta_val_dataset.close()
valid_data = Ultrasound_Dataset(valid_df,transform=transforms_valid) valid_loader = DataLoader(valid_data , batch_size = 4,shuffle=False) # Checking GPU Avalaibility use_cuda = True if use_cuda and torch.cuda.is_available(): print('yes') print(torch.cuda.is_available()) # Model Initialization model = Unet(1,net_type='semi_inception',version='b',add_residual=True) if use_cuda and torch.cuda.is_available(): model.cuda() criterion = CustomLoss(0.5,1) optimizer = optim.Adam(model.parameters(),5e-6) scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, 'min', patience = 3) training_loss,valid_loss,model,saved_model=train_(model,optimizer,scheduler,criterion,train_loader,valid_loader,epochs=5) plot_learning_curve(training_loss,valid_loss) # save model for further use torch.save(model.state_dict(),'../Mymodel')
lr_schedule = 0.985 def adjust_lr(optimizer, current_lr, schedule): current_lr = current_lr * schedule for param_group in optimizer.param_groups: param_group['lr'] = current_lr return current_lr if __name__ == "__main__": args, unparsed = config.get_args() model = Unet(args) model = model.cuda() model.train() optimizer = torch.optim.Adam(model.parameters(), lr=lr, weight_decay=weight_decay) loss = MulticlassDiceLoss() train = get_file_list(brats_preprocessed_folder, train_ids_path) val = get_file_list(brats_preprocessed_folder, valid_ids_path) shapes = [brats_dataloader.load_patient(i)[0].shape[1:] for i in train] max_shape = np.max(shapes, 0) max_shape = list(np.max((max_shape, patch_size), 0)) dataloader_train = brats_dataloader(train, batch_size, max_shape, num_threads,