def __init__(self, opt):
super(ModelGAN, self).__init__(opt)
# ------------------------------------
# define network
# ------------------------------------
self.netG = define_G(opt).to(self.device)
self.netG = DataParallel(self.netG)
if self.is_train:
self.netF = define_F(opt).to(self.device)
self.netD = define_D(opt).to(self.device)
self.netF = DataParallel(self.netF)
self.netD = DataParallel(self.netD)
def define_loss(self):
# ------------------------------------
# G_loss
# ------------------------------------
if self.opt_train['G_lossfn_weight'] > 0:
G_lossfn_type = self.opt_train['G_lossfn_type']
if G_lossfn_type == 'l1':
self.G_lossfn = nn.L1Loss().to(self.device)
elif G_lossfn_type == 'l2':
self.G_lossfn = nn.MSELoss().to(self.device)
elif G_lossfn_type == 'l2sum':
self.G_lossfn = nn.MSELoss(reduction='sum').to(self.device)
elif G_lossfn_type == 'ssim':
self.G_lossfn = SSIMLoss().to(self.device)
else:
raise NotImplementedError(
'Loss type [{:s}] is not found.'.format(G_lossfn_type))
self.G_lossfn_weight = self.opt_train['G_lossfn_weight']
else:
print('Do not use pixel loss.')
self.G_lossfn = None
# ------------------------------------
# F_loss
# ------------------------------------
if self.opt_train['F_lossfn_weight'] > 0:
F_lossfn_type = self.opt_train['F_lossfn_type']
if F_lossfn_type == 'l1':
self.F_lossfn = nn.L1Loss().to(self.device)
elif F_lossfn_type == 'l2':
self.F_lossfn = nn.MSELoss().to(self.device)
else:
raise NotImplementedError(
'Loss type [{:s}] not recognized.'.format(F_lossfn_type))
self.F_lossfn_weight = self.opt_train['F_lossfn_weight']
self.netF = define_F(self.opt, use_bn=False).to(self.device)
else:
print('Do not use feature loss.')
self.F_lossfn = None
# ------------------------------------
# D_loss
# ------------------------------------
self.D_lossfn = GANLoss(self.opt_train['gan_type'], 1.0,
0.0).to(self.device)
self.D_lossfn_weight = self.opt_train['D_lossfn_weight']
self.D_update_ratio = self.opt_train[
'D_update_ratio'] if self.opt_train['D_update_ratio'] else 1
self.D_init_iters = self.opt_train['D_init_iters'] if self.opt_train[
'D_init_iters'] else 0