def __init__(self, params):
self.params = params
self.model_dict = {}
self.opt_dict = {}
self.current_epoch = 0
self.current_iter = 0
self.preview_noise = helper.new_random_z(16, params['z_size'], seed=3)
self.transform = load.NormDenorm([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])
self.train_loader, self.data_len = load.data_load(
f'data/{params["dataset"]}/{params["train_folder"]}/',
self.transform,
params["batch_size"],
shuffle=True,
perc=params["data_perc"],
output_res=params["output_size"])
print(f'Data Loader Initialized: {self.data_len} Images')
self.model_dict["G"] = n.Generator(layers=params["gen_layers"],
filts=params["gen_filters"],
channels=params["in_channels"],
z_size=params['z_size'])
self.model_dict["D"] = n.Discriminator(layers=params["disc_layers"],
filts=params["disc_filters"],
channels=params["in_channels"])
for i in self.model_dict.keys():
self.model_dict[i].apply(helper.weights_init_normal)
self.model_dict[i].cuda()
self.model_dict[i].train()
print('Networks Initialized')
# setup optimizers #
self.opt_dict["G"] = optim.RMSprop(self.model_dict["G"].parameters(),
lr=params['lr_gen'])
self.opt_dict["D"] = optim.RMSprop(self.model_dict["D"].parameters(),
lr=params['lr_disc'])
print('Optimizers Initialized')
# setup history storage #
self.losses = ['G_Loss', 'D_Loss']
self.loss_batch_dict = {}
self.loss_epoch_dict = {}
self.train_hist_dict = {}
for loss in self.losses:
self.train_hist_dict[loss] = []
self.loss_epoch_dict[loss] = []
self.loss_batch_dict[loss] = []
def __init__(self, params):
self.params = params
self.model_dict = {}
self.opt_dict = {}
self.current_epoch = 0
self.current_iter = 0
self.current_cycle = 0
# Setup data loaders
self.transform = load.NormDenorm([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])
self.train_loader, data_len = load.data_load(f'/data/{params["dataset"]}/{params["train_folder"]}/',
self.transform,
params["batch_size"],
shuffle=True,
output_res=params["img_output_size"],
perc=params["test_perc"])
self.test_loader, test_data_len = load.data_load(f'/data/{params["dataset"]}/{params["test_folder"]}/',
self.transform,
1,
shuffle=False,
perc=params["test_perc"],
output_res=params["img_output_size"],
train=False)
# Set learning rate schedule
self.set_lr_sched(params['train_epoch'],
math.ceil(float(data_len) / float(params['batch_size'])),
params['lr_cycle_mult'])
# Setup models
self.model_dict["G"] = n.Generator(layers=params["gen_layers"],
filts=params["gen_filters"],
channels=params["in_channels"])
self.model_dict["D"] = n.Discriminator(layers=params["disc_layers"],
filts=params["disc_filters"],
channels=params["in_channels"] * 2)
for i in self.model_dict.keys():
self.model_dict[i].apply(helper.weights_init_normal)
self.model_dict[i].cuda()
self.model_dict[i].train()
print('Networks Initialized')
# Setup losses
self.BCE_loss = nn.BCELoss()
self.L1_loss = nn.L1Loss()
# Setup optimizers
self.opt_dict["G"] = optim.Adam(self.model_dict["G"].parameters(),
lr=params['lr_gen'],
betas=(params['beta1'], params['beta2']),
weight_decay=.00001)
self.opt_dict["D"] = optim.Adam(self.model_dict["D"].parameters(),
lr=params['lr_disc'],
betas=(params['beta1'], params['beta2']),
weight_decay=.00001)
print('Losses Initialized')
# Setup history storage
self.losses = ['D_loss', 'G_D_loss', 'G_L_loss']
self.loss_batch_dict = {}
self.loss_batch_dict_test = {}
self.loss_epoch_dict = {}
self.loss_epoch_dict_test = {}
self.train_hist_dict = {}
self.train_hist_dict_test = {}
for loss in self.losses:
self.train_hist_dict[loss] = []
self.loss_epoch_dict[loss] = []
self.loss_batch_dict[loss] = []
self.train_hist_dict_test[loss] = []
self.loss_epoch_dict_test[loss] = []
self.loss_batch_dict_test[loss] = []
def main(opt):
# Training config
print(opt)
t.manual_seed(0)
# Parameters
lambda_FM = 10
lambda_P = 10
lambda_2 = opt.lambda_second
nf = 64 # 64
n_blocks = 6 # 6
# Load the networks
if t.cuda.is_available():
device = "cuda"
else:
device = 'cpu'
print(f"Device: {device}")
if opt.segment:
disc = networks.MultiScaleDisc(input_nc=4, ndf=nf).to(device)
gen = networks.Generator(input_nc=6,
output_nc=1,
ngf=nf,
n_blocks=n_blocks,
transposed=opt.transposed).to(device)
else:
disc = networks.MultiScaleDisc(input_nc=1, ndf=nf).to(device)
gen = networks.Generator(input_nc=3,
output_nc=1,
ngf=nf,
n_blocks=n_blocks,
transposed=opt.transposed).to(device)
if opt.current_epoch != 0:
disc.load_state_dict(
t.load(
os.path.join(opt.checkpoints_file,
f"e_{opt.current_epoch:0>3d}_discriminator.pth")))
gen.load_state_dict(
t.load(
os.path.join(opt.checkpoints_file,
f"e_{opt.current_epoch:0>3d}_generator.pth")))
print(f"- e_{opt.current_epoch:0>3d}_generator.pth was loaded! -")
print(f"- e_{opt.current_epoch:0>3d}_discriminator.pth was loaded! -")
else:
disc.apply(utils.weights_init)
gen.apply(utils.weights_init)
print("- Weights are initialized from scratch -")
# Losses to track
# # Main losses
loss_change_g = []
loss_change_d = []
# # Components
loss_change_fm1 = []
loss_change_fm2 = []
loss_change_d1 = []
loss_change_d2 = []
loss_change_g1 = []
loss_change_g2 = []
loss_change_p = []
# Create optimizers (Notice the lr of discriminator)
optim_g = optim.Adam(gen.parameters(),
lr=opt.learning_rate / 5,
betas=(0.5, 0.999))
optim_d = optim.Adam(disc.parameters(),
lr=opt.learning_rate,
betas=(0.5, 0.999),
weight_decay=0.0001)
# Create Schedulers
# g_scheduler = t.optim.lr_scheduler.LambdaLR(optim_g, utils.lr_lambda)
# d_scheduler = t.optim.lr_scheduler.LambdaLR(optim_d, utils.lr_lambda)
# Create loss functions
loss = losses.GanLoss()
loss_fm = losses.FeatureMatchingLoss()
loss_p = losses.VGGLoss(device) # perceptual loss
# Create dataloader
ds = dataset.CustomDataset(opt.data_dir,
is_segment=opt.segment,
sf=opt.scale_factor)
dataloader = DataLoader(ds,
batch_size=opt.batch_size,
shuffle=True,
num_workers=2)
# Start to training
print("Training is starting...")
i = 0
for e in range(1 + opt.current_epoch,
1 + opt.training_epoch + opt.current_epoch):
print(f"---- Epoch #{e} ----")
start = time.time()
for data in tqdm(dataloader):
i += 1
rgb = data[0].to(device)
ir = data[1].to(device)
if opt.segment:
segment = data[2].to(device)
condition = t.cat([rgb, segment], dim=1)
ir_ = t.cat([ir, segment], dim=1)
else:
condition = rgb
ir_ = ir
out1, out2 = disc(ir_)
ir_pred = gen(condition)
# # # Updating Discriminator # # #
optim_d.zero_grad()
if opt.segment:
ir_pred_ = t.cat([ir_pred, segment], dim=1)
else:
ir_pred_ = ir_pred
out1_pred, out2_pred = disc(
ir_pred_.detach()) # It returns a list [fms... + output]
l_d_pred1, l_d_pred2 = loss(out1_pred[-1],
out2_pred[-1],
is_real=False)
l_d_real1, l_d_real2 = loss(out1[-1], out2[-1], is_real=True)
l_d_scale1 = l_d_pred1 + l_d_real1
l_d_scale2 = l_d_pred2 + l_d_real2
disc_loss = l_d_scale1 + l_d_scale2 * lambda_2
# Normalize the loss, and track
loss_change_d += [disc_loss.item() / opt.batch_size]
loss_change_d1 += [l_d_scale1.item() / opt.batch_size]
loss_change_d2 += [l_d_scale2.item() / opt.batch_size]
disc_loss.backward()
optim_d.step()
# # # Updating Generator # # #
optim_g.zero_grad()
out1_pred, out2_pred = disc(
ir_pred_) # It returns a list [fms... + output]
fm_scale1 = loss_fm(out1_pred[:-1], out1[:-1])
fm_scale2 = loss_fm(out2_pred[:-1], out2[:-1])
fm = fm_scale1 + fm_scale2 * lambda_2
perceptual = loss_p(ir_pred, ir)
loss_change_fm1 += [fm_scale1.item() / opt.batch_size]
loss_change_fm2 += [fm_scale2.item() / opt.batch_size]
loss_change_p += [perceptual.item() / opt.batch_size]
l_g_scale1, l_g_scale2 = loss(out1_pred[-1],
out2_pred[-1],
is_real=True)
gen_loss = l_g_scale1 + l_g_scale2 * lambda_2 + fm * lambda_FM + perceptual * lambda_P
loss_change_g += [gen_loss.item() / opt.batch_size]
loss_change_g1 += [l_g_scale1.item() / opt.batch_size]
loss_change_g2 += [l_g_scale2.item() / opt.batch_size]
gen_loss.backward()
optim_g.step()
# Save images
if i % opt.img_save_freq == 1:
utils.save_tensor_images(ir_pred, i, opt.results_file, 'pred')
utils.save_tensor_images(ir, i, opt.results_file, 'ir')
utils.save_tensor_images(rgb, i, opt.results_file, 'rgb')
utils.save_tensor_images(segment, i, opt.results_file,
'segment')
print('\nExample images saved')
print("Losses:")
print(
f"G: {loss_change_g[-1]:.4f}; D: {loss_change_d[-1]:.4f}")
print(
f"G1: {loss_change_g1[-1]:.4f}; G2: {loss_change_g2[-1]:.4f}"
)
print(
f"D1: {loss_change_d1[-1]:.4f}; D2: {loss_change_d2[-1]:.4f}"
)
print(
f"FM1: {loss_change_fm1[-1]:.4f}; FM2: {loss_change_fm2[-1]:.4f}; P: {loss_change_p[-1]:.4f}"
)
# g_scheduler.step()
# d_scheduler.step()
print(
f"Epoch duration: {int((time.time() - start) // 60):5d}m {(time.time() - start) % 60:.1f}s"
)
if i % opt.model_save_freq == 0:
utils.save_model(disc, gen, e, opt.checkpoints_file)
# End of training
# Main losses are g and d, but I want to save all components separately
utils.save_loss(d=loss_change_d,
d1=loss_change_d1,
d2=loss_change_d2,
g=loss_change_g,
g1=loss_change_g1,
g2=loss_change_g2,
fm1=loss_change_fm1,
fm2=loss_change_fm2,
p=loss_change_p,
path=opt.loss_file,
e=e)
utils.save_model(disc, gen, e, opt.checkpoints_file)
utils.show_loss(opt.checkpoints_file)
print("Done!")
def main(opt):
print(opt)
nf = 64
n_blocks = 6
# Load the networks
if t.cuda.is_available():
device = "cuda"
else:
device = 'cpu'
print(f"Device: {device}")
if opt.segment:
gen = networks.Generator(input_nc=6,
output_nc=1,
ngf=nf,
n_blocks=n_blocks,
transposed=opt.transposed).to(device)
else:
gen = networks.Generator(input_nc=3,
output_nc=1,
ngf=nf,
n_blocks=n_blocks,
transposed=opt.transposed).to(device)
gen.load_state_dict(
t.load(
os.path.join(opt.checkpoints_file,
f"e_{opt.current_epoch:0>3d}_generator.pth")))
gen.eval()
try:
ds = dataset.CustomDataset(root_dir=opt.inp_file,
sf=opt.scale_factor,
is_segment=opt.segment)
except:
raise NotImplementedError
print("IR images not found but it is ok")
ds = dataset.TestDataset(root_dir=opt.inp_file)
dataloader = DataLoader(ds, batch_size=1, shuffle=False, num_workers=2)
if opt.batch_size != 1:
print("Batch size other than 1, is not implemented yet")
i = 0
for data in tqdm(dataloader):
i += 1
with t.no_grad():
rgb = data[0].to(device)
ir = data[1]
if opt.segment:
segment = data[-1].to(device)
condition = t.cat([rgb, segment], dim=1)
else:
condition = rgb
ir_pred = gen(condition)
if opt.segment:
utils.save_all_images(rgb,
ir,
ir_pred,
i,
opt.out_file,
segment=segment,
resize_factor=0.5)
else:
utils.save_all_images(rgb,
ir,
ir_pred,
i,
opt.out_file,
resize_factor=0.5)
print("Done!")
def __init__(self, params):
self.params = params
self.model_dict = {}
self.opt_dict = {}
self.current_epoch = 0
self.current_iter = 0
# Setup data loaders
self.transform = load.NormDenorm([.485, .456, .406],
[.229, .224, .225])
self.train_data = load.GenericDataset(f'data/{params["dataset"]}',
self.transform,
output_res=params["res"],
test_perc=params["test_perc"],
data_perc=params["data_perc"])
self.data_len = self.train_data.__len__()
self.test_data = copy.deepcopy(self.train_data)
self.test_data.train = False
self.train_loader = torch.utils.data.DataLoader(
self.train_data,
batch_size=params["batch_size"],
num_workers=params["workers"],
shuffle=True,
drop_last=True)
self.test_loader = torch.utils.data.DataLoader(
self.test_data,
batch_size=params["batch_size"],
num_workers=params["workers"],
shuffle=True,
drop_last=True)
print('Data Loaders Initialized')
# Setup models
self.model_dict["G"] = n.Generator(layers=params["gen_layers"],
filts=params["gen_filters"],
channels=params["in_channels"],
res_layers=params["res_blocks"])
self.tensor_transform = n.TensorTransform(res=params["res"],
mean=[.485, .456, .406],
std=[.229, .224, .225])
self.tensor_transform.cuda()
for i in self.model_dict.keys():
self.model_dict[i].apply(helper.weights_init_normal)
self.model_dict[i].cuda()
self.model_dict[i].train()
print('Networks Initialized')
# Setup loss
self.style = load.open_style(f'style/{params["style_image"]}',
self.transform,
batch_size=params["batch_size"],
size=params["res"]).cuda()
self.vgg = n.make_vgg()
self.vgg.cuda()
self.cs_loss = n.ContStyleLoss(self.vgg, self.style,
params['content_weight'],
params['style_weight'],
params['vgg_layers_s'],
params['vgg_layers_c'])
# Setup optimizers
self.opt_dict["G"] = optim.Adam(self.model_dict["G"].parameters(),
lr=params['lr'],
betas=(params['beta1'],
params['beta2']))
print('Losses Initialized')
# Setup history storage
self.losses = ['S_Loss', 'C_Loss']
self.loss_batch_dict = {}
self.loss_batch_dict_test = {}
self.loss_epoch_dict = {}
self.loss_epoch_dict_test = {}
self.train_hist_dict = {}
self.train_hist_dict_test = {}
for loss in self.losses:
self.train_hist_dict[loss] = []
self.loss_epoch_dict[loss] = []
self.loss_batch_dict[loss] = []
self.train_hist_dict_test[loss] = []
self.loss_epoch_dict_test[loss] = []
self.loss_batch_dict_test[loss] = []
def __init__(self, params):
self.params = params
self.model_dict = {}
self.opt_dict = {}
self.current_epoch = 0
self.current_iter = 0
self.current_cycle = 0
transform = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize(mean=(.5, .5, .5), std=(.5, .5, .5))
])
self.train_loader, data_len = load.data_load(
f'data/{params["dataset"]}/{params["train_folder"]}/{params["A"]}',
f'data/{params["dataset"]}/{params["train_folder"]}/{params["B"]}',
transform,
params["batch_size"],
shuffle=True,
cache=True,
cache_file=f'{params["dataset"]}_content_cache.pickle',
close=params["similar_distance"],
input_res=params["img_input_size"],
output_res=params["img_output_size"])
self.set_lr_sched(
params['train_epoch'],
math.ceil(float(data_len) / float(params['batch_size'])),
params['lr_cycle_mult'])
self.model_dict["G_A"] = n.Generator(layers=params["gen_layers"],
filts=params["gen_filters"],
channels=params["in_channels"],
res_layers=params["res_blocks"])
self.model_dict["G_B"] = n.Generator(layers=params["gen_layers"],
filts=params["gen_filters"],
channels=params["in_channels"],
res_layers=params["res_blocks"])
self.model_dict["D_A"] = n.Discriminator(
layers=params["disc_layers"],
filts=params["disc_filters"],
channels=params["in_channels"])
self.model_dict["D_B"] = n.Discriminator(
layers=params["disc_layers"],
filts=params["disc_filters"],
channels=params["in_channels"])
for i in self.model_dict.keys():
self.model_dict[i].apply(helper.weights_init_normal)
self.model_dict[i].cuda()
self.model_dict[i].train()
print('Networks Initialized')
# setup losses #
self.BCE_loss = nn.BCELoss()
self.L1_loss = nn.L1Loss()
# setup optimizers #
self.opt_dict["G"] = optim.Adam(
itertools.chain(self.model_dict["G_A"].parameters(),
self.model_dict["G_B"].parameters()),
lr=params['lr_gen'],
betas=(params['beta1'], params['beta2']),
weight_decay=.00001)
self.opt_dict["D_A"] = optim.Adam(self.model_dict["D_A"].parameters(),
lr=params['lr_disc'],
betas=(params['beta1'],
params['beta2']),
weight_decay=.00001)
self.opt_dict["D_B"] = optim.Adam(self.model_dict["D_B"].parameters(),
lr=params['lr_disc'],
betas=(params['beta1'],
params['beta2']),
weight_decay=.00001)
# setup fake image pool #
self.fakeA_pool = helper.ImagePool(50)
self.fakeB_pool = helper.ImagePool(50)
print('Losses and Pools Initialized')
# setup history storage #
self.losses = [
'D_A_feat_loss', 'D_B_feat_loss', 'D_A_loss', 'D_B_loss',
'G_A_loss', 'G_B_loss', 'Cycle_A_loss', 'Cycle_B_loss'
]
self.loss_batch_dict = {}
self.loss_epoch_dict = {}
self.train_hist_dict = {'per_epoch_ptimes': [], 'total_ptime': {}}
for loss in self.losses:
self.train_hist_dict[loss] = []
self.loss_epoch_dict[loss] = []
self.loss_batch_dict[loss] = []
def __init__(self, params):
self.params = params
self.model_dict = {}
self.opt_dict = {}
self.current_epoch = 0
self.current_iter = 0
self.preview_noise = helper.new_random_z(16, params['z_size'], seed=3)
self.transform = load.NormDenorm([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])
self.train_data = load.MountainDataset(params['dataset'],
self.transform,
output_res=params["res"],
perc=params['data_perc'])
self.datalen = self.train_data.__len__()
self.train_loader = torch.utils.data.DataLoader(
self.train_data,
batch_size=params["batch_size"],
num_workers=params["workers"],
shuffle=True,
drop_last=True)
print('Data Loader Initialized: ' + str(self.datalen) + ' Images')
self.model_dict["G"] = n.Generator(
layers=int(math.log(params["res"], 2) - 3),
filts=params["gen_stretch_z_filts"],
max_filts=params["gen_max_filts"],
min_filts=params["gen_min_filts"],
attention=params["attention"])
self.model_dict["D"] = n.Discriminator(
channels=3,
layers=params["disc_layers"],
filts_min=params["disc_min_filts"],
filts=params["disc_max_filts"],
attention=params["attention"])
for i in self.model_dict.keys():
self.model_dict[i].apply(helper.weights_init_normal)
self.model_dict[i].cuda()
self.model_dict[i].train()
print('Networks Initialized')
self.opt_dict["G"] = optim.Adam(self.model_dict["G"].parameters(),
lr=params['lr_gen'])
self.opt_dict["D"] = optim.Adam(self.model_dict["D"].parameters(),
lr=params['lr_disc'])
print('Optimizers Initialized')
# setup history storage #
self.losses = ['G_Loss', 'D_Loss']
self.loss_batch_dict = {}
self.loss_epoch_dict = {}
self.train_hist_dict = {}
for loss in self.losses:
self.train_hist_dict[loss] = []
self.loss_epoch_dict[loss] = []
self.loss_batch_dict[loss] = []