val_loader = Loader(val_path, arg.batch_size, transform=preprocess, sampler=arg.sampler, torch_type=arg.dtype, cpus=arg.cpus, shuffle=False, drop_last=False) test_loader = Loader(test_path, 1, torch_type=arg.dtype, cpus=arg.cpus, shuffle=False, drop_last=False) norm_layer = nn.BatchNorm2d act = nn.ReLU if arg.model == "unet": net = Unet2D(feature_scale=arg.feature_scale, act=act) elif arg.model == "unetres": net = UnetRes2D(1, nn.InstanceNorm2d, is_pool=arg.pool) net = nn.DataParallel(net).to(torch_device) recon_loss = nn.BCEWithLogitsLoss() model = CNNTrainer(arg, net, torch_device, recon_loss=recon_loss) if arg.test is False: model.train(train_loader, val_loader) model.test(test_loader, val_loader)
if arg.model == "fusion": net = Fusionnet(arg.in_channel, arg.out_channel, arg.ngf, arg.clamp) elif arg.model == "unet": net = Unet2D(feature_scale=arg.feature_scale) elif arg.model == "unet_sh": net = UnetSH2D(arg.sh_size, feature_scale=arg.feature_scale) else: raise NotImplementedError("Not Implemented Model") net = nn.DataParallel(net).to(torch_device) if arg.loss == "l2": recon_loss = nn.L2Loss() elif arg.loss == "l1": recon_loss = nn.L1Loss() recon_loss = EdgeWeightedLoss(1, 10) #recon_loss=nn.BCELoss() model = CNNTrainer(arg, net, torch_device, recon_loss=recon_loss, logger=logger) model.load(filename="epoch[0123]_losssum[0.005018].pth.tar") if arg.test == 0: model.train(train_loader, test_loader) if arg.test == 1: model.test(test_loader) # utils.slack_alarm("zsef123", "Model %s Done"%(arg.save_dir))
shuffle=False, drop_last=False) test_loader = Loader(test_path, 1, torch_type=arg.dtype, cpus=arg.cpus, shuffle=False, drop_last=False) norm_layer = nn.BatchNorm2d act = nn.ReLU if arg.model == "unet": net = Unet2D(feature_scale=arg.feature_scale, act=act) elif arg.model == "unetres": net = UnetRes2D(1, nn.InstanceNorm2d, is_pool=arg.pool) elif arg.model == "unetbr": net = UnetBR2D(1, nn.InstanceNorm2d, is_pool=arg.pool) elif arg.model == "vnet": net = VNet(elu=False) net = nn.DataParallel(net).to(torch_device) recon_loss = nn.BCEWithLogitsLoss() if arg.model == "vnet": model = VNetTrainer(arg, net, torch_device, recon_loss=recon_loss) else: model = CNNTrainer(arg, net, torch_device, recon_loss=recon_loss) if arg.test is False: model.train(train_loader, val_loader) model.test(test_loader, val_loader)
valid_path, arg.batch_size, image_type=arg.output, cpus=arg.cpus, ) train_loader = MARLoader( train_path, arg.batch_size, image_type=arg.output, cpus=arg.cpus, ) if arg.model == "fusion": fusionnet = Fusionnet(1, 1, arg.ngf, arg.output, arg.clamp).cuda() model = CNNTrainer(arg, fusionnet, recon_loss=nn.L1Loss(), mse_loss=nn.MSELoss()) elif arg.model == "unet": # for example # unet = Unet(1, 1, arg.ngf, arg.output, arg.clamp).cuda() # model = CNNTrainer(arg, Unet, recon_loss=nn.L1Loss(), mse_loss=nn.MSELoss()) raise NotImplementedError("Not Implemented Unet") elif arg.model == "pix2pix": # G = ...Net # D = ...Discriminator # model = GANTrainer(...) raise NotImplementedError("Not Implemented Model") else: raise NotImplementedError("Not Implemented Model")
elif arg.model == "unet_nonlocal": net = Unet_NonLocal2D(feature_scale=arg.feature_scale) elif arg.model == "unet_gridatt": net = Unet_GridAttention2D(feature_scale=arg.feature_scale) elif arg.model == "unet_multiatt": net = Unet_CT_multi_attention_dsv_2D(feature_scale=arg.feature_scale) elif arg.model == "unet_dilated": net = Unet_Dilation_2D(feature_scale=arg.feature_scale, dilation=arg.dilation) else: raise NotImplementedError("Not Implemented Model") net = nn.DataParallel(net).to(torch_device) model = CNNTrainer(arg, net, torch_device, recon_loss=nn.BCEWithLogitsLoss(), metric=jaccard_similarity_score) if arg.infer: train_loader = NucleusLoader(train_path, 1, transform=preprocess, channel=arg.in_channel, torch_type=arg.dtype, cpus=arg.cpus, shuffle=False, drop_last=False) valid_loader = NucleusLoader(valid_path, 1, transform=preprocess, channel=arg.in_channel,
#if arg.loss == "l2": # recon_loss = nn.L2Loss() #elif arg.loss == "l1": # recon_loss = nn.L1Loss() backzero = 1 recon_loss = FocalLoss3d_ver1(backzero=backzero) val_loss = FocalLoss3d_ver1(gamma=2, pw=1, threshold=1.0, erode=3, backzero=backzero) model = CNNTrainer(arg, net, torch_device, recon_loss=recon_loss, val_loss=val_loss, logger=logger) #model.load(filename="epoch[0402]_losssum[0.016887].pth.tar") ##1013 version, quite nice #arg.save_dir = "nets_1008_unet_glob1007_disdice_FRE_pw10_erode2_feat1_trans64" #model.load(filename="epoch[0483]_losssum[0.016087].pth.tar") ##1014 version #model.load(filename="epoch[0418]_losssum[0.049294].pth.tar") model.load(filename="epoch[0324]_losssum[0.053456].pth.tar") #filename = "epoch[0401]_losssum[0.059813].pth.tar" #model.load(filename=filename) sdir = "/data1/0000Deployed_JY/181013_analCART"