def __init__(self, opt):
BaseModel.__init__(self, opt)
self.net_names = ['G']
self.net_G = Unet(inchannel=3,
outchannel=3,
ndf=64,
enc_blocks=1,
dec_blocks=1,
depth=3,
concat=True,
bilinear=self.opt.bilinear,
norm_layer='LN')
if self.isTrain:
self.loss_names = ['G_ad_gen', 'D', 'G', 'G_style', 'G_scene']
self.net_names += ['D']
self.net_D = ResDiscriminator(ndf=32,
img_f=128,
layers=4,
use_spect=True)
self.GANloss = AdversarialLoss('lsgan').to(self.device)
self.Styleloss = vgg_style_loss().to(self.device)
self.sceneloss = scene_loss(self.opt.scenepath).to(self.device)
self.optim_G = torch.optim.Adam(self.net_G.parameters(),
lr=opt.lr,
betas=(0.0, 0.999))
self.optim_D = torch.optim.Adam(self.net_D.parameters(),
lr=opt.lr * opt.lrg2d,
betas=(0.0, 0.999))
def __init__(self, opt):
BaseModel.__init__(self, opt)
self.net_names = ['E_s', 'E_e', 'G_se']
self.net_E_s = Encoder_S(n_downsample=3, ndf=64, norm_layer='LN')
self.net_E_e = Encoder_E(inc=3,
n_downsample=4,
outc=self.opt.dimen_e,
ndf=64,
usekl=opt.kl)
self.net_G_se = Decoder_SE(s_inc=self.net_E_s.outc,
e_inc=self.opt.dimen_e,
n_upsample=self.net_E_s.n_downsample,
norm_layer='LN')
if self.isTrain:
self.l1loss = nn.L1Loss()
self.set_Adam_optims(self.net_names)
self.Styleloss = vgg_style_loss().to(self.device)
self.loss_names = ['kl', 'style', 'recon', 'ident', 'total']
def __init__(self, opt):
BaseModel.__init__(self, opt)
self.net_names = ['E_s', 'E_e', 'D_se']
self.net_E_s = S_Encoder()
self.net_E_e = E_Encoder(outc=self.opt.dimen_e, ndf=64, usekl=True)
self.net_D_se = SE_Decoder(e_inc=self.opt.dimen_e)
if self.isTrain:
self.l1loss = nn.L1Loss()
self.Styleloss = vgg_style_loss().to(self.device)
self.loss_names = ['kl', 'style', 'recon', 'ident', 'total']
self.optim_E_s = torch.optim.Adam(self.net_E_s.parameters(),
lr=opt.lr,
betas=(0.0, 0.999))
self.optim_E_e = torch.optim.Adam(self.net_E_e.parameters(),
lr=opt.lr,
betas=(0.0, 0.999))
self.optim_D_se = torch.optim.Adam(self.net_D_se.parameters(),
lr=opt.lr,
betas=(0.0, 0.999))
def __init__(self, opt):
BaseModel.__init__(self, opt)
self.net_names = ['G']
self.net_G = Unet(inchannel=3,
outchannel=3,
ndf=64,
enc_blocks=2,
dec_blocks=2,
depth=3,
bilinear=self.opt.bilinear,
norm_layer='LN')
if self.isTrain:
self.loss_names = ['G_ad_gen', 'D', 'G', 'G_style', 'G_scene']
self.net_names += ['D']
self.net_D = ResDiscriminator(ndf=32,
img_f=128,
layers=4,
use_spect=True)
self.GANloss = AdversarialLoss('lsgan').to(self.device)
self.Styleloss = vgg_style_loss().to(self.device)
self.sceneloss = scene_loss('./checkpoints/net_E_s.path').to(
self.device)
self.set_Adam_optims(self.net_names)