def main(): # patches_file = "50_patches_dataset.h5" # hf = h5py.File(patches_file, 'r') # We obtain a list with all the IDs of the patches # all_groups = list(hf) # Dividing the dataset into train and validation. Shuffle has to be false otherwise the model might be trained # on what was previously validation set and validated on what was previously train set. # X_train, X_validation = train_test_split(all_groups, test_size=0.2, shuffle=False) # print(X_train, X_validation) # for testing datapath = "Data/" train_file = datapath + "patches_dataset_test.h5" val_file = datapath + "val250.h5" # Loader Parameters params = {'batch_size': 2, 'shuffle': False, 'num_workers': 0} train_dataset = PatchDataset(train_file, n_classes=3) print(len(train_dataset)) val_dataset = PatchDataset(val_file, n_classes=3) print(len(val_dataset)) train_loader = DataLoader(train_dataset, **params) val_loader = DataLoader(val_dataset, **params) loaders = {'train': train_loader, 'val': val_loader} # Model and param model = Modified3DUNet(in_channels=1, n_classes=3) optimizer = optim.Adam(model.parameters()) max_epochs = 10 # Median foreground percentage = 0.2 (= class 1,2) # Median cancer percentage = 0.01 (= class 2) # Median pancreas percentage = 0.2 - 0.01 = 0.19 (= class 1) # Median background percentage = 1-0.2 = 99.8 (=class 0) # [99.8, 0.19, 0.01] => corresponding class weights = [1, 525, 9980] # class_weights = torch.tensor([1., 525., 9980.]) # loss_criterion = GeneralizedDiceLoss(weight=class_weights) # loss_criterion = WeightedCrossEntropyLoss(weight=class_weights) weights = [1, 100, 500] class_weights = torch.FloatTensor(weights) loss_criterion = nn.CrossEntropyLoss(weight=class_weights) # trainer = UNetTrainer(model, optimizer, loaders, max_epochs, loss_criterion=loss_criterion) # trainer.train() # Load from last epoch checkpoint_trainer = UNetTrainer.load_checkpoint( "WCEL_1_10_50_last_model", model, optimizer, loaders, max_epochs, loss_criterion=loss_criterion) pred = checkpoint_trainer.single_image_forward(val_dataset[0][0])
def main(): # setup environments and seeds os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu torch.manual_seed(args.seed) torch.cuda.manual_seed(args.seed) random.seed(args.seed) np.random.seed(args.seed) # setup networks #Network = getattr(models, args.net) #model = Network(**args.net_params) model = Modified3DUNet(in_channels=1, n_classes=2, base_n_filter=16) model = model.cuda() '''optimizer = getattr(torch.optim, args.opt)( model.parameters(), **args.opt_params)''' optimizer = torch.optim.Adam(model.parameters(), lr=0.001, weight_decay=0.0001) #optimizer = torch.optim.SGD(model.parameters(),lr = 0.1,momentum=0.9) criterion = getattr(criterions, args.criterion) msg = '-------------- New training session -----------------' msg += '\n' + str(args) logging.info(msg) num_gpus = len(args.gpu.split(',')) args.batch_size *= num_gpus args.workers *= num_gpus args.opt_params['lr'] *= num_gpus # create dataloaders #Dataset = getattr(datasets, args.dataset) dset = cell_training('/home/tom/Modified-3D-UNet-Pytorch/PNAS/') train_loader = DataLoader(dset, batch_size=args.batch_size, shuffle=True, num_workers=args.workers, pin_memory=True) file_name_best = os.path.join(ckpts, 'cell/model_best.tar') for epoch in range(args.start_epoch, args.epochs): adjust_learning_rate(optimizer, epoch) # train for one epoch train_loss = train(train_loader, model, criterion, optimizer, epoch) # remember best lost and save checkpoint ckpt = { 'epoch': epoch + 1, 'state_dict': model.state_dict(), 'optim_dict': optimizer.state_dict(), 'train_loss': train_loss, } file_name = os.path.join(ckpts, 'model_last.tar') torch.save(ckpt, file_name) msg = 'Epoch: {:02d} Train loss {:.4f}'.format(epoch + 1, train_loss) logging.info(msg)