def __init__(self, opt):
BaseModel.__init__(self, opt)
self.opt = opt
# specify the training losses you want to print out. The training/test scripts will call <BaseModel.get_current_losses>
self.loss_names = ['G', 'G_L1', 'G_R', 'G_R_SSIM', "IF"]
self.visual_names = ['hdr', 'real', 'fake_u', 'fake_o']
self.model_names = ['G']
self.opt.device = self.device
self.netG = networks.define_G(opt.netG, opt.init_type, opt.init_gain,
self.opt)
self.cur_device = torch.cuda.current_device()
self.ismaster = du.is_master_proc(opt.NUM_GPUS)
print(self.netG)
if self.isTrain:
util.saveprint(self.opt, 'netG', str(self.netG))
# define loss functions
self.criterionL1 = torch.nn.L1Loss()
self.criterionL2 = torch.nn.MSELoss()
self.criterionSSIM = ssim.SSIM()
self.criterionDSSIM_CS = ssim.DSSIM(mode='c_s').to(self.device)
self.optimizer_G = torch.optim.Adam(self.netG.parameters(),
lr=opt.lr,
betas=(opt.beta1, 0.999))
self.optimizers.append(self.optimizer_G)
def __init__(self, opt):
"""Initialize the pix2pix class.
Parameters:
opt (Option class)-- stores all the experiment flags; needs to be a subclass of BaseOptions
"""
BaseModel.__init__(self, opt)
self.opt = opt
# specify the training losses you want to print out. The training/test scripts will call <BaseModel.get_current_losses>
self.loss_names = ['G','G_L1','G_R','G_I_L2','G_I_smooth',"G_L"]
# specify the images you want to save/display. The training/test scripts will call <BaseModel.get_current_visuals>
self.visual_names = ['harmonized','real','fake','reflectance','illumination']
# specify the models you want to save to the disk. The training/test scripts will call <BaseModel.save_networks> and <BaseModel.load_networks>
self.model_names = ['G']
self.opt.device = self.device
self.netG = networks.define_G(opt.netG, opt.init_type, opt.init_gain, self.opt)
self.cur_device = torch.cuda.current_device()
self.ismaster = du.is_master_proc(opt.NUM_GPUS)
if self.ismaster:
print(self.netG)
if self.isTrain:
util.saveprint(self.opt, 'netG', str(self.netG))
# define loss functions
self.criterionL1 = torch.nn.L1Loss()
self.criterionL2 = torch.nn.MSELoss()
self.criterionSSIM = ssim.SSIM()
# initialize optimizers; schedulers will be automatically created by function <BaseModel.setup>.
self.optimizer_G = torch.optim.Adam(self.netG.parameters(), lr=opt.lr, betas=(opt.beta1, 0.999))
self.optimizers.append(self.optimizer_G)
def __init__(self, opt):
BaseModel.__init__(self, opt)
self.opt = opt
self.loss_names = ['G', 'G_L1']
if opt.loss_RH:
self.loss_names.append("G_R_grident")
if opt.loss_IH:
self.loss_names.append("G_I_L2")
if opt.loss_IS:
self.loss_names.append("G_I_smooth")
self.visual_names = [
'mask', 'harmonized', 'comp', 'real', 'reflectance', 'illumination'
]
self.model_names = ['G']
self.opt.device = self.device
self.netG = networks.define_G(opt.netG, opt.init_type, opt.init_gain,
self.opt)
self.cur_device = torch.cuda.current_device()
self.ismaster = du.is_master_proc(opt.NUM_GPUS)
if self.ismaster:
print(self.netG)
if self.isTrain:
if self.ismaster == 0:
util.saveprint(self.opt, 'netG', str(self.netG))
self.criterionL1 = torch.nn.L1Loss().cuda(self.cur_device)
self.criterionL2 = torch.nn.MSELoss().cuda(self.cur_device)
# initialize optimizers; schedulers will be automatically created by function <BaseModel.setup>.
self.optimizer_G = torch.optim.Adam(self.netG.parameters(),
lr=opt.lr,
betas=(opt.beta1, 0.999))
self.optimizers.append(self.optimizer_G)
def __init__(self, opt):
BaseModel.__init__(self, opt)
self.opt = opt
# specify the training losses you want to print out. The training/test scripts will call <BaseModel.get_current_losses>
self.loss_names = ['G','G_L1','G_R','G_R_SSIM',"IF"]
# specify the images you want to save/display. The training/test scripts will call <BaseModel.get_current_visuals>
self.visual_names = ['enhanced','real','fake','reconstruct','illumination']
# specify the models you want to save to the disk. The training/test scripts will call <BaseModel.save_networks> and <BaseModel.load_networks>
self.model_names = ['G']
self.opt.device = self.device
self.netG = networks.define_G(opt.netG, opt.init_type, opt.init_gain, self.opt)
self.cur_device = torch.cuda.current_device()
self.ismaster = du.is_master_proc(opt.NUM_GPUS)
print(self.netG)
if self.isTrain:
# if self.ismaster == 0:
util.saveprint(self.opt, 'netG', str(self.netG))
# define loss functions
self.criterionL1 = torch.nn.L1Loss()
self.criterionL2 = torch.nn.MSELoss()
self.criterionSSIM = ssim.SSIM()
self.criterionDSSIM_CS = ssim.DSSIM(mode='c_s').to(self.device)
# initialize optimizers; schedulers will be automatically created by function <BaseModel.setup>.
self.optimizer_G = torch.optim.Adam(self.netG.parameters(), lr=opt.lr, betas=(opt.beta1, 0.999))
self.optimizers.append(self.optimizer_G)
def test(cfg):
dataset = create_dataset(cfg) # create a dataset given cfg.dataset_mode and other options
model = create_model(cfg) # create a model given cfg.model and other options
model.setup(cfg) # regular setup: load and print networks; create schedulers
# create a website
web_dir = os.path.join(cfg.results_dir, cfg.name, '%s_%s' % (cfg.phase, cfg.epoch)) # define the website directory
webpage = html.HTML(web_dir, 'Experiment = %s, Phase = %s, Epoch = %s' % (cfg.name, cfg.phase, cfg.epoch))
if cfg.eval:
model.eval()
ismaster = du.is_master_proc(cfg.NUM_GPUS)
num_image = 0
for i, data in enumerate(dataset):
model.set_input(data) # unpack data from data loader
model.test() # run inference
visuals = model.get_current_visuals() # get image results
img_path = model.get_image_paths() # get image paths # Added by Mia
if i % 5 == 0 and ismaster: # save images to an HTML file
print('processing (%04d)-th image... %s' % (i, img_path))
visuals_ones = OrderedDict()
for j in range(len(img_path)):
img_path_one = []
for label, im_data in visuals.items():
visuals_ones[label] = im_data[j:j+1, :, :, :]
img_path_one.append(img_path[j])
save_images(webpage, visuals_ones, img_path_one, aspect_ratio=cfg.aspect_ratio, width=cfg.display_winsize)
num_image += 1
visuals_ones.clear()
webpage.save() # save the HTML
def __init__(self, opt):
BaseModel.__init__(self, opt)
self.opt = opt
# specify the training losses you want to print out. The training/test scripts will call <BaseModel.get_current_losses>
self.loss_names = ['G','G_L1',"IF"]
if opt.loss_RH:
self.loss_names.append("G_R_grident")
if opt.loss_IH:
self.loss_names.append("G_I_L2")
if opt.loss_IS:
self.loss_names.append("G_I_smooth")
# specify the images you want to save/display. The training/test scripts will call <BaseModel.get_current_visuals>
self.visual_names = ['mask', 'harmonized','comp','real','reflectance','illumination','ifm_mean']
# specify the models you want to save to the disk. The training/test scripts will call <BaseModel.save_networks> and <BaseModel.load_networks>
self.model_names = ['G']
self.opt.device = self.device
self.netG = networks.define_G(opt.netG, opt.init_type, opt.init_gain, self.opt)
self.cur_device = torch.cuda.current_device()
self.ismaster = du.is_master_proc(opt.NUM_GPUS)
if self.ismaster:
print(self.netG)
if self.isTrain:
util.saveprint(self.opt, 'netG', str(self.netG))
# define loss functions
self.criterionL1 = torch.nn.L1Loss().cuda(self.cur_device)
self.criterionL2 = torch.nn.MSELoss().cuda(self.cur_device)
self.criterionDSSIM_CS = ssim.DSSIM(mode='c_s').to(self.device)
self.optimizer_G = torch.optim.Adam(self.netG.parameters(), lr=opt.lr, betas=(opt.beta1, 0.999))
self.optimizers.append(self.optimizer_G)
def test(cfg):
dataset = create_dataset(
cfg) # create a dataset given cfg.dataset_mode and other options
model = create_model(
cfg) # create a model given cfg.model and other options
model.setup(
cfg) # regular setup: load and print networks; create schedulers
# create a website
web_dir = os.path.join(
cfg.results_dir, cfg.name,
'%s_%s' % (cfg.phase, cfg.epoch)) # define the website directory
webpage = html.HTML(
web_dir, 'Experiment = %s, Phase = %s, Epoch = %s' %
(cfg.name, cfg.phase, cfg.epoch))
# test with eval mode. This only affects layers like batchnorm and dropout.
# For [pix2pix]: we use batchnorm and dropout in the original pix2pix. You can experiment it with and without eval() mode.
# For [CycleGAN]: It should not affect CycleGAN as CycleGAN uses instancenorm without dropout.
if cfg.eval:
model.eval()
ismaster = du.is_master_proc(cfg.NUM_GPUS)
fmse_score_list = []
mse_scores = 0
fmse_scores = 0
num_image = 0
for i, data in enumerate(dataset):
model.set_input(data) # unpack data from data loader
model.test() # run inference
visuals = model.get_current_visuals() # get image results
img_path = model.get_image_paths() # get image paths # Added by Mia
if i % 5 == 0 and ismaster: # save images to an HTML file
print('processing (%04d)-th image... %s' % (i, img_path))
visuals_ones = OrderedDict()
harmonized = None
real = None
for j in range(len(img_path)):
img_path_one = []
for label, im_data in visuals.items():
visuals_ones[label] = im_data[j:j + 1, :, :, :]
img_path_one.append(img_path[j])
save_images(webpage,
visuals_ones,
img_path_one,
aspect_ratio=cfg.aspect_ratio,
width=cfg.display_winsize)
num_image += 1
visuals_ones.clear()
webpage.save() # save the HTML
def __init__(self, opt):
"""Initialize the pix2pix class.
Parameters:
opt (Option class)-- stores all the experiment flags; needs to be a subclass of BaseOptions
"""
BaseModel.__init__(self, opt)
self.opt = opt
self.loss_names = [
'G', 'G_L1', 'G_R_grident', 'G_I_L2', 'G_I_smooth', "IF"
]
# specify the images you want to save/display. The training/test scripts will call <BaseModel.get_current_visuals>
self.visual_names = [
'mask', 'harmonized', 'comp', 'real', 'reflectance',
'illumination', 'ifm_mean'
]
# specify the models you want to save to the disk. The training/test scripts will call <BaseModel.save_networks> and <BaseModel.load_networks>
self.model_names = ['G']
self.opt.device = self.device
self.netG = networks.define_G(opt.netG, opt.init_type, opt.init_gain,
self.opt)
self.cur_device = torch.cuda.current_device()
self.ismaster = du.is_master_proc(opt.NUM_GPUS)
if self.ismaster:
print(self.netG)
if self.isTrain:
if self.ismaster == 0:
util.saveprint(self.opt, 'netG', str(self.netG))
# define loss functions
self.criterionL1 = torch.nn.L1Loss().cuda(self.cur_device)
self.criterionL2 = torch.nn.MSELoss().cuda(self.cur_device)
self.criterionDSSIM_CS = ssim.DSSIM(mode='c_s').to(self.device)
self.optimizer_G = torch.optim.Adam(self.netG.parameters(),
lr=opt.lr,
betas=(opt.beta1, 0.999))
self.optimizers.append(self.optimizer_G)
def train(cfg):
#init
du.init_distributed_training(cfg)
# Set random seed from configs.
np.random.seed(cfg.RNG_SEED)
torch.manual_seed(cfg.RNG_SEED)
#init dataset
dataset = create_dataset(
cfg) # create a dataset given cfg.dataset_mode and other options
dataset_size = len(dataset) # get the number of images in the dataset.
print('The number of training images = %d' % dataset_size)
model = create_model(
cfg) # create a model given cfg.model and other options
model.setup(
cfg) # regular setup: load and print networks; create schedulers
visualizer = Visualizer(
cfg) # create a visualizer that display/save images and plots
total_iters = 0 # the total number of training iterations
# cur_device = torch.cuda.current_device()
is_master = du.is_master_proc(cfg.NUM_GPUS)
for epoch in range(
cfg.epoch_count, cfg.niter + cfg.niter_decay + 1
): # outer loop for different epochs; we save the model by <epoch_count>, <epoch_count>+<save_latest_freq>
if is_master:
epoch_start_time = time.time() # timer for entire epoch
iter_data_time = time.time(
) # timer for data loading per iteration
epoch_iter = 0 # the number of training iterations in current epoch, reset to 0 every epoch
shuffle_dataset(dataset, epoch)
for i, data in enumerate(dataset): # inner loop within one epoch
if is_master:
iter_start_time = time.time(
) # timer for computation per iteration
if total_iters % cfg.print_freq == 0:
t_data = iter_start_time - iter_data_time
iter_data_time = time.time()
visualizer.reset()
total_iters += cfg.batch_size
epoch_iter += cfg.batch_size
model.set_input(
data) # unpack data from dataset and apply preprocessing
model.optimize_parameters(
) # calculate loss functions, get gradients, update network weights
if total_iters % cfg.display_freq == 0 and is_master: # display images on visdom and save images to a HTML file
save_result = total_iters % cfg.update_html_freq == 0
model.compute_visuals()
visualizer.display_current_results(model.get_current_visuals(),
epoch, save_result)
losses = model.get_current_losses()
if cfg.NUM_GPUS > 1:
losses = du.all_reduce(losses)
if total_iters % cfg.print_freq == 0 and is_master: # print training losses and save logging information to the disk
t_comp = (time.time() - iter_start_time) / cfg.batch_size
visualizer.print_current_losses(epoch, epoch_iter, losses,
t_comp, t_data)
if cfg.display_id > 0:
visualizer.plot_current_losses(
epoch,
float(epoch_iter) / dataset_size, losses)
if total_iters % cfg.save_latest_freq == 0 and is_master: # cache our latest model every <save_latest_freq> iterations
print('saving the latest model (epoch %d, total_iters %d)' %
(epoch, total_iters))
save_suffix = 'iter_%d' % total_iters if cfg.save_by_iter else 'latest'
model.save_networks(save_suffix)
if epoch % cfg.save_epoch_freq == 0 and is_master: # cache our model every <save_epoch_freq> epochs
print('saving the model at the end of epoch %d, iters %d' %
(epoch, total_iters))
model.save_networks('latest')
if cfg.save_iter_model and epoch >= 80:
model.save_networks(epoch)
if is_master:
print('End of epoch %d / %d \t Time Taken: %d sec' %
(epoch, cfg.niter + cfg.niter_decay,
time.time() - epoch_start_time))
model.update_learning_rate(
) # update learning rates at the end of every epoch.
