if arg.norm == "bn":
norm_layer = nn.BatchNorm2d
elif arg.norm == "in":
norm_layer = nn.InstanceNorm2d
if arg.act == "relu":
act = nn.ReLU
elif arg.act == "elu":
act = nn.ELU
elif arg.act == "leaky":
act = nn.LeakyReLU
elif arg.act == "prelu":
act = nn.PReLU
if arg.model == "fusion":
net = Fusionnet(1, 1, arg.ngf, arg.clamp)
elif arg.model == "unet":
net = Unet2D(feature_scale=arg.feature_scale,
is_pool=arg.pool,
act=act)
elif arg.model == "unetgn":
net = UnetGN2D(feature_scale=arg.feature_scale, is_pool=arg.pool)
elif arg.model == "unetslim":
net = UnetSlim(feature_scale=arg.feature_scale, norm=norm_layer)
elif arg.model == "unet_sh":
net = UnetSH2D(arg.sh_size,
feature_scale=arg.feature_scale,
is_pool=arg.pool)
elif arg.model == "unetres":
net = UnetRes2D(1, nn.InstanceNorm2d, is_pool=arg.pool)
elif arg.model == "unetgcn":
os.environ["CUDA_VISIBLE_DEVICES"] = arg.gpus
torch_device = torch.device("cuda")
preprocess = preprocess.get_preprocess(arg.augment)
train_loader = nucleusloader(f_path_train, arg.batch_size, transform=preprocess,
cpus=arg.cpus,
shuffle=True, drop_last=True)
valid_loader = nucleusloader3(f_path_valid, batch_size=1, transform=None,
cpus=arg.cpus, shuffle=False,
drop_last=True)
if arg.model == "fusion":
net = Fusionnet(arg.in_channel, arg.out_channel, arg.ngf, arg.clamp)
elif arg.model == "unet":
net = Unet3D(feature_scale=arg.feature_scale)
elif arg.model == "unet_gh":
## "nets_1004_unet_glob_absloss_FRE_pw10_erode2_feat1_trans30"
#net = Unet3D_glob2(feature_scale=arg.feature_scale, trans_feature=64)
#net = Unet3D_glob(feature_scale=arg.feature_scale, trans_feature=64)
#net = Unet3D_glob1106(feature_scale=arg.feature_scale, trans_feature=64)
#net = Unet3D_glob1106(feature_scale=arg.feature_scale, trans_feature=24)
#for arg.save_dir = "nets_1110_unet_glob1109_class_FRE_pw1_erode2_feat2_trans24_signdiff_newdata"
#net=Unet3D_glob1109(feature_scale=arg.feature_scale, trans_feature=24)
#net = Unet3D_glob1109(feature_scale=arg.feature_scale, trans_feature=32)
net=Unet3D_glob1111(feature_scale=arg.feature_scale, trans_feature=24)
elif arg.model == "unet_reduced":
valid_loader = MARLoader(
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")