def main(): args = get_arguments() os.environ["CUDA_VISIBLE_DEVICES"] = "0" ## FOR REPRODUCIBILITY OF RESULTS seed = 1777777 utils.reproducibility(args, seed) utils.make_dirs(args.save) utils.save_arguments(args, args.save) training_generator, val_generator, full_volume, affine = medical_loaders.generate_datasets( args, path='.././datasets') model, optimizer = medzoo.create_model(args) criterion = create_loss('CrossEntropyLoss') criterion = DiceLoss(classes=args.classes, weight=torch.tensor([0.1, 1, 1, 1]).cuda()) if args.cuda: model = model.cuda() print("Model transferred in GPU.....") trainer = Trainer(args, model, criterion, optimizer, train_data_loader=training_generator, valid_data_loader=val_generator, lr_scheduler=None) print("START TRAINING...") trainer.training()
def train(): # args = iseg2019_arguments() print(args) utils.reproducibility(args, seed) utils.make_dirs(args.save) ( training_generator, val_generator, full_volume, affine, ) = medical_loaders.generate_datasets(args) model, optimizer = medzoo.create_model(args) val_criterion = DiceLoss(classes=args.classes) criterion = torch.nn.CrossEntropyLoss() if args.cuda: model = model.cuda() print("Model transferred in GPU.....") trainer = train.Trainer( args, model, criterion, optimizer, val_criterion=val_criterion, train_data_loader=training_generator, valid_data_loader=val_generator, lr_scheduler=None, ) print("START TRAINING...") trainer.training()
def main(): args = get_arguments() #os.environ["CUDA_VISIBLE_DEVICES"] = "0" ## FOR REPRODUCIBILITY OF RESULTS seed = 1777777 utils.reproducibility(args, seed) training_generator, val_generator, full_volume, affine = medical_loaders.generate_datasets( args, path='./datasets') model, optimizer = medzoo.create_model(args) # criterion = DiceLoss(classes=args.classes) # # ## TODO LOAD PRETRAINED MODEL print(affine.shape) model.restore_checkpoint(args.pretrained) if args.cuda: model = model.cuda() full_volume = full_volume.cuda() print("Model transferred in GPU.....") x = torch.randn(3, 156, 240, 240).cuda() print(full_volume.shape) output = non_overlap_padding(args, full_volume, model, criterion, kernel_dim=(32, 32, 32))
def main(): args = get_arguments() utils.reproducibility(args, seed) utils.make_dirs(args.save) training_generator, val_generator, full_volume, affine = medical_loaders.generate_datasets( args, path='.././datasets') model, optimizer = medzoo.create_model(args) criterion = DiceLoss(classes=args.classes) if args.cuda: model = model.cuda() print("Model transferred in GPU.....") trainer = train.Trainer(args, model, criterion, optimizer, train_data_loader=training_generator, valid_data_loader=val_generator, lr_scheduler=None) print("START TRAINING...") trainer.training() visualize_3D_no_overlap_new(args, full_volume, affine, model, 10, args.dim)
def main(): args = get_arguments() utils.reproducibility(args, seed) utils.make_dirs(args.save) training_generator, val_generator, full_volume, affine = medical_loaders.generate_datasets(args, path='.././datasets') model, optimizer = medzoo.create_model(args) criterion = DiceLoss(classes=args.classes) if args.cuda: model = model.cuda() trainer = train.Trainer(args, model, criterion, optimizer, train_data_loader=training_generator, valid_data_loader=val_generator) trainer.training()
def main(): args = get_arguments() utils.reproducibility(args, seed) utils.make_dirs(args.save) training_generator, val_generator, full_volume, affine = medical_loaders.generate_datasets( args, path='.././datasets') model, optimizer = medzoo.create_model(args) criterion = DiceLoss( classes=args.classes ) # ,skip_index_after=2,weight=torch.tensor([0.00001,1,1,1]).cuda()) if args.cuda: model = model.cuda() print("Model transferred in GPU.....") trainer = train.Trainer(args, model, criterion, optimizer, train_data_loader=training_generator, valid_data_loader=val_generator) print("START TRAINING...") trainer.training()
def test(): # args = iseg2019_arguments() print(args) utils.reproducibility(args, seed) utils.make_dirs(args.save) params = {"batch_size": args.batchSz, "shuffle": True, "num_workers": 2} print(params) samples_train = args.samples_train samples_val = args.samples_val test_loader = MRIDatasetISEG2019( args, "test", dataset_path=dataset_dir, crop_dim=args.dim, split_id=0, samples=samples_train, load=args.loadData, ) model_name = args.model lr = args.lr in_channels = args.inChannels num_classes = args.classes weight_decay = 0.0000000001 print("Building Model . . . . . . . ." + model_name) model = UNet3D(in_channels=in_channels, n_classes=num_classes, base_n_filter=8) print( model_name, "Number of params: {}".format( sum([p.data.nelement() for p in model.parameters()]) ), ) model.restore_checkpoint( "/home/kyle/results/UNET3D/iseg2019_9_06-08_21-25/iseg2019_9_06-08_21-25_BEST.pth" ) criterion = DiceLoss(classes=args.classes) # model = model.cuda() # print("Model transferred in GPU.....") print("TESTING...") model.eval() confusion_matrix = [[0] * 4 for i in range(4)] for batch_idx, input_tuple in enumerate(test_loader): with torch.no_grad(): img_t1, img_t2, target = input_tuple target = torch.reshape(target, (-1, 1, 64, 64, 64)) img_t1 = torch.reshape(img_t1, (-1, 1, 64, 64, 64)) img_t2 = torch.reshape(img_t2, (-1, 1, 64, 64, 64)) input_tensor = torch.cat((img_t1, img_t2), dim=1) # print(input_tensor.size()) input_tensor.requires_grad = False output = model(input_tensor) output = torch.argmax(output, dim=1) output = torch.reshape(output, (-1, 1, 64, 64, 64)) assert target.size() == output.size() output = torch.reshape(output, (-1,)).tolist() target = torch.reshape(target, (-1,)).tolist() assert len(output) == len(target) for gt, pred in zip(target, output): confusion_matrix[int(gt)][int(pred)] += 1 pprint(confusion_matrix)