def initialize(self, opt):
BaseModel.initialize(self, opt)
self.isTrain = opt.isTrain
self.loss_names = ['loss']
self.model_names = ['fc1', 'fc2']
self.fc1 = init_net(MyLinear(opt.max_idx, 4096), init_type='normal', gpu=self.opt.gpu)
self.fc2 = init_net(MyLinear(4096, 1), init_type='normal', gpu=self.opt.gpu)
if opt.isTrain:
if opt.propensity == 'no':
self.criterion = nn.BCELoss(size_average=True)
if opt.gpu >= 0:
self.criterion.cuda(opt.gpu)
else:
self.criterion = None
self.schedulers = []
self.optimizers = []
self.optimizer = torch.optim.Adam(self.parameters(), lr=opt.lr, weight_decay=1e-5)
# self.optimizer = torch.optim.SparseAdam(self.parameters(), lr=opt.lr)
self.optimizers.append(self.optimizer)
for optimizer in self.optimizers:
self.schedulers.append(networks.get_scheduler(optimizer, opt))
if not self.isTrain or opt.continue_train:
self.load_networks(opt.which_epoch)
self.print_networks()
def initialize(self):
self.isTrain = cfg.isTrain
self.gpu_ids = cfg.gpu_ids
self.device = torch.device('cuda:{}'.format(
cfg.gpu_ids[0])) if int(cfg.gpu_ids) > -1 else torch.device('cpu')
self.save_dir = os.path.join(cfg.checkpoints_dir, cfg.load)
if not os.path.exists(self.save_dir):
os.makedirs(self.save_dir)
print("Create file path: ", self.save_dir)
self.loss_names = []
self.model_names = []
self.visual_names = []
self.image_paths = []
self.loss_names = ['D_A', 'G_A', 'cycle_A', 'D_B', 'G_B', 'cycle_B']
self.loss_dict = {}
# specify the models you want to save to the disk. The program will call base_model.save_networks and base_model.load_networks
if self.isTrain:
self.model_names = ['G_A', 'G_B', 'D_A', 'D_B']
else: # during test time, only load Gs
self.model_names = ['G_A', 'G_B']
# load/define networks
# The naming conversion is different from those used in the paper
# Code (paper): G_A (G), G_B (F), D_A (D_Y), D_B (D_X)
self.netG_A = init_net(L2H_RDN.L2H_RDN())
self.netG_B = init_net(H2L_RDN.H2L_RDN())
if self.isTrain:
self.netD_A = init_net(NLayerDiscriminator())
self.netD_B = init_net(NLayerDiscriminator())
if self.isTrain:
# define loss functions
self.criterionGAN = networks.GANLoss().to(self.device)
self.criterionCycle = torch.nn.L1Loss()
self.criterionIdt = torch.nn.L1Loss()
# initialize optimizers
self.optimizer_G = torch.optim.Adam(itertools.chain(
self.netG_A.parameters(), self.netG_B.parameters()),
lr=cfg.lr,
betas=(cfg.beta, 0.999))
self.optimizer_D = torch.optim.Adam(itertools.chain(
self.netD_A.parameters(), self.netD_B.parameters()),
lr=cfg.lr,
betas=(cfg.beta, 0.999))
self.optimizers = []
self.optimizers.append(self.optimizer_G)
self.optimizers.append(self.optimizer_D)
def initialize(self, opt):
BaseModel.initialize(self, opt)
self.isTrain = opt.isTrain
self.loss_names = ['loss']
self.info_names = []
self.model_names = ['fc1', 'fc2']
self.fc1 = init_net(nn.Linear(opt.max_idx, 4096),
init_type='normal',
gpu=self.opt.gpu)
self.fc2 = init_net(nn.Linear(4096, 1),
init_type='normal',
gpu=self.opt.gpu)
class POEMLoss(nn.Module):
def __init__(self):
super(POEMLoss, self).__init__()
def forward(self, weight, theta):
return torch.sum(theta * weight) / weight.shape[0]
if opt.isTrain:
self.criterion = POEMLoss()
self.schedulers = []
self.optimizers = []
self.optimizer = torch.optim.Adam(self.parameters(),
lr=opt.lr,
weight_decay=1e-3)
self.optimizers.append(self.optimizer)
for optimizer in self.optimizers:
self.schedulers.append(networks.get_scheduler(optimizer, opt))
if not self.isTrain or opt.continue_train:
self.load_networks(opt.which_epoch)
self.print_networks()