def save_crossover(folder, number=4, scale=1.0):
dataset.clean(folder)
seeds = [scale * get_random_seeds(1) for i in range(number)]
for i in range(len(seeds)):
for j in range(len(seeds)):
im = model.generator.generate(seeds[i], seeds[j], 2).numpy()[0]
dataset.save_image(folder, str(i) + "." + str(j), im)
def main(filename='image.tif', k=32, epochs=10):
img = ds.image2(append_index=True)
centres, clusters = cluster_img(img, k, epochs)
new_img = []
for c, cluster in zip(centres, clusters):
for element in cluster:
if tuple(element[:3]) == ds.WHITE:
new_img.append(element)
else:
new_img.append(c[:3] + tuple(element[-1:]))
new_img = sorted(new_img, key=lambda i: i[-1])
out_img = ds.format_square(new_img)
ds.save_image(filename, out_img)
def train(opt):
lib.print_model_settings(locals().copy())
if 'Attn' in opt.Prediction:
converter = AttnLabelConverter(opt.character)
text_len = opt.batch_max_length+2
else:
converter = CTCLabelConverter(opt.character)
text_len = opt.batch_max_length
opt.classes = converter.character
""" dataset preparation """
if not opt.data_filtering_off:
print('Filtering the images containing characters which are not in opt.character')
print('Filtering the images whose label is longer than opt.batch_max_length')
# see https://github.com/clovaai/deep-text-recognition-benchmark/blob/6593928855fb7abb999a99f428b3e4477d4ae356/dataset.py#L130
log = open(os.path.join(opt.exp_dir,opt.exp_name,'log_dataset.txt'), 'a')
AlignCollate_valid = AlignPairCollate(imgH=opt.imgH, imgW=opt.imgW, keep_ratio_with_pad=opt.PAD)
train_dataset = LmdbStyleDataset(root=opt.train_data, opt=opt)
train_loader = torch.utils.data.DataLoader(
train_dataset, batch_size=opt.batch_size*2, #*2 to sample different images from training encoder and discriminator real images
shuffle=True, # 'True' to check training progress with validation function.
num_workers=int(opt.workers),
collate_fn=AlignCollate_valid, pin_memory=True, drop_last=True)
print('-' * 80)
valid_dataset = LmdbStyleDataset(root=opt.valid_data, opt=opt)
valid_loader = torch.utils.data.DataLoader(
valid_dataset, batch_size=opt.batch_size*2, #*2 to sample different images from training encoder and discriminator real images
shuffle=False, # 'True' to check training progress with validation function.
num_workers=int(opt.workers),
collate_fn=AlignCollate_valid, pin_memory=True, drop_last=True)
print('-' * 80)
log.write('-' * 80 + '\n')
log.close()
text_dataset = text_gen(opt)
text_loader = torch.utils.data.DataLoader(
text_dataset, batch_size=opt.batch_size,
shuffle=True,
num_workers=int(opt.workers),
pin_memory=True, drop_last=True)
opt.num_class = len(converter.character)
c_code_size = opt.latent
cEncoder = GlobalContentEncoder(opt.num_class, text_len, opt.char_embed_size, c_code_size)
ocrModel = ModelV1(opt)
genModel = styleGANGen(opt.size, opt.latent, opt.latent, opt.n_mlp, channel_multiplier=opt.channel_multiplier)
g_ema = styleGANGen(opt.size, opt.latent, opt.latent, opt.n_mlp, channel_multiplier=opt.channel_multiplier)
disEncModel = styleGANDis(opt.size, channel_multiplier=opt.channel_multiplier, input_dim=opt.input_channel, code_s_dim=c_code_size)
accumulate(g_ema, genModel, 0)
# uCriterion = torch.nn.MSELoss()
# sCriterion = torch.nn.MSELoss()
# if opt.contentLoss == 'vis' or opt.contentLoss == 'seq':
# ocrCriterion = torch.nn.L1Loss()
# else:
if 'CTC' in opt.Prediction:
ocrCriterion = torch.nn.CTCLoss(zero_infinity=True).to(device)
else:
print('Not implemented error')
sys.exit()
# ocrCriterion = torch.nn.CrossEntropyLoss(ignore_index=0).to(device) # ignore [GO] token = ignore index 0
cEncoder= torch.nn.DataParallel(cEncoder).to(device)
cEncoder.train()
genModel = torch.nn.DataParallel(genModel).to(device)
g_ema = torch.nn.DataParallel(g_ema).to(device)
genModel.train()
g_ema.eval()
disEncModel = torch.nn.DataParallel(disEncModel).to(device)
disEncModel.train()
ocrModel = torch.nn.DataParallel(ocrModel).to(device)
if opt.ocrFixed:
if opt.Transformation == 'TPS':
ocrModel.module.Transformation.eval()
ocrModel.module.FeatureExtraction.eval()
ocrModel.module.AdaptiveAvgPool.eval()
# ocrModel.module.SequenceModeling.eval()
ocrModel.module.Prediction.eval()
else:
ocrModel.train()
g_reg_ratio = opt.g_reg_every / (opt.g_reg_every + 1)
d_reg_ratio = opt.d_reg_every / (opt.d_reg_every + 1)
optimizer = optim.Adam(
list(genModel.parameters())+list(cEncoder.parameters()),
lr=opt.lr * g_reg_ratio,
betas=(0 ** g_reg_ratio, 0.99 ** g_reg_ratio),
)
dis_optimizer = optim.Adam(
disEncModel.parameters(),
lr=opt.lr * d_reg_ratio,
betas=(0 ** d_reg_ratio, 0.99 ** d_reg_ratio),
)
ocr_optimizer = optim.Adam(
ocrModel.parameters(),
lr=opt.lr,
betas=(0.9, 0.99),
)
## Loading pre-trained files
if opt.modelFolderFlag:
if len(glob.glob(os.path.join(opt.exp_dir,opt.exp_name,"iter_*_synth.pth")))>0:
opt.saved_synth_model = glob.glob(os.path.join(opt.exp_dir,opt.exp_name,"iter_*_synth.pth"))[-1]
if opt.saved_ocr_model !='' and opt.saved_ocr_model !='None':
print(f'loading pretrained ocr model from {opt.saved_ocr_model}')
checkpoint = torch.load(opt.saved_ocr_model)
ocrModel.load_state_dict(checkpoint)
# if opt.saved_gen_model !='' and opt.saved_gen_model !='None':
# print(f'loading pretrained gen model from {opt.saved_gen_model}')
# checkpoint = torch.load(opt.saved_gen_model, map_location=lambda storage, loc: storage)
# genModel.module.load_state_dict(checkpoint['g'])
# g_ema.module.load_state_dict(checkpoint['g_ema'])
if opt.saved_synth_model != '' and opt.saved_synth_model != 'None':
print(f'loading pretrained synth model from {opt.saved_synth_model}')
checkpoint = torch.load(opt.saved_synth_model)
# styleModel.load_state_dict(checkpoint['styleModel'])
# mixModel.load_state_dict(checkpoint['mixModel'])
genModel.load_state_dict(checkpoint['genModel'])
g_ema.load_state_dict(checkpoint['g_ema'])
disEncModel.load_state_dict(checkpoint['disEncModel'])
ocrModel.load_state_dict(checkpoint['ocrModel'])
optimizer.load_state_dict(checkpoint["optimizer"])
dis_optimizer.load_state_dict(checkpoint["dis_optimizer"])
ocr_optimizer.load_state_dict(checkpoint["ocr_optimizer"])
# if opt.imgReconLoss == 'l1':
# recCriterion = torch.nn.L1Loss()
# elif opt.imgReconLoss == 'ssim':
# recCriterion = ssim
# elif opt.imgReconLoss == 'ms-ssim':
# recCriterion = msssim
# loss averager
loss_avg_dis = Averager()
loss_avg_gen = Averager()
loss_avg_unsup = Averager()
loss_avg_sup = Averager()
log_r1_val = Averager()
log_avg_path_loss_val = Averager()
log_avg_mean_path_length_avg = Averager()
log_ada_aug_p = Averager()
loss_avg_ocr_sup = Averager()
loss_avg_ocr_unsup = Averager()
""" final options """
with open(os.path.join(opt.exp_dir,opt.exp_name,'opt.txt'), 'a') as opt_file:
opt_log = '------------ Options -------------\n'
args = vars(opt)
for k, v in args.items():
opt_log += f'{str(k)}: {str(v)}\n'
opt_log += '---------------------------------------\n'
print(opt_log)
opt_file.write(opt_log)
""" start training """
start_iter = 0
if opt.saved_synth_model != '' and opt.saved_synth_model != 'None':
try:
start_iter = int(opt.saved_synth_model.split('_')[-2].split('.')[0])
print(f'continue to train, start_iter: {start_iter}')
except:
pass
#get schedulers
scheduler = get_scheduler(optimizer,opt)
dis_scheduler = get_scheduler(dis_optimizer,opt)
ocr_scheduler = get_scheduler(ocr_optimizer,opt)
start_time = time.time()
iteration = start_iter
cntr=0
mean_path_length = 0
d_loss_val = 0
r1_loss = torch.tensor(0.0, device=device)
g_loss_val = 0
path_loss = torch.tensor(0.0, device=device)
path_lengths = torch.tensor(0.0, device=device)
mean_path_length_avg = 0
# loss_dict = {}
accum = 0.5 ** (32 / (10 * 1000))
ada_augment = torch.tensor([0.0, 0.0], device=device)
ada_aug_p = opt.augment_p if opt.augment_p > 0 else 0.0
ada_aug_step = opt.ada_target / opt.ada_length
r_t_stat = 0
epsilon = 10e-50
# sample_z = torch.randn(opt.n_sample, opt.latent, device=device)
while(True):
# print(cntr)
# train part
if opt.lr_policy !="None":
scheduler.step()
dis_scheduler.step()
ocr_scheduler.step()
image_input_tensors, _, labels, _ = iter(train_loader).next()
labels_z_c = iter(text_loader).next()
image_input_tensors = image_input_tensors.to(device)
gt_image_tensors = image_input_tensors[:opt.batch_size].detach()
real_image_tensors = image_input_tensors[opt.batch_size:].detach()
labels_gt = labels[:opt.batch_size]
requires_grad(cEncoder, False)
requires_grad(genModel, False)
requires_grad(disEncModel, True)
requires_grad(ocrModel, False)
text_z_c, length_z_c = converter.encode(labels_z_c, batch_max_length=opt.batch_max_length)
text_gt, length_gt = converter.encode(labels_gt, batch_max_length=opt.batch_max_length)
z_c_code = cEncoder(text_z_c)
noise_style = mixing_noise_style(opt.batch_size, opt.latent, opt.mixing, device)
style=[]
style.append(noise_style[0]*z_c_code)
if len(noise_style)>1:
style.append(noise_style[1]*z_c_code)
if opt.zAlone:
#to validate orig style gan results
newstyle = []
newstyle.append(style[0][:,:opt.latent])
if len(style)>1:
newstyle.append(style[1][:,:opt.latent])
style = newstyle
fake_img,_ = genModel(style, input_is_latent=opt.input_latent)
# #unsupervised code prediction on generated image
# u_pred_code = disEncModel(fake_img, mode='enc')
# uCost = uCriterion(u_pred_code, z_code)
# #supervised code prediction on gt image
# s_pred_code = disEncModel(gt_image_tensors, mode='enc')
# sCost = uCriterion(s_pred_code, gt_phoc_tensors)
#Domain discriminator
fake_pred = disEncModel(fake_img)
real_pred = disEncModel(real_image_tensors)
disCost = d_logistic_loss(real_pred, fake_pred)
# dis_cost = disCost + opt.gamma_e*uCost + opt.beta*sCost
loss_avg_dis.add(disCost)
# loss_avg_sup.add(opt.beta*sCost)
# loss_avg_unsup.add(opt.gamma_e * uCost)
disEncModel.zero_grad()
disCost.backward()
dis_optimizer.step()
d_regularize = cntr % opt.d_reg_every == 0
if d_regularize:
real_image_tensors.requires_grad = True
real_pred = disEncModel(real_image_tensors)
r1_loss = d_r1_loss(real_pred, real_image_tensors)
disEncModel.zero_grad()
(opt.r1 / 2 * r1_loss * opt.d_reg_every + 0 * real_pred[0]).backward()
dis_optimizer.step()
log_r1_val.add(r1_loss)
# Recognizer update
if not opt.ocrFixed and not opt.zAlone:
requires_grad(disEncModel, False)
requires_grad(ocrModel, True)
if 'CTC' in opt.Prediction:
preds_recon = ocrModel(gt_image_tensors, text_gt, is_train=True)
preds_size = torch.IntTensor([preds_recon.size(1)] * opt.batch_size)
preds_recon_softmax = preds_recon.log_softmax(2).permute(1, 0, 2)
ocrCost = ocrCriterion(preds_recon_softmax, text_gt, preds_size, length_gt)
else:
print("Not implemented error")
sys.exit()
ocrModel.zero_grad()
ocrCost.backward()
# torch.nn.utils.clip_grad_norm_(ocrModel.parameters(), opt.grad_clip) # gradient clipping with 5 (Default)
ocr_optimizer.step()
loss_avg_ocr_sup.add(ocrCost)
else:
loss_avg_ocr_sup.add(torch.tensor(0.0))
# [Word Generator] update
# image_input_tensors, _, labels, _ = iter(train_loader).next()
labels_z_c = iter(text_loader).next()
# image_input_tensors = image_input_tensors.to(device)
# gt_image_tensors = image_input_tensors[:opt.batch_size]
# real_image_tensors = image_input_tensors[opt.batch_size:]
# labels_gt = labels[:opt.batch_size]
requires_grad(cEncoder, True)
requires_grad(genModel, True)
requires_grad(disEncModel, False)
requires_grad(ocrModel, False)
text_z_c, length_z_c = converter.encode(labels_z_c, batch_max_length=opt.batch_max_length)
z_c_code = cEncoder(text_z_c)
noise_style = mixing_noise_style(opt.batch_size, opt.latent, opt.mixing, device)
style=[]
style.append(noise_style[0]*z_c_code)
if len(noise_style)>1:
style.append(noise_style[1]*z_c_code)
if opt.zAlone:
#to validate orig style gan results
newstyle = []
newstyle.append(style[0][:,:opt.latent])
if len(style)>1:
newstyle.append(style[1][:,:opt.latent])
style = newstyle
fake_img,_ = genModel(style, input_is_latent=opt.input_latent)
fake_pred = disEncModel(fake_img)
disGenCost = g_nonsaturating_loss(fake_pred)
if opt.zAlone:
ocrCost = torch.tensor(0.0)
else:
#Compute OCR prediction (Reconstruction of content)
# text_for_pred = torch.LongTensor(opt.batch_size, opt.batch_max_length + 1).fill_(0).to(device)
# length_for_pred = torch.IntTensor([opt.batch_max_length] * opt.batch_size).to(device)
if 'CTC' in opt.Prediction:
preds_recon = ocrModel(fake_img, text_z_c, is_train=False)
preds_size = torch.IntTensor([preds_recon.size(1)] * opt.batch_size)
preds_recon_softmax = preds_recon.log_softmax(2).permute(1, 0, 2)
ocrCost = ocrCriterion(preds_recon_softmax, text_z_c, preds_size, length_z_c)
else:
print("Not implemented error")
sys.exit()
genModel.zero_grad()
cEncoder.zero_grad()
gen_enc_cost = disGenCost + opt.ocrWeight * ocrCost
grad_fake_OCR = torch.autograd.grad(ocrCost, fake_img, retain_graph=True)[0]
loss_grad_fake_OCR = 10**6*torch.mean(grad_fake_OCR**2)
grad_fake_adv = torch.autograd.grad(disGenCost, fake_img, retain_graph=True)[0]
loss_grad_fake_adv = 10**6*torch.mean(grad_fake_adv**2)
if opt.grad_balance:
gen_enc_cost.backward(retain_graph=True)
grad_fake_OCR = torch.autograd.grad(ocrCost, fake_img, create_graph=True, retain_graph=True)[0]
grad_fake_adv = torch.autograd.grad(disGenCost, fake_img, create_graph=True, retain_graph=True)[0]
a = opt.ocrWeight * torch.div(torch.std(grad_fake_adv), epsilon+torch.std(grad_fake_OCR))
if a is None:
print(ocrCost, disGenCost, torch.std(grad_fake_adv), torch.std(grad_fake_OCR))
if a>1000 or a<0.0001:
print(a)
ocrCost = a.detach() * ocrCost
gen_enc_cost = disGenCost + ocrCost
gen_enc_cost.backward(retain_graph=True)
grad_fake_OCR = torch.autograd.grad(ocrCost, fake_img, create_graph=False, retain_graph=True)[0]
grad_fake_adv = torch.autograd.grad(disGenCost, fake_img, create_graph=False, retain_graph=True)[0]
loss_grad_fake_OCR = 10 ** 6 * torch.mean(grad_fake_OCR ** 2)
loss_grad_fake_adv = 10 ** 6 * torch.mean(grad_fake_adv ** 2)
with torch.no_grad():
gen_enc_cost.backward()
else:
gen_enc_cost.backward()
loss_avg_gen.add(disGenCost)
loss_avg_ocr_unsup.add(opt.ocrWeight * ocrCost)
optimizer.step()
g_regularize = cntr % opt.g_reg_every == 0
if g_regularize:
path_batch_size = max(1, opt.batch_size // opt.path_batch_shrink)
# image_input_tensors, _, labels, _ = iter(train_loader).next()
labels_z_c = iter(text_loader).next()
# image_input_tensors = image_input_tensors.to(device)
# gt_image_tensors = image_input_tensors[:path_batch_size]
# labels_gt = labels[:path_batch_size]
text_z_c, length_z_c = converter.encode(labels_z_c[:path_batch_size], batch_max_length=opt.batch_max_length)
# text_gt, length_gt = converter.encode(labels_gt, batch_max_length=opt.batch_max_length)
z_c_code = cEncoder(text_z_c)
noise_style = mixing_noise_style(path_batch_size, opt.latent, opt.mixing, device)
style=[]
style.append(noise_style[0]*z_c_code)
if len(noise_style)>1:
style.append(noise_style[1]*z_c_code)
if opt.zAlone:
#to validate orig style gan results
newstyle = []
newstyle.append(style[0][:,:opt.latent])
if len(style)>1:
newstyle.append(style[1][:,:opt.latent])
style = newstyle
fake_img, grad = genModel(style, return_latents=True, g_path_regularize=True, mean_path_length=mean_path_length)
decay = 0.01
path_lengths = torch.sqrt(grad.pow(2).sum(2).mean(1))
mean_path_length_orig = mean_path_length + decay * (path_lengths.mean() - mean_path_length)
path_loss = (path_lengths - mean_path_length_orig).pow(2).mean()
mean_path_length = mean_path_length_orig.detach().item()
genModel.zero_grad()
cEncoder.zero_grad()
weighted_path_loss = opt.path_regularize * opt.g_reg_every * path_loss
if opt.path_batch_shrink:
weighted_path_loss += 0 * fake_img[0, 0, 0, 0]
weighted_path_loss.backward()
optimizer.step()
# mean_path_length_avg = (
# reduce_sum(mean_path_length).item() / get_world_size()
# )
#commented above for multi-gpu , non-distributed setting
mean_path_length_avg = mean_path_length
accumulate(g_ema, genModel, accum)
log_avg_path_loss_val.add(path_loss)
log_avg_mean_path_length_avg.add(torch.tensor(mean_path_length_avg))
log_ada_aug_p.add(torch.tensor(ada_aug_p))
if get_rank() == 0:
if wandb and opt.wandb:
wandb.log(
{
"Generator": g_loss_val,
"Discriminator": d_loss_val,
"Augment": ada_aug_p,
"Rt": r_t_stat,
"R1": r1_val,
"Path Length Regularization": path_loss_val,
"Mean Path Length": mean_path_length,
"Real Score": real_score_val,
"Fake Score": fake_score_val,
"Path Length": path_length_val,
}
)
# validation part
if (iteration + 1) % opt.valInterval == 0 or iteration == 0: # To see training progress, we also conduct validation when 'iteration == 0'
#generate paired content with similar style
labels_z_c_1 = iter(text_loader).next()
labels_z_c_2 = iter(text_loader).next()
text_z_c_1, length_z_c_1 = converter.encode(labels_z_c_1, batch_max_length=opt.batch_max_length)
text_z_c_2, length_z_c_2 = converter.encode(labels_z_c_2, batch_max_length=opt.batch_max_length)
z_c_code_1 = cEncoder(text_z_c_1)
z_c_code_2 = cEncoder(text_z_c_2)
style_c1_s1 = []
style_c2_s1 = []
style_s1 = mixing_noise_style(opt.batch_size, opt.latent, opt.mixing, device)
style_c1_s1.append(style_s1[0]*z_c_code_1)
style_c2_s1.append(style_s1[0]*z_c_code_2)
if len(style_s1)>1:
style_c1_s1.append(style_s1[1]*z_c_code_1)
style_c2_s1.append(style_s1[1]*z_c_code_2)
noise_style = mixing_noise_style(opt.batch_size, opt.latent, opt.mixing, device)
style_c1_s2 = []
style_c1_s2.append(noise_style[0]*z_c_code_1)
if len(noise_style)>1:
style_c1_s2.append(noise_style[1]*z_c_code_1)
if opt.zAlone:
#to validate orig style gan results
newstyle = []
newstyle.append(style_c1_s1[0][:,:opt.latent])
if len(style_c1_s1)>1:
newstyle.append(style_c1_s1[1][:,:opt.latent])
style_c1_s1 = newstyle
style_c2_s1 = newstyle
style_c1_s2 = newstyle
fake_img_c1_s1, _ = g_ema(style_c1_s1, input_is_latent=opt.input_latent)
fake_img_c2_s1, _ = g_ema(style_c2_s1, input_is_latent=opt.input_latent)
fake_img_c1_s2, _ = g_ema(style_c1_s2, input_is_latent=opt.input_latent)
if not opt.zAlone:
#Run OCR prediction
if 'CTC' in opt.Prediction:
preds = ocrModel(fake_img_c1_s1, text_z_c_1, is_train=False)
preds_size = torch.IntTensor([preds.size(1)] * opt.batch_size)
_, preds_index = preds.max(2)
preds_str_fake_img_c1_s1 = converter.decode(preds_index.data, preds_size.data)
preds = ocrModel(fake_img_c2_s1, text_z_c_2, is_train=False)
preds_size = torch.IntTensor([preds.size(1)] * opt.batch_size)
_, preds_index = preds.max(2)
preds_str_fake_img_c2_s1 = converter.decode(preds_index.data, preds_size.data)
preds = ocrModel(fake_img_c1_s2, text_z_c_1, is_train=False)
preds_size = torch.IntTensor([preds.size(1)] * opt.batch_size)
_, preds_index = preds.max(2)
preds_str_fake_img_c1_s2 = converter.decode(preds_index.data, preds_size.data)
preds = ocrModel(gt_image_tensors, text_gt, is_train=False)
preds_size = torch.IntTensor([preds.size(1)] * gt_image_tensors.shape[0])
_, preds_index = preds.max(2)
preds_str_gt = converter.decode(preds_index.data, preds_size.data)
else:
print("Not implemented error")
sys.exit()
else:
preds_str_fake_img_c1_s1 = [':None:'] * fake_img_c1_s1.shape[0]
preds_str_gt = [':None:'] * fake_img_c1_s1.shape[0]
os.makedirs(os.path.join(opt.trainDir,str(iteration)), exist_ok=True)
for trImgCntr in range(opt.batch_size):
try:
save_image(tensor2im(fake_img_c1_s1[trImgCntr].detach()),os.path.join(opt.trainDir,str(iteration),str(trImgCntr)+'_c1_s1_'+labels_z_c_1[trImgCntr]+'_ocr:'+preds_str_fake_img_c1_s1[trImgCntr]+'.png'))
if not opt.zAlone:
save_image(tensor2im(fake_img_c2_s1[trImgCntr].detach()),os.path.join(opt.trainDir,str(iteration),str(trImgCntr)+'_c2_s1_'+labels_z_c_2[trImgCntr]+'_ocr:'+preds_str_fake_img_c2_s1[trImgCntr]+'.png'))
save_image(tensor2im(fake_img_c1_s2[trImgCntr].detach()),os.path.join(opt.trainDir,str(iteration),str(trImgCntr)+'_c1_s2_'+labels_z_c_1[trImgCntr]+'_ocr:'+preds_str_fake_img_c1_s2[trImgCntr]+'.png'))
if trImgCntr<gt_image_tensors.shape[0]:
save_image(tensor2im(gt_image_tensors[trImgCntr].detach()),os.path.join(opt.trainDir,str(iteration),str(trImgCntr)+'_gt_act:'+labels_gt[trImgCntr]+'_ocr:'+preds_str_gt[trImgCntr]+'.png'))
except:
print('Warning while saving training image')
elapsed_time = time.time() - start_time
# for log
with open(os.path.join(opt.exp_dir,opt.exp_name,'log_train.txt'), 'a') as log:
# training loss and validation loss
loss_log = f'[{iteration+1}/{opt.num_iter}] \
Train Dis loss: {loss_avg_dis.val():0.5f}, Train Gen loss: {loss_avg_gen.val():0.5f},\
Train UnSup OCR loss: {loss_avg_ocr_unsup.val():0.5f}, Train Sup OCR loss: {loss_avg_ocr_sup.val():0.5f}, \
Train R1-val loss: {log_r1_val.val():0.5f}, Train avg-path-loss: {log_avg_path_loss_val.val():0.5f}, \
Train mean-path-length loss: {log_avg_mean_path_length_avg.val():0.5f}, Train ada-aug-p: {log_ada_aug_p.val():0.5f}, \
Elapsed_time: {elapsed_time:0.5f}'
#plotting
lib.plot.plot(os.path.join(opt.plotDir,'Train-Dis-Loss'), loss_avg_dis.val().item())
lib.plot.plot(os.path.join(opt.plotDir,'Train-Gen-Loss'), loss_avg_gen.val().item())
lib.plot.plot(os.path.join(opt.plotDir,'Train-UnSup-OCR-Loss'), loss_avg_ocr_unsup.val().item())
lib.plot.plot(os.path.join(opt.plotDir,'Train-Sup-OCR-Loss'), loss_avg_ocr_sup.val().item())
lib.plot.plot(os.path.join(opt.plotDir,'Train-r1_val'), log_r1_val.val().item())
lib.plot.plot(os.path.join(opt.plotDir,'Train-path_loss_val'), log_avg_path_loss_val.val().item())
lib.plot.plot(os.path.join(opt.plotDir,'Train-mean_path_length_avg'), log_avg_mean_path_length_avg.val().item())
lib.plot.plot(os.path.join(opt.plotDir,'Train-ada_aug_p'), log_ada_aug_p.val().item())
print(loss_log)
loss_avg_dis.reset()
loss_avg_gen.reset()
loss_avg_ocr_unsup.reset()
loss_avg_ocr_sup.reset()
log_r1_val.reset()
log_avg_path_loss_val.reset()
log_avg_mean_path_length_avg.reset()
log_ada_aug_p.reset()
lib.plot.flush()
lib.plot.tick()
# save model per 1e+5 iter.
if (iteration) % 1e+4 == 0:
torch.save({
'cEncoder':cEncoder.state_dict(),
'genModel':genModel.state_dict(),
'g_ema':g_ema.state_dict(),
'ocrModel':ocrModel.state_dict(),
'disEncModel':disEncModel.state_dict(),
'optimizer':optimizer.state_dict(),
'ocr_optimizer':ocr_optimizer.state_dict(),
'dis_optimizer':dis_optimizer.state_dict()},
os.path.join(opt.exp_dir,opt.exp_name,'iter_'+str(iteration+1)+'_synth.pth'))
if (iteration + 1) == opt.num_iter:
print('end the training')
sys.exit()
iteration += 1
cntr+=1
dt_pan_d = gdal.Open(pan_d)
img_mul = dt_mul.ReadAsArray() # (c, h, w)
img_lr = dt_lr.ReadAsArray()
img_lr_u = dt_lr_u.ReadAsArray()
img_pan = dt_pan.ReadAsArray()
img_pan_d = dt_pan_d.ReadAsArray()
XSize = dt_lr.RasterXSize
YSize = dt_lr.RasterYSize
for _ in range(patch_num[num]):
x = random.randint(0, XSize - blk)
y = random.randint(0, YSize - blk)
save_image('%s/%d_mul.tif' % (trainDir, trainCount),
img_mul[:, y * 4:(y + blk) * 4, x * 4:(x + blk) * 4], 4)
save_image('%s/%d_lr_u.tif' % (trainDir, trainCount),
img_lr_u[:, y * 4:(y + blk) * 4, x * 4:(x + blk) * 4], 4)
save_image('%s/%d_lr.tif' % (trainDir, trainCount),
img_lr[:, y:(y + blk), x:(x + blk)], 4)
save_image('%s/%d_pan.tif' % (trainDir, trainCount),
img_pan[y * 4:(y + blk) * 4, x * 4:(x + blk) * 4], 1)
save_image('%s/%d_pan_d.tif' % (trainDir, trainCount),
img_pan_d[y * 4:(y + blk) * 4, x * 4:(x + blk) * 4], 1)
trainCount += 1
print ("done %d" % num)
record.write("%d\n" % trainCount)
record.close()
def train(opt):
lib.print_model_settings(locals().copy())
""" dataset preparation """
if not opt.data_filtering_off:
print('Filtering the images containing characters which are not in opt.character')
print('Filtering the images whose label is longer than opt.batch_max_length')
# see https://github.com/clovaai/deep-text-recognition-benchmark/blob/6593928855fb7abb999a99f428b3e4477d4ae356/dataset.py#L130
log = open(os.path.join(opt.exp_dir,opt.exp_name,'log_dataset.txt'), 'a')
AlignCollate_valid = AlignPairCollate(imgH=opt.imgH, imgW=opt.imgW, keep_ratio_with_pad=opt.PAD)
train_dataset, train_dataset_log = hierarchical_dataset(root=opt.train_data, opt=opt)
train_loader = torch.utils.data.DataLoader(
train_dataset, batch_size=opt.batch_size,
sampler=data_sampler(train_dataset, shuffle=True, distributed=opt.distributed),
num_workers=int(opt.workers),
collate_fn=AlignCollate_valid, pin_memory=True, drop_last=True)
log.write(train_dataset_log)
print('-' * 80)
valid_dataset, valid_dataset_log = hierarchical_dataset(root=opt.valid_data, opt=opt)
valid_loader = torch.utils.data.DataLoader(
valid_dataset, batch_size=opt.batch_size,
sampler=data_sampler(train_dataset, shuffle=False, distributed=opt.distributed),
num_workers=int(opt.workers),
collate_fn=AlignCollate_valid, pin_memory=True, drop_last=True)
log.write(valid_dataset_log)
print('-' * 80)
log.write('-' * 80 + '\n')
log.close()
if 'Attn' in opt.Prediction:
converter = AttnLabelConverter(opt.character)
else:
converter = CTCLabelConverter(opt.character)
opt.num_class = len(converter.character)
# styleModel = StyleTensorEncoder(input_dim=opt.input_channel)
# genModel = AdaIN_Tensor_WordGenerator(opt)
# disModel = MsImageDisV2(opt)
# styleModel = StyleLatentEncoder(input_dim=opt.input_channel, norm='none')
# mixModel = Mixer(opt,nblk=3, dim=opt.latent)
genModel = styleGANGen(opt.size, opt.latent, opt.n_mlp, opt.num_class, channel_multiplier=opt.channel_multiplier).to(device)
disModel = styleGANDis(opt.size, channel_multiplier=opt.channel_multiplier, input_dim=opt.input_channel).to(device)
g_ema = styleGANGen(opt.size, opt.latent, opt.n_mlp, opt.num_class, channel_multiplier=opt.channel_multiplier).to(device)
ocrModel = ModelV1(opt).to(device)
accumulate(g_ema, genModel, 0)
# # weight initialization
# for currModel in [styleModel, mixModel]:
# for name, param in currModel.named_parameters():
# if 'localization_fc2' in name:
# print(f'Skip {name} as it is already initialized')
# continue
# try:
# if 'bias' in name:
# init.constant_(param, 0.0)
# elif 'weight' in name:
# init.kaiming_normal_(param)
# except Exception as e: # for batchnorm.
# if 'weight' in name:
# param.data.fill_(1)
# continue
if opt.contentLoss == 'vis' or opt.contentLoss == 'seq':
ocrCriterion = torch.nn.L1Loss()
else:
if 'CTC' in opt.Prediction:
ocrCriterion = torch.nn.CTCLoss(zero_infinity=True).to(device)
else:
ocrCriterion = torch.nn.CrossEntropyLoss(ignore_index=0).to(device) # ignore [GO] token = ignore index 0
# vggRecCriterion = torch.nn.L1Loss()
# vggModel = VGGPerceptualLossModel(models.vgg19(pretrained=True), vggRecCriterion)
print('model input parameters', opt.imgH, opt.imgW, opt.input_channel, opt.output_channel,
opt.hidden_size, opt.num_class, opt.batch_max_length)
if opt.distributed:
genModel = torch.nn.parallel.DistributedDataParallel(
genModel,
device_ids=[opt.local_rank],
output_device=opt.local_rank,
broadcast_buffers=False,
)
disModel = torch.nn.parallel.DistributedDataParallel(
disModel,
device_ids=[opt.local_rank],
output_device=opt.local_rank,
broadcast_buffers=False,
)
ocrModel = torch.nn.parallel.DistributedDataParallel(
ocrModel,
device_ids=[opt.local_rank],
output_device=opt.local_rank,
broadcast_buffers=False
)
# styleModel = torch.nn.DataParallel(styleModel).to(device)
# styleModel.train()
# mixModel = torch.nn.DataParallel(mixModel).to(device)
# mixModel.train()
# genModel = torch.nn.DataParallel(genModel).to(device)
# g_ema = torch.nn.DataParallel(g_ema).to(device)
genModel.train()
g_ema.eval()
# disModel = torch.nn.DataParallel(disModel).to(device)
disModel.train()
# vggModel = torch.nn.DataParallel(vggModel).to(device)
# vggModel.eval()
# ocrModel = torch.nn.DataParallel(ocrModel).to(device)
# if opt.distributed:
# ocrModel.module.Transformation.eval()
# ocrModel.module.FeatureExtraction.eval()
# ocrModel.module.AdaptiveAvgPool.eval()
# # ocrModel.module.SequenceModeling.eval()
# ocrModel.module.Prediction.eval()
# else:
# ocrModel.Transformation.eval()
# ocrModel.FeatureExtraction.eval()
# ocrModel.AdaptiveAvgPool.eval()
# # ocrModel.SequenceModeling.eval()
# ocrModel.Prediction.eval()
ocrModel.eval()
if opt.distributed:
g_module = genModel.module
d_module = disModel.module
else:
g_module = genModel
d_module = disModel
g_reg_ratio = opt.g_reg_every / (opt.g_reg_every + 1)
d_reg_ratio = opt.d_reg_every / (opt.d_reg_every + 1)
optimizer = optim.Adam(
genModel.parameters(),
lr=opt.lr * g_reg_ratio,
betas=(0 ** g_reg_ratio, 0.99 ** g_reg_ratio),
)
dis_optimizer = optim.Adam(
disModel.parameters(),
lr=opt.lr * d_reg_ratio,
betas=(0 ** d_reg_ratio, 0.99 ** d_reg_ratio),
)
## Loading pre-trained files
if opt.modelFolderFlag:
if len(glob.glob(os.path.join(opt.exp_dir,opt.exp_name,"iter_*_synth.pth")))>0:
opt.saved_synth_model = glob.glob(os.path.join(opt.exp_dir,opt.exp_name,"iter_*_synth.pth"))[-1]
if opt.saved_ocr_model !='' and opt.saved_ocr_model !='None':
if not opt.distributed:
ocrModel = torch.nn.DataParallel(ocrModel)
print(f'loading pretrained ocr model from {opt.saved_ocr_model}')
checkpoint = torch.load(opt.saved_ocr_model)
ocrModel.load_state_dict(checkpoint)
#temporary fix
if not opt.distributed:
ocrModel = ocrModel.module
if opt.saved_gen_model !='' and opt.saved_gen_model !='None':
print(f'loading pretrained gen model from {opt.saved_gen_model}')
checkpoint = torch.load(opt.saved_gen_model, map_location=lambda storage, loc: storage)
genModel.module.load_state_dict(checkpoint['g'])
g_ema.module.load_state_dict(checkpoint['g_ema'])
if opt.saved_synth_model != '' and opt.saved_synth_model != 'None':
print(f'loading pretrained synth model from {opt.saved_synth_model}')
checkpoint = torch.load(opt.saved_synth_model)
# styleModel.load_state_dict(checkpoint['styleModel'])
# mixModel.load_state_dict(checkpoint['mixModel'])
genModel.load_state_dict(checkpoint['genModel'])
g_ema.load_state_dict(checkpoint['g_ema'])
disModel.load_state_dict(checkpoint['disModel'])
optimizer.load_state_dict(checkpoint["optimizer"])
dis_optimizer.load_state_dict(checkpoint["dis_optimizer"])
# if opt.imgReconLoss == 'l1':
# recCriterion = torch.nn.L1Loss()
# elif opt.imgReconLoss == 'ssim':
# recCriterion = ssim
# elif opt.imgReconLoss == 'ms-ssim':
# recCriterion = msssim
# loss averager
loss_avg = Averager()
loss_avg_dis = Averager()
loss_avg_gen = Averager()
loss_avg_imgRecon = Averager()
loss_avg_vgg_per = Averager()
loss_avg_vgg_sty = Averager()
loss_avg_ocr = Averager()
log_r1_val = Averager()
log_avg_path_loss_val = Averager()
log_avg_mean_path_length_avg = Averager()
log_ada_aug_p = Averager()
""" final options """
with open(os.path.join(opt.exp_dir,opt.exp_name,'opt.txt'), 'a') as opt_file:
opt_log = '------------ Options -------------\n'
args = vars(opt)
for k, v in args.items():
opt_log += f'{str(k)}: {str(v)}\n'
opt_log += '---------------------------------------\n'
print(opt_log)
opt_file.write(opt_log)
""" start training """
start_iter = 0
if opt.saved_synth_model != '' and opt.saved_synth_model != 'None':
try:
start_iter = int(opt.saved_synth_model.split('_')[-2].split('.')[0])
print(f'continue to train, start_iter: {start_iter}')
except:
pass
#get schedulers
scheduler = get_scheduler(optimizer,opt)
dis_scheduler = get_scheduler(dis_optimizer,opt)
start_time = time.time()
iteration = start_iter
cntr=0
mean_path_length = 0
d_loss_val = 0
r1_loss = torch.tensor(0.0, device=device)
g_loss_val = 0
path_loss = torch.tensor(0.0, device=device)
path_lengths = torch.tensor(0.0, device=device)
mean_path_length_avg = 0
loss_dict = {}
accum = 0.5 ** (32 / (10 * 1000))
ada_augment = torch.tensor([0.0, 0.0], device=device)
ada_aug_p = opt.augment_p if opt.augment_p > 0 else 0.0
ada_aug_step = opt.ada_target / opt.ada_length
r_t_stat = 0
sample_z = torch.randn(opt.n_sample, opt.latent, device=device)
while(True):
# print(cntr)
# train part
if opt.lr_policy !="None":
scheduler.step()
dis_scheduler.step()
image_input_tensors, image_gt_tensors, labels_1, labels_2 = iter(train_loader).next()
image_input_tensors = image_input_tensors.to(device)
image_gt_tensors = image_gt_tensors.to(device)
batch_size = image_input_tensors.size(0)
requires_grad(genModel, False)
# requires_grad(styleModel, False)
# requires_grad(mixModel, False)
requires_grad(disModel, True)
text_1, length_1 = converter.encode(labels_1, batch_max_length=opt.batch_max_length)
text_2, length_2 = converter.encode(labels_2, batch_max_length=opt.batch_max_length)
#forward pass from style and word generator
# style = styleModel(image_input_tensors).squeeze(2).squeeze(2)
style = mixing_noise(opt.batch_size, opt.latent, opt.mixing, device)
# scInput = mixModel(style,text_2)
if 'CTC' in opt.Prediction:
images_recon_2,_ = genModel(style, text_2, input_is_latent=opt.input_latent)
else:
images_recon_2,_ = genModel(style, text_2[:,1:-1], input_is_latent=opt.input_latent)
#Domain discriminator: Dis update
if opt.augment:
image_gt_tensors_aug, _ = augment(image_gt_tensors, ada_aug_p)
images_recon_2, _ = augment(images_recon_2, ada_aug_p)
else:
image_gt_tensors_aug = image_gt_tensors
fake_pred = disModel(images_recon_2)
real_pred = disModel(image_gt_tensors_aug)
disCost = d_logistic_loss(real_pred, fake_pred)
loss_dict["d"] = disCost*opt.disWeight
loss_dict["real_score"] = real_pred.mean()
loss_dict["fake_score"] = fake_pred.mean()
loss_avg_dis.add(disCost)
disModel.zero_grad()
disCost.backward()
dis_optimizer.step()
if opt.augment and opt.augment_p == 0:
ada_augment += torch.tensor(
(torch.sign(real_pred).sum().item(), real_pred.shape[0]), device=device
)
ada_augment = reduce_sum(ada_augment)
if ada_augment[1] > 255:
pred_signs, n_pred = ada_augment.tolist()
r_t_stat = pred_signs / n_pred
if r_t_stat > opt.ada_target:
sign = 1
else:
sign = -1
ada_aug_p += sign * ada_aug_step * n_pred
ada_aug_p = min(1, max(0, ada_aug_p))
ada_augment.mul_(0)
d_regularize = cntr % opt.d_reg_every == 0
if d_regularize:
image_gt_tensors.requires_grad = True
image_input_tensors.requires_grad = True
cat_tensor = image_gt_tensors
real_pred = disModel(cat_tensor)
r1_loss = d_r1_loss(real_pred, cat_tensor)
disModel.zero_grad()
(opt.r1 / 2 * r1_loss * opt.d_reg_every + 0 * real_pred[0]).backward()
dis_optimizer.step()
loss_dict["r1"] = r1_loss
# #[Style Encoder] + [Word Generator] update
image_input_tensors, image_gt_tensors, labels_1, labels_2 = iter(train_loader).next()
image_input_tensors = image_input_tensors.to(device)
image_gt_tensors = image_gt_tensors.to(device)
batch_size = image_input_tensors.size(0)
requires_grad(genModel, True)
# requires_grad(styleModel, True)
# requires_grad(mixModel, True)
requires_grad(disModel, False)
text_1, length_1 = converter.encode(labels_1, batch_max_length=opt.batch_max_length)
text_2, length_2 = converter.encode(labels_2, batch_max_length=opt.batch_max_length)
# style = styleModel(image_input_tensors).squeeze(2).squeeze(2)
# scInput = mixModel(style,text_2)
# images_recon_2,_ = genModel([scInput], input_is_latent=opt.input_latent)
style = mixing_noise(batch_size, opt.latent, opt.mixing, device)
if 'CTC' in opt.Prediction:
images_recon_2, _ = genModel(style, text_2)
else:
images_recon_2, _ = genModel(style, text_2[:,1:-1])
if opt.augment:
images_recon_2, _ = augment(images_recon_2, ada_aug_p)
fake_pred = disModel(images_recon_2)
disGenCost = g_nonsaturating_loss(fake_pred)
loss_dict["g"] = disGenCost
# # #Adversarial loss
# # disGenCost = disModel.module.calc_gen_loss(torch.cat((images_recon_2,image_input_tensors),dim=1))
# #Input reconstruction loss
# recCost = recCriterion(images_recon_2,image_gt_tensors)
# #vgg loss
# vggPerCost, vggStyleCost = vggModel(image_gt_tensors, images_recon_2)
#ocr loss
text_for_pred = torch.LongTensor(batch_size, opt.batch_max_length + 1).fill_(0).to(device)
length_for_pred = torch.IntTensor([opt.batch_max_length] * batch_size).to(device)
if opt.contentLoss == 'vis' or opt.contentLoss == 'seq':
preds_recon = ocrModel(images_recon_2, text_for_pred, is_train=False, returnFeat=opt.contentLoss)
preds_gt = ocrModel(image_gt_tensors, text_for_pred, is_train=False, returnFeat=opt.contentLoss)
ocrCost = ocrCriterion(preds_recon, preds_gt)
else:
if 'CTC' in opt.Prediction:
preds_recon = ocrModel(images_recon_2, text_for_pred, is_train=False)
# preds_o = preds_recon[:, :text_1.shape[1], :]
preds_size = torch.IntTensor([preds_recon.size(1)] * batch_size)
preds_recon_softmax = preds_recon.log_softmax(2).permute(1, 0, 2)
ocrCost = ocrCriterion(preds_recon_softmax, text_2, preds_size, length_2)
#predict ocr recognition on generated images
# preds_recon_size = torch.IntTensor([preds_recon.size(1)] * batch_size)
_, preds_recon_index = preds_recon.max(2)
labels_o_ocr = converter.decode(preds_recon_index.data, preds_size.data)
#predict ocr recognition on gt style images
preds_s = ocrModel(image_input_tensors, text_for_pred, is_train=False)
# preds_s = preds_s[:, :text_1.shape[1] - 1, :]
preds_s_size = torch.IntTensor([preds_s.size(1)] * batch_size)
_, preds_s_index = preds_s.max(2)
labels_s_ocr = converter.decode(preds_s_index.data, preds_s_size.data)
#predict ocr recognition on gt stylecontent images
preds_sc = ocrModel(image_gt_tensors, text_for_pred, is_train=False)
# preds_sc = preds_sc[:, :text_2.shape[1] - 1, :]
preds_sc_size = torch.IntTensor([preds_sc.size(1)] * batch_size)
_, preds_sc_index = preds_sc.max(2)
labels_sc_ocr = converter.decode(preds_sc_index.data, preds_sc_size.data)
else:
preds_recon = ocrModel(images_recon_2, text_for_pred[:, :-1], is_train=False) # align with Attention.forward
target_2 = text_2[:, 1:] # without [GO] Symbol
ocrCost = ocrCriterion(preds_recon.view(-1, preds_recon.shape[-1]), target_2.contiguous().view(-1))
#predict ocr recognition on generated images
_, preds_o_index = preds_recon.max(2)
labels_o_ocr = converter.decode(preds_o_index, length_for_pred)
for idx, pred in enumerate(labels_o_ocr):
pred_EOS = pred.find('[s]')
labels_o_ocr[idx] = pred[:pred_EOS] # prune after "end of sentence" token ([s])
#predict ocr recognition on gt style images
preds_s = ocrModel(image_input_tensors, text_for_pred, is_train=False)
_, preds_s_index = preds_s.max(2)
labels_s_ocr = converter.decode(preds_s_index, length_for_pred)
for idx, pred in enumerate(labels_s_ocr):
pred_EOS = pred.find('[s]')
labels_s_ocr[idx] = pred[:pred_EOS] # prune after "end of sentence" token ([s])
#predict ocr recognition on gt stylecontent images
preds_sc = ocrModel(image_gt_tensors, text_for_pred, is_train=False)
_, preds_sc_index = preds_sc.max(2)
labels_sc_ocr = converter.decode(preds_sc_index, length_for_pred)
for idx, pred in enumerate(labels_sc_ocr):
pred_EOS = pred.find('[s]')
labels_sc_ocr[idx] = pred[:pred_EOS] # prune after "end of sentence" token ([s])
# cost = opt.reconWeight*recCost + opt.disWeight*disGenCost + opt.vggPerWeight*vggPerCost + opt.vggStyWeight*vggStyleCost + opt.ocrWeight*ocrCost
cost = opt.disWeight*disGenCost + opt.ocrWeight*ocrCost
# styleModel.zero_grad()
genModel.zero_grad()
# mixModel.zero_grad()
disModel.zero_grad()
# vggModel.zero_grad()
ocrModel.zero_grad()
cost.backward()
optimizer.step()
loss_avg.add(cost)
g_regularize = cntr % opt.g_reg_every == 0
if g_regularize:
image_input_tensors, image_gt_tensors, labels_1, labels_2 = iter(train_loader).next()
image_input_tensors = image_input_tensors.to(device)
image_gt_tensors = image_gt_tensors.to(device)
batch_size = image_input_tensors.size(0)
text_1, length_1 = converter.encode(labels_1, batch_max_length=opt.batch_max_length)
text_2, length_2 = converter.encode(labels_2, batch_max_length=opt.batch_max_length)
path_batch_size = max(1, batch_size // opt.path_batch_shrink)
# style = styleModel(image_input_tensors).squeeze(2).squeeze(2)
# scInput = mixModel(style,text_2)
# images_recon_2, latents = genModel([scInput],input_is_latent=opt.input_latent, return_latents=True)
style = mixing_noise(path_batch_size, opt.latent, opt.mixing, device)
if 'CTC' in opt.Prediction:
images_recon_2, latents = genModel(style, text_2[:path_batch_size], return_latents=True)
else:
images_recon_2, latents = genModel(style, text_2[:path_batch_size,1:-1], return_latents=True)
path_loss, mean_path_length, path_lengths = g_path_regularize(
images_recon_2, latents, mean_path_length
)
genModel.zero_grad()
weighted_path_loss = opt.path_regularize * opt.g_reg_every * path_loss
if opt.path_batch_shrink:
weighted_path_loss += 0 * images_recon_2[0, 0, 0, 0]
weighted_path_loss.backward()
optimizer.step()
mean_path_length_avg = (
reduce_sum(mean_path_length).item() / get_world_size()
)
loss_dict["path"] = path_loss
loss_dict["path_length"] = path_lengths.mean()
accumulate(g_ema, g_module, accum)
loss_reduced = reduce_loss_dict(loss_dict)
d_loss_val = loss_reduced["d"].mean().item()
g_loss_val = loss_reduced["g"].mean().item()
r1_val = loss_reduced["r1"].mean().item()
path_loss_val = loss_reduced["path"].mean().item()
real_score_val = loss_reduced["real_score"].mean().item()
fake_score_val = loss_reduced["fake_score"].mean().item()
path_length_val = loss_reduced["path_length"].mean().item()
#Individual losses
loss_avg_gen.add(opt.disWeight*disGenCost)
loss_avg_imgRecon.add(torch.tensor(0.0))
loss_avg_vgg_per.add(torch.tensor(0.0))
loss_avg_vgg_sty.add(torch.tensor(0.0))
loss_avg_ocr.add(opt.ocrWeight*ocrCost)
log_r1_val.add(loss_reduced["path"])
log_avg_path_loss_val.add(loss_reduced["path"])
log_avg_mean_path_length_avg.add(torch.tensor(mean_path_length_avg))
log_ada_aug_p.add(torch.tensor(ada_aug_p))
if get_rank() == 0:
# pbar.set_description(
# (
# f"d: {d_loss_val:.4f}; g: {g_loss_val:.4f}; r1: {r1_val:.4f}; "
# f"path: {path_loss_val:.4f}; mean path: {mean_path_length_avg:.4f}; "
# f"augment: {ada_aug_p:.4f}"
# )
# )
if wandb and opt.wandb:
wandb.log(
{
"Generator": g_loss_val,
"Discriminator": d_loss_val,
"Augment": ada_aug_p,
"Rt": r_t_stat,
"R1": r1_val,
"Path Length Regularization": path_loss_val,
"Mean Path Length": mean_path_length,
"Real Score": real_score_val,
"Fake Score": fake_score_val,
"Path Length": path_length_val,
}
)
# if cntr % 100 == 0:
# with torch.no_grad():
# g_ema.eval()
# sample, _ = g_ema([scInput[:,:opt.latent],scInput[:,opt.latent:]])
# utils.save_image(
# sample,
# os.path.join(opt.trainDir, f"sample_{str(cntr).zfill(6)}.png"),
# nrow=int(opt.n_sample ** 0.5),
# normalize=True,
# range=(-1, 1),
# )
# validation part
if (iteration + 1) % opt.valInterval == 0 or iteration == 0: # To see training progress, we also conduct validation when 'iteration == 0'
#Save training images
curr_batch_size = style[0].shape[0]
images_recon_2, _ = g_ema(style, text_2[:curr_batch_size], input_is_latent=opt.input_latent)
os.makedirs(os.path.join(opt.trainDir,str(iteration)), exist_ok=True)
for trImgCntr in range(batch_size):
try:
if opt.contentLoss == 'vis' or opt.contentLoss == 'seq':
save_image(tensor2im(image_input_tensors[trImgCntr].detach()),os.path.join(opt.trainDir,str(iteration),str(trImgCntr)+'_sInput_'+labels_1[trImgCntr]+'.png'))
save_image(tensor2im(image_gt_tensors[trImgCntr].detach()),os.path.join(opt.trainDir,str(iteration),str(trImgCntr)+'_csGT_'+labels_2[trImgCntr]+'.png'))
save_image(tensor2im(images_recon_2[trImgCntr].detach()),os.path.join(opt.trainDir,str(iteration),str(trImgCntr)+'_csRecon_'+labels_2[trImgCntr]+'.png'))
else:
save_image(tensor2im(image_input_tensors[trImgCntr].detach()),os.path.join(opt.trainDir,str(iteration),str(trImgCntr)+'_sInput_'+labels_1[trImgCntr]+'_'+labels_s_ocr[trImgCntr]+'.png'))
save_image(tensor2im(image_gt_tensors[trImgCntr].detach()),os.path.join(opt.trainDir,str(iteration),str(trImgCntr)+'_csGT_'+labels_2[trImgCntr]+'_'+labels_sc_ocr[trImgCntr]+'.png'))
save_image(tensor2im(images_recon_2[trImgCntr].detach()),os.path.join(opt.trainDir,str(iteration),str(trImgCntr)+'_csRecon_'+labels_2[trImgCntr]+'_'+labels_o_ocr[trImgCntr]+'.png'))
except:
print('Warning while saving training image')
elapsed_time = time.time() - start_time
# for log
with open(os.path.join(opt.exp_dir,opt.exp_name,'log_train.txt'), 'a') as log:
# styleModel.eval()
genModel.eval()
g_ema.eval()
# mixModel.eval()
disModel.eval()
with torch.no_grad():
valid_loss, infer_time, length_of_data = validation_synth_v6(
iteration, g_ema, ocrModel, disModel, ocrCriterion, valid_loader, converter, opt)
# styleModel.train()
genModel.train()
# mixModel.train()
disModel.train()
# training loss and validation loss
loss_log = f'[{iteration+1}/{opt.num_iter}] Train Synth loss: {loss_avg.val():0.5f}, \
Train Dis loss: {loss_avg_dis.val():0.5f}, Train Gen loss: {loss_avg_gen.val():0.5f},\
Train OCR loss: {loss_avg_ocr.val():0.5f}, \
Train R1-val loss: {log_r1_val.val():0.5f}, Train avg-path-loss: {log_avg_path_loss_val.val():0.5f}, \
Train mean-path-length loss: {log_avg_mean_path_length_avg.val():0.5f}, Train ada-aug-p: {log_ada_aug_p.val():0.5f}, \
Valid Synth loss: {valid_loss[0]:0.5f}, \
Valid Dis loss: {valid_loss[1]:0.5f}, Valid Gen loss: {valid_loss[2]:0.5f}, \
Valid OCR loss: {valid_loss[6]:0.5f}, Elapsed_time: {elapsed_time:0.5f}'
#plotting
lib.plot.plot(os.path.join(opt.plotDir,'Train-Synth-Loss'), loss_avg.val().item())
lib.plot.plot(os.path.join(opt.plotDir,'Train-Dis-Loss'), loss_avg_dis.val().item())
lib.plot.plot(os.path.join(opt.plotDir,'Train-Gen-Loss'), loss_avg_gen.val().item())
# lib.plot.plot(os.path.join(opt.plotDir,'Train-ImgRecon1-Loss'), loss_avg_imgRecon.val().item())
# lib.plot.plot(os.path.join(opt.plotDir,'Train-VGG-Per-Loss'), loss_avg_vgg_per.val().item())
# lib.plot.plot(os.path.join(opt.plotDir,'Train-VGG-Sty-Loss'), loss_avg_vgg_sty.val().item())
lib.plot.plot(os.path.join(opt.plotDir,'Train-OCR-Loss'), loss_avg_ocr.val().item())
lib.plot.plot(os.path.join(opt.plotDir,'Train-r1_val'), log_r1_val.val().item())
lib.plot.plot(os.path.join(opt.plotDir,'Train-path_loss_val'), log_avg_path_loss_val.val().item())
lib.plot.plot(os.path.join(opt.plotDir,'Train-mean_path_length_avg'), log_avg_mean_path_length_avg.val().item())
lib.plot.plot(os.path.join(opt.plotDir,'Train-ada_aug_p'), log_ada_aug_p.val().item())
lib.plot.plot(os.path.join(opt.plotDir,'Valid-Synth-Loss'), valid_loss[0].item())
lib.plot.plot(os.path.join(opt.plotDir,'Valid-Dis-Loss'), valid_loss[1].item())
lib.plot.plot(os.path.join(opt.plotDir,'Valid-Gen-Loss'), valid_loss[2].item())
# lib.plot.plot(os.path.join(opt.plotDir,'Valid-ImgRecon1-Loss'), valid_loss[3].item())
# lib.plot.plot(os.path.join(opt.plotDir,'Valid-VGG-Per-Loss'), valid_loss[4].item())
# lib.plot.plot(os.path.join(opt.plotDir,'Valid-VGG-Sty-Loss'), valid_loss[5].item())
lib.plot.plot(os.path.join(opt.plotDir,'Valid-OCR-Loss'), valid_loss[6].item())
print(loss_log)
loss_avg.reset()
loss_avg_dis.reset()
loss_avg_gen.reset()
loss_avg_imgRecon.reset()
loss_avg_vgg_per.reset()
loss_avg_vgg_sty.reset()
loss_avg_ocr.reset()
log_r1_val.reset()
log_avg_path_loss_val.reset()
log_avg_mean_path_length_avg.reset()
log_ada_aug_p.reset()
lib.plot.flush()
lib.plot.tick()
# save model per 1e+5 iter.
if (iteration) % 1e+4 == 0:
torch.save({
# 'styleModel':styleModel.state_dict(),
# 'mixModel':mixModel.state_dict(),
'genModel':g_module.state_dict(),
'g_ema':g_ema.state_dict(),
'disModel':d_module.state_dict(),
'optimizer':optimizer.state_dict(),
'dis_optimizer':dis_optimizer.state_dict()},
os.path.join(opt.exp_dir,opt.exp_name,'iter_'+str(iteration+1)+'_synth.pth'))
if (iteration + 1) == opt.num_iter:
print('end the training')
sys.exit()
iteration += 1
cntr+=1
def train(opt):
plotDir = os.path.join(opt.exp_dir, opt.exp_name, 'plots')
if not os.path.exists(plotDir):
os.makedirs(plotDir)
lib.print_model_settings(locals().copy())
""" dataset preparation """
if not opt.data_filtering_off:
print(
'Filtering the images containing characters which are not in opt.character'
)
print(
'Filtering the images whose label is longer than opt.batch_max_length'
)
# see https://github.com/clovaai/deep-text-recognition-benchmark/blob/6593928855fb7abb999a99f428b3e4477d4ae356/dataset.py#L130
opt.select_data = opt.select_data.split('-')
opt.batch_ratio = opt.batch_ratio.split('-')
log = open(os.path.join(opt.exp_dir, opt.exp_name, 'log_dataset.txt'), 'a')
AlignCollate_valid = AlignPairCollate(imgH=opt.imgH,
imgW=opt.imgW,
keep_ratio_with_pad=opt.PAD)
train_dataset, train_dataset_log = hierarchical_dataset(
root=opt.train_data, opt=opt)
train_loader = torch.utils.data.DataLoader(
train_dataset,
batch_size=opt.batch_size,
shuffle=
True, # 'True' to check training progress with validation function.
num_workers=int(opt.workers),
collate_fn=AlignCollate_valid,
pin_memory=True)
log.write(train_dataset_log)
print('-' * 80)
valid_dataset, valid_dataset_log = hierarchical_dataset(
root=opt.valid_data, opt=opt)
valid_loader = torch.utils.data.DataLoader(
valid_dataset,
batch_size=opt.batch_size,
shuffle=
False, # 'True' to check training progress with validation function.
num_workers=int(opt.workers),
collate_fn=AlignCollate_valid,
pin_memory=True)
log.write(valid_dataset_log)
print('-' * 80)
log.write('-' * 80 + '\n')
log.close()
converter = CTCLabelConverter(opt.character)
opt.num_class = len(converter.character)
if opt.rgb:
opt.input_channel = 3
styleModel = StyleTensorEncoder(input_dim=opt.input_channel)
genModel = AdaIN_Tensor_WordGenerator(opt)
disModel = MsImageDisV2(opt)
vggRecCriterion = torch.nn.L1Loss()
vggModel = VGGPerceptualLossModel(models.vgg19(pretrained=True),
vggRecCriterion)
print('model input parameters', opt.imgH, opt.imgW, opt.input_channel,
opt.output_channel, opt.hidden_size, opt.num_class,
opt.batch_max_length)
# weight initialization
for currModel in [styleModel, genModel, disModel]:
for name, param in currModel.named_parameters():
if 'localization_fc2' in name:
print(f'Skip {name} as it is already initialized')
continue
try:
if 'bias' in name:
init.constant_(param, 0.0)
elif 'weight' in name:
init.kaiming_normal_(param)
except Exception as e: # for batchnorm.
if 'weight' in name:
param.data.fill_(1)
continue
styleModel = torch.nn.DataParallel(styleModel).to(device)
styleModel.train()
genModel = torch.nn.DataParallel(genModel).to(device)
genModel.train()
disModel = torch.nn.DataParallel(disModel).to(device)
disModel.train()
vggModel = torch.nn.DataParallel(vggModel).to(device)
vggModel.eval()
if opt.modelFolderFlag:
if len(
glob.glob(
os.path.join(opt.exp_dir, opt.exp_name,
"iter_*_synth.pth"))) > 0:
opt.saved_synth_model = glob.glob(
os.path.join(opt.exp_dir, opt.exp_name,
"iter_*_synth.pth"))[-1]
if opt.saved_synth_model != '' and opt.saved_synth_model != 'None':
print(f'loading pretrained synth model from {opt.saved_synth_model}')
checkpoint = torch.load(opt.saved_synth_model)
styleModel.load_state_dict(checkpoint['styleModel'])
genModel.load_state_dict(checkpoint['genModel'])
disModel.load_state_dict(checkpoint['disModel'])
if opt.imgReconLoss == 'l1':
recCriterion = torch.nn.L1Loss()
elif opt.imgReconLoss == 'ssim':
recCriterion = ssim
elif opt.imgReconLoss == 'ms-ssim':
recCriterion = msssim
if opt.styleLoss == 'l1':
styleRecCriterion = torch.nn.L1Loss()
elif opt.styleLoss == 'triplet':
styleRecCriterion = torch.nn.TripletMarginLoss(
margin=opt.tripletMargin, p=1)
#for validation; check only positive pairs
styleTestRecCriterion = torch.nn.L1Loss()
# loss averager
loss_avg = Averager()
loss_avg_dis = Averager()
loss_avg_gen = Averager()
loss_avg_imgRecon = Averager()
loss_avg_vgg_per = Averager()
loss_avg_vgg_sty = Averager()
##---------------------------------------##
# filter that only require gradient decent
filtered_parameters = []
params_num = []
for p in filter(lambda p: p.requires_grad, styleModel.parameters()):
filtered_parameters.append(p)
params_num.append(np.prod(p.size()))
for p in filter(lambda p: p.requires_grad, genModel.parameters()):
filtered_parameters.append(p)
params_num.append(np.prod(p.size()))
print('Trainable style and generator params num : ', sum(params_num))
# setup optimizer
if opt.optim == 'adam':
optimizer = optim.Adam(filtered_parameters,
lr=opt.lr,
betas=(opt.beta1, opt.beta2),
weight_decay=opt.weight_decay)
else:
optimizer = optim.Adadelta(filtered_parameters,
lr=opt.lr,
rho=opt.rho,
eps=opt.eps,
weight_decay=opt.weight_decay)
print("SynthOptimizer:")
print(optimizer)
#filter parameters for Dis training
dis_filtered_parameters = []
dis_params_num = []
for p in filter(lambda p: p.requires_grad, disModel.parameters()):
dis_filtered_parameters.append(p)
dis_params_num.append(np.prod(p.size()))
print('Dis Trainable params num : ', sum(dis_params_num))
# setup optimizer
if opt.optim == 'adam':
dis_optimizer = optim.Adam(dis_filtered_parameters,
lr=opt.lr,
betas=(opt.beta1, opt.beta2),
weight_decay=opt.weight_decay)
else:
dis_optimizer = optim.Adadelta(dis_filtered_parameters,
lr=opt.lr,
rho=opt.rho,
eps=opt.eps,
weight_decay=opt.weight_decay)
print("DisOptimizer:")
print(dis_optimizer)
##---------------------------------------##
""" final options """
with open(os.path.join(opt.exp_dir, opt.exp_name, 'opt.txt'),
'a') as opt_file:
opt_log = '------------ Options -------------\n'
args = vars(opt)
for k, v in args.items():
opt_log += f'{str(k)}: {str(v)}\n'
opt_log += '---------------------------------------\n'
print(opt_log)
opt_file.write(opt_log)
""" start training """
start_iter = 0
if opt.saved_synth_model != '' and opt.saved_synth_model != 'None':
try:
start_iter = int(
opt.saved_synth_model.split('_')[-2].split('.')[0])
print(f'continue to train, start_iter: {start_iter}')
except:
pass
#get schedulers
scheduler = get_scheduler(optimizer, opt)
dis_scheduler = get_scheduler(dis_optimizer, opt)
start_time = time.time()
iteration = start_iter
cntr = 0
while (True):
# train part
if opt.lr_policy != "None":
scheduler.step()
dis_scheduler.step()
image_input_tensors, image_gt_tensors, labels_1, labels_2 = iter(
train_loader).next()
cntr += 1
image_input_tensors = image_input_tensors.to(device)
image_gt_tensors = image_gt_tensors.to(device)
batch_size = image_input_tensors.size(0)
text_2, length_2 = converter.encode(
labels_2, batch_max_length=opt.batch_max_length)
#forward pass from style and word generator
style = styleModel(image_input_tensors)
images_recon_2 = genModel(style, text_2)
#Domain discriminator: Dis update
disModel.zero_grad()
disCost = opt.disWeight * (disModel.module.calc_dis_loss(
torch.cat((images_recon_2.detach(), image_input_tensors), dim=1),
torch.cat((image_gt_tensors, image_input_tensors), dim=1)))
disCost.backward()
dis_optimizer.step()
loss_avg_dis.add(disCost)
# #[Style Encoder] + [Word Generator] update
#Adversarial loss
disGenCost = disModel.module.calc_gen_loss(
torch.cat((images_recon_2, image_input_tensors), dim=1))
#Input reconstruction loss
recCost = recCriterion(images_recon_2, image_gt_tensors)
#vgg loss
vggPerCost, vggStyleCost = vggModel(image_gt_tensors, images_recon_2)
cost = opt.reconWeight * recCost + opt.disWeight * disGenCost + opt.vggPerWeight * vggPerCost + opt.vggStyWeight * vggStyleCost
styleModel.zero_grad()
genModel.zero_grad()
disModel.zero_grad()
vggModel.zero_grad()
cost.backward()
optimizer.step()
loss_avg.add(cost)
#Individual losses
loss_avg_gen.add(opt.disWeight * disGenCost)
loss_avg_imgRecon.add(opt.reconWeight * recCost)
loss_avg_vgg_per.add(opt.vggPerWeight * vggPerCost)
loss_avg_vgg_sty.add(opt.vggStyWeight * vggStyleCost)
# validation part
if (
iteration + 1
) % opt.valInterval == 0 or iteration == 0: # To see training progress, we also conduct validation when 'iteration == 0'
#Save training images
os.makedirs(os.path.join(opt.exp_dir, opt.exp_name, 'trainImages',
str(iteration)),
exist_ok=True)
for trImgCntr in range(batch_size):
try:
save_image(
tensor2im(image_input_tensors[trImgCntr].detach()),
os.path.join(
opt.exp_dir, opt.exp_name, 'trainImages',
str(iteration),
str(trImgCntr) + '_sInput_' + labels_1[trImgCntr] +
'.png'))
save_image(
tensor2im(image_gt_tensors[trImgCntr].detach()),
os.path.join(
opt.exp_dir, opt.exp_name, 'trainImages',
str(iteration),
str(trImgCntr) + '_csGT_' + labels_2[trImgCntr] +
'.png'))
save_image(
tensor2im(images_recon_2[trImgCntr].detach()),
os.path.join(
opt.exp_dir, opt.exp_name, 'trainImages',
str(iteration),
str(trImgCntr) + '_csRecon_' +
labels_2[trImgCntr] + '.png'))
except:
print('Warning while saving training image')
elapsed_time = time.time() - start_time
# for log
with open(os.path.join(opt.exp_dir, opt.exp_name, 'log_train.txt'),
'a') as log:
styleModel.eval()
genModel.eval()
disModel.eval()
with torch.no_grad():
valid_loss, infer_time, length_of_data = validation_synth_v3(
iteration, styleModel, genModel, vggModel, disModel,
recCriterion, valid_loader, converter, opt)
styleModel.train()
genModel.train()
disModel.train()
# training loss and validation loss
loss_log = f'[{iteration+1}/{opt.num_iter}] Train Synth loss: {loss_avg.val():0.5f}, \
Train Dis loss: {loss_avg_dis.val():0.5f}, Train Gen loss: {loss_avg_gen.val():0.5f},\
Train ImgRecon loss: {loss_avg_imgRecon.val():0.5f}, Train VGG-Per loss: {loss_avg_vgg_per.val():0.5f},\
Train VGG-Sty loss: {loss_avg_vgg_sty.val():0.5f}, Valid Synth loss: {valid_loss[1]:0.5f}, \
Valid Dis loss: {valid_loss[2]:0.5f}, Elapsed_time: {elapsed_time:0.5f}'
#plotting
lib.plot.plot(os.path.join(plotDir, 'Train-Synth-Loss'),
loss_avg.val().item())
lib.plot.plot(os.path.join(plotDir, 'Train-Dis-Loss'),
loss_avg_dis.val().item())
lib.plot.plot(os.path.join(plotDir, 'Train-Gen-Loss'),
loss_avg_gen.val().item())
lib.plot.plot(os.path.join(plotDir, 'Train-ImgRecon1-Loss'),
loss_avg_imgRecon.val().item())
lib.plot.plot(os.path.join(plotDir, 'Train-VGG-Per-Loss'),
loss_avg_vgg_per.val().item())
lib.plot.plot(os.path.join(plotDir, 'Train-VGG-Sty-Loss'),
loss_avg_vgg_sty.val().item())
lib.plot.plot(os.path.join(plotDir, 'Valid-Synth-Loss'),
valid_loss[0].item())
lib.plot.plot(os.path.join(plotDir, 'Valid-Dis-Loss'),
valid_loss[1].item())
lib.plot.plot(os.path.join(plotDir, 'Valid-Gen-Loss'),
valid_loss[2].item())
lib.plot.plot(os.path.join(plotDir, 'Valid-ImgRecon1-Loss'),
valid_loss[3].item())
lib.plot.plot(os.path.join(plotDir, 'Valid-VGG-Per-Loss'),
valid_loss[4].item())
lib.plot.plot(os.path.join(plotDir, 'Valid-VGG-Sty-Loss'),
valid_loss[5].item())
print(loss_log)
loss_avg.reset()
loss_avg_dis.reset()
loss_avg_gen.reset()
loss_avg_imgRecon.reset()
loss_avg_vgg_per.reset()
loss_avg_vgg_sty.reset()
lib.plot.flush()
lib.plot.tick()
# save model per 1e+5 iter.
if (iteration) % 1e+4 == 0:
torch.save(
{
'styleModel': styleModel.state_dict(),
'genModel': genModel.state_dict(),
'disModel': disModel.state_dict()
},
os.path.join(opt.exp_dir, opt.exp_name,
'iter_' + str(iteration + 1) + '_synth.pth'))
if (iteration + 1) == opt.num_iter:
print('end the training')
sys.exit()
iteration += 1
def train(opt):
""" dataset preparation """
if not opt.data_filtering_off:
print('Filtering the images containing characters which are not in opt.character')
print('Filtering the images whose label is longer than opt.batch_max_length')
# see https://github.com/clovaai/deep-text-recognition-benchmark/blob/6593928855fb7abb999a99f428b3e4477d4ae356/dataset.py#L130
opt.select_data = opt.select_data.split('-')
opt.batch_ratio = opt.batch_ratio.split('-')
#considering the real images for discriminator
opt.batch_size = opt.batch_size*2
train_dataset = Batch_Balanced_Dataset(opt)
log = open(os.path.join(opt.exp_dir,opt.exp_name,'log_dataset.txt'), 'a')
AlignCollate_valid = AlignCollate(imgH=opt.imgH, imgW=opt.imgW, keep_ratio_with_pad=opt.PAD)
valid_dataset, valid_dataset_log = hierarchical_dataset(root=opt.valid_data, opt=opt)
valid_loader = torch.utils.data.DataLoader(
valid_dataset, batch_size=opt.batch_size,
shuffle=True, # 'True' to check training progress with validation function.
num_workers=int(opt.workers),
collate_fn=AlignCollate_valid, pin_memory=True)
log.write(valid_dataset_log)
print('-' * 80)
log.write('-' * 80 + '\n')
log.close()
""" model configuration """
if 'CTC' in opt.Prediction:
converter = CTCLabelConverter(opt.character)
else:
converter = AttnLabelConverter(opt.character)
opt.num_class = len(converter.character)
if opt.rgb:
opt.input_channel = 3
model = AdaINGen(opt)
ocrModel = Model(opt)
disModel = MsImageDis()
print('model input parameters', opt.imgH, opt.imgW, opt.num_fiducial, opt.input_channel, opt.output_channel,
opt.hidden_size, opt.num_class, opt.batch_max_length, opt.Transformation, opt.FeatureExtraction,
opt.SequenceModeling, opt.Prediction)
# Synthesizer weight initialization
for name, param in model.named_parameters():
if 'localization_fc2' in name:
print(f'Skip {name} as it is already initialized')
continue
try:
if 'bias' in name:
init.constant_(param, 0.0)
elif 'weight' in name:
init.kaiming_normal_(param)
except Exception as e: # for batchnorm.
if 'weight' in name:
param.data.fill_(1)
continue
# Recognizer weight initialization
for name, param in ocrModel.named_parameters():
if 'localization_fc2' in name:
print(f'Skip {name} as it is already initialized')
continue
try:
if 'bias' in name:
init.constant_(param, 0.0)
elif 'weight' in name:
init.kaiming_normal_(param)
except Exception as e: # for batchnorm.
if 'weight' in name:
param.data.fill_(1)
continue
# Discriminator weight initialization
for name, param in disModel.named_parameters():
if 'localization_fc2' in name:
print(f'Skip {name} as it is already initialized')
continue
try:
if 'bias' in name:
init.constant_(param, 0.0)
elif 'weight' in name:
init.kaiming_normal_(param)
except Exception as e: # for batchnorm.
if 'weight' in name:
param.data.fill_(1)
continue
# data parallel for multi-GPU
model = torch.nn.DataParallel(model).to(device)
model.train()
ocrModel = torch.nn.DataParallel(ocrModel).to(device)
ocrModel.train()
disModel = torch.nn.DataParallel(disModel).to(device)
disModel.train()
if opt.saved_synth_model != '':
print(f'loading pretrained synth model from {opt.saved_synth_model}')
if opt.FT:
model.load_state_dict(torch.load(opt.saved_synth_model), strict=False)
else:
model.load_state_dict(torch.load(opt.saved_synth_model))
print("Model:")
print(model)
if opt.saved_ocr_model != '':
print(f'loading pretrained ocr model from {opt.saved_ocr_model}')
if opt.FT:
ocrModel.load_state_dict(torch.load(opt.saved_ocr_model), strict=False)
else:
ocrModel.load_state_dict(torch.load(opt.saved_ocr_model))
# ocrModel.eval() #as we can't call RNN.backward in eval mode
print("OCRModel:")
print(ocrModel)
if opt.saved_dis_model != '':
print(f'loading pretrained discriminator model from {opt.saved_dis_model}')
if opt.FT:
disModel.load_state_dict(torch.load(opt.saved_dis_model), strict=False)
else:
disModel.load_state_dict(torch.load(opt.saved_dis_model))
# ocrModel.eval() #as we can't call RNN.backward in eval mode
print("DisModel:")
print(disModel)
""" setup loss """
if 'CTC' in opt.Prediction:
ocrCriterion = torch.nn.CTCLoss(zero_infinity=True).to(device)
else:
ocrCriterion = torch.nn.CrossEntropyLoss(ignore_index=0).to(device) # ignore [GO] token = ignore index 0
recCriterion = torch.nn.L1Loss()
# loss averager
loss_avg = Averager()
loss_avg_ocr = Averager() ##----------
loss_avg_dis = Averager()
# filter that only require gradient decent
filtered_parameters = []
params_num = []
for p in filter(lambda p: p.requires_grad, model.parameters()):
filtered_parameters.append(p)
params_num.append(np.prod(p.size()))
print('Trainable params num : ', sum(params_num))
# [print(name, p.numel()) for name, p in filter(lambda p: p[1].requires_grad, model.named_parameters())]
# setup optimizer
if opt.adam:
optimizer = optim.Adam(filtered_parameters, lr=opt.lr, betas=(opt.beta1, 0.999))
else:
optimizer = optim.Adadelta(filtered_parameters, lr=opt.lr, rho=opt.rho, eps=opt.eps)
print("SynthOptimizer:")
print(optimizer)
#filter parameters for OCR training
# filter that only require gradient decent
ocr_filtered_parameters = []
ocr_params_num = []
for p in filter(lambda p: p.requires_grad, ocrModel.parameters()):
ocr_filtered_parameters.append(p)
ocr_params_num.append(np.prod(p.size()))
print('OCR Trainable params num : ', sum(ocr_params_num))
# setup optimizer
if opt.adam:
ocr_optimizer = optim.Adam(ocr_filtered_parameters, lr=opt.lr, betas=(opt.beta1, 0.999))
else:
ocr_optimizer = optim.Adadelta(ocr_filtered_parameters, lr=opt.lr, rho=opt.rho, eps=opt.eps)
print("OCROptimizer:")
print(ocr_optimizer)
#filter parameters for OCR training
# filter that only require gradient decent
dis_filtered_parameters = []
dis_params_num = []
for p in filter(lambda p: p.requires_grad, disModel.parameters()):
dis_filtered_parameters.append(p)
dis_params_num.append(np.prod(p.size()))
print('Dis Trainable params num : ', sum(dis_params_num))
# setup optimizer
if opt.adam:
dis_optimizer = optim.Adam(dis_filtered_parameters, lr=opt.lr, betas=(opt.beta1, 0.999))
else:
dis_optimizer = optim.Adadelta(dis_filtered_parameters, lr=opt.lr, rho=opt.rho, eps=opt.eps)
print("DisOptimizer:")
print(dis_optimizer)
""" final options """
# print(opt)
with open(os.path.join(opt.exp_dir,opt.exp_name,'opt.txt'), 'a') as opt_file:
opt_log = '------------ Options -------------\n'
args = vars(opt)
for k, v in args.items():
opt_log += f'{str(k)}: {str(v)}\n'
opt_log += '---------------------------------------\n'
print(opt_log)
opt_file.write(opt_log)
""" start training """
start_iter = 0
if opt.saved_synth_model != '':
try:
start_iter = int(opt.saved_synth_model.split('_')[-1].split('.')[0])
print(f'continue to train, start_iter: {start_iter}')
except:
pass
start_time = time.time()
best_accuracy = -1
best_norm_ED = -1
best_accuracy_ocr = -1
best_norm_ED_ocr = -1
iteration = start_iter
while(True):
# train part
image_tensors_all, labels_1_all, labels_2_all = train_dataset.get_batch()
# ## comment
# pdb.set_trace()
# for imgCntr in range(image_tensors.shape[0]):
# save_image(tensor2im(image_tensors[imgCntr]),'temp/'+str(imgCntr)+'.png')
# pdb.set_trace()
# ###
disCnt = int(image_tensors_all.size(0)/2)
image_tensors, image_tensors_real, labels_1, labels_2 = image_tensors_all[:disCnt], image_tensors_all[disCnt:disCnt+disCnt], labels_1_all[:disCnt], labels_2_all[:disCnt]
image = image_tensors.to(device)
image_real = image_tensors_real.to(device)
text_1, length_1 = converter.encode(labels_1, batch_max_length=opt.batch_max_length)
text_2, length_2 = converter.encode(labels_2, batch_max_length=opt.batch_max_length)
batch_size = image.size(0)
images_recon_1, images_recon_2, _ = model(image, text_1, text_2)
if 'CTC' in opt.Prediction:
#ocr training
preds_ocr = ocrModel(image, text_1)
preds_size_ocr = torch.IntTensor([preds_ocr.size(1)] * batch_size)
preds_ocr = preds_ocr.log_softmax(2).permute(1, 0, 2)
ocrCost_train = ocrCriterion(preds_ocr, text_1, preds_size_ocr, length_1)
#dis training
#Check: Using alternate real images
disCost = opt.disWeight*0.5*(disModel.module.calc_dis_loss(images_recon_1.detach(), image_real) + disModel.module.calc_dis_loss(images_recon_2.detach(), image))
#synth training
preds_1 = ocrModel(images_recon_1, text_1)
preds_size_1 = torch.IntTensor([preds_1.size(1)] * batch_size)
preds_1 = preds_1.log_softmax(2).permute(1, 0, 2)
preds_2 = ocrModel(images_recon_2, text_2)
preds_size_2 = torch.IntTensor([preds_2.size(1)] * batch_size)
preds_2 = preds_2.log_softmax(2).permute(1, 0, 2)
ocrCost = 0.5*(ocrCriterion(preds_1, text_1, preds_size_1, length_1) + ocrCriterion(preds_2, text_2, preds_size_2, length_2))
#gen training
disGenCost = 0.5*(disModel.module.calc_gen_loss(images_recon_1)+disModel.module.calc_gen_loss(images_recon_2))
else:
preds = model(image, text[:, :-1]) # align with Attention.forward
target = text[:, 1:] # without [GO] Symbol
ocrCost = ocrCriterion(preds.view(-1, preds.shape[-1]), target.contiguous().view(-1))
recCost = recCriterion(images_recon_1,image)
cost = opt.ocrWeight*ocrCost + opt.reconWeight*recCost + opt.disWeight*disGenCost
disModel.zero_grad()
disCost.backward()
torch.nn.utils.clip_grad_norm_(disModel.parameters(), opt.grad_clip) # gradient clipping with 5 (Default)
dis_optimizer.step()
loss_avg_dis.add(disCost)
model.zero_grad()
ocrModel.zero_grad()
disModel.zero_grad()
cost.backward()
torch.nn.utils.clip_grad_norm_(model.parameters(), opt.grad_clip) # gradient clipping with 5 (Default)
optimizer.step()
loss_avg.add(cost)
#training OCR
ocrModel.zero_grad()
ocrCost_train.backward()
torch.nn.utils.clip_grad_norm_(ocrModel.parameters(), opt.grad_clip) # gradient clipping with 5 (Default)
ocr_optimizer.step()
loss_avg_ocr.add(ocrCost_train)
#START HERE
# validation part
if (iteration + 1) % opt.valInterval == 0 or iteration == 0: # To see training progress, we also conduct validation when 'iteration == 0'
#Save training images
os.makedirs(os.path.join(opt.exp_dir,opt.exp_name,'trainImages',str(iteration)), exist_ok=True)
for trImgCntr in range(batch_size):
try:
save_image(tensor2im(image[trImgCntr].detach()),os.path.join(opt.exp_dir,opt.exp_name,'trainImages',str(iteration),str(trImgCntr)+'_input_'+labels_1[trImgCntr]+'.png'))
save_image(tensor2im(images_recon_1[trImgCntr].detach()),os.path.join(opt.exp_dir,opt.exp_name,'trainImages',str(iteration),str(trImgCntr)+'_recon_'+labels_1[trImgCntr]+'.png'))
save_image(tensor2im(images_recon_2[trImgCntr].detach()),os.path.join(opt.exp_dir,opt.exp_name,'trainImages',str(iteration),str(trImgCntr)+'_pair_'+labels_2[trImgCntr]+'.png'))
except:
print('Warning while saving training image')
elapsed_time = time.time() - start_time
# for log
with open(os.path.join(opt.exp_dir,opt.exp_name,'log_train.txt'), 'a') as log:
model.eval()
ocrModel.eval()
disModel.eval()
with torch.no_grad():
# valid_loss, current_accuracy, current_norm_ED, preds, confidence_score, labels, infer_time, length_of_data = validation(
# model, criterion, valid_loader, converter, opt)
valid_loss, current_accuracy, current_norm_ED, preds, confidence_score, labels, infer_time, length_of_data = validation_synth_adv(
iteration, model, ocrModel, disModel, recCriterion, ocrCriterion, valid_loader, converter, opt)
model.train()
ocrModel.train()
disModel.train()
# training loss and validation loss
loss_log = f'[{iteration+1}/{opt.num_iter}] Train OCR loss: {loss_avg_ocr.val():0.5f}, Train Synth loss: {loss_avg.val():0.5f}, Train Dis loss: {loss_avg_dis.val():0.5f}, Valid OCR loss: {valid_loss[0]:0.5f}, Valid Synth loss: {valid_loss[1]:0.5f}, Valid Dis loss: {valid_loss[2]:0.5f}, Elapsed_time: {elapsed_time:0.5f}'
loss_avg_ocr.reset()
loss_avg.reset()
loss_avg_dis.reset()
current_model_log_ocr = f'{"Current_accuracy_OCR":17s}: {current_accuracy[0]:0.3f}, {"Current_norm_ED_OCR":17s}: {current_norm_ED[0]:0.2f}'
current_model_log_1 = f'{"Current_accuracy_recon":17s}: {current_accuracy[1]:0.3f}, {"Current_norm_ED_recon":17s}: {current_norm_ED[1]:0.2f}'
current_model_log_2 = f'{"Current_accuracy_pair":17s}: {current_accuracy[2]:0.3f}, {"Current_norm_ED_pair":17s}: {current_norm_ED[2]:0.2f}'
# keep best accuracy model (on valid dataset)
if current_accuracy[1] > best_accuracy:
best_accuracy = current_accuracy[1]
torch.save(model.state_dict(), os.path.join(opt.exp_dir,opt.exp_name,'best_accuracy.pth'))
torch.save(disModel.state_dict(), os.path.join(opt.exp_dir,opt.exp_name,'best_accuracy_dis.pth'))
if current_norm_ED[1] > best_norm_ED:
best_norm_ED = current_norm_ED[1]
torch.save(model.state_dict(), os.path.join(opt.exp_dir,opt.exp_name,'best_norm_ED.pth'))
torch.save(disModel.state_dict(), os.path.join(opt.exp_dir,opt.exp_name,'best_norm_ED_dis.pth'))
best_model_log = f'{"Best_accuracy_Recon":17s}: {best_accuracy:0.3f}, {"Best_norm_ED_Recon":17s}: {best_norm_ED:0.2f}'
# keep best accuracy model (on valid dataset)
if current_accuracy[0] > best_accuracy_ocr:
best_accuracy_ocr = current_accuracy[0]
torch.save(ocrModel.state_dict(), os.path.join(opt.exp_dir,opt.exp_name,'best_accuracy_ocr.pth'))
if current_norm_ED[0] > best_norm_ED_ocr:
best_norm_ED_ocr = current_norm_ED[0]
torch.save(ocrModel.state_dict(), os.path.join(opt.exp_dir,opt.exp_name,'best_norm_ED_ocr.pth'))
best_model_log_ocr = f'{"Best_accuracy_ocr":17s}: {best_accuracy_ocr:0.3f}, {"Best_norm_ED_ocr":17s}: {best_norm_ED_ocr:0.2f}'
loss_model_log = f'{loss_log}\n{current_model_log_ocr}\n{current_model_log_1}\n{current_model_log_2}\n{best_model_log_ocr}\n{best_model_log}'
print(loss_model_log)
log.write(loss_model_log + '\n')
# show some predicted results
dashed_line = '-' * 80
head = f'{"Ground Truth":32s} | {"Prediction":25s} | Confidence Score & T/F'
predicted_result_log = f'{dashed_line}\n{head}\n{dashed_line}\n'
for gt_ocr, pred_ocr, confidence_ocr, gt_1, pred_1, confidence_1, gt_2, pred_2, confidence_2 in zip(labels[0][:5], preds[0][:5], confidence_score[0][:5], labels[1][:5], preds[1][:5], confidence_score[1][:5], labels[2][:5], preds[2][:5], confidence_score[2][:5]):
if 'Attn' in opt.Prediction:
gt = gt[:gt.find('[s]')]
pred = pred[:pred.find('[s]')]
predicted_result_log += f'{"ocr"}: {gt_ocr:27s} | {pred_ocr:25s} | {confidence_ocr:0.4f}\t{str(pred_ocr == gt_ocr)}\n'
predicted_result_log += f'{"recon"}: {gt_1:25s} | {pred_1:25s} | {confidence_1:0.4f}\t{str(pred_1 == gt_1)}\n'
predicted_result_log += f'{"pair"}: {gt_2:26s} | {pred_2:25s} | {confidence_2:0.4f}\t{str(pred_2 == gt_2)}\n'
predicted_result_log += f'{dashed_line}'
print(predicted_result_log)
log.write(predicted_result_log + '\n')
# save model per 1e+5 iter.
if (iteration + 1) % 1e+5 == 0:
torch.save(
model.state_dict(), os.path.join(opt.exp_dir,opt.exp_name,'iter_{iteration+1}.pth'))
torch.save(
ocrModel.state_dict(), os.path.join(opt.exp_dir,opt.exp_name,'iter_{iteration+1}_ocr.pth'))
torch.save(
disModel.state_dict(), os.path.join(opt.exp_dir,opt.exp_name,'iter_{iteration+1}_dis.pth'))
if (iteration + 1) == opt.num_iter:
print('end the training')
sys.exit()
iteration += 1
def train(**kwargs):
# Retrieve training configuration
data_loader = kwargs['data_loader']
net = kwargs['net']
loss = kwargs['loss']
loss_metric = kwargs['loss_metric']
optimizer = kwargs['optimizer']
feature_center = kwargs['feature_center']
epoch = kwargs['epoch']
save_freq = kwargs['save_freq']
save_dir = kwargs['save_dir']
verbose = kwargs['verbose']
global step
# Attention Regularization: LA Loss
l2_loss = nn.MSELoss()
# Default Parameters
beta = 0.05
theta_c = 0.5
theta_d = 0.5
crop_size = (256, 256) # size of cropped images for 'See Better'
# metrics initialization
batches = 0
epoch_loss_raw = np.array(
[0, 0, 0, 0, 0],
dtype='float') # loss_sum(classify, center, h**o, heter)
epoch_loss_crop = np.array([0, 0, 0, 0],
dtype='float') # loss_sum(classify, triplet)
epoch_loss_drop = np.array([0], dtype='float')
epoch_acc = np.array(
[[0, 0, 0], [0, 0, 0], [0, 0, 0]],
dtype='float') # Top-1/3/5 Accuracy for Raw/Crop/Drop Images
# begin training
start_time = time.time()
logging.info('Epoch %03d, Learning Rate %g' %
(epoch + 1, optimizer.param_groups[0]['lr']))
net.train()
for i, batch in enumerate(data_loader):
batch_start = time.time()
# obtain data for training
X = batch[0].to(torch.device("cuda"))
y = batch[1].to(torch.device("cuda"))
if save_crop_image:
fns = batch[2][0]
##################################
# Raw Image
##################################
y_pred, embeddings, attention_map = net(X)
#print('RAW y_pred:\n', y_pred[0, ...], '\n', y_pred[2, ...])
# loss
loss_triplet = loss_metric(embeddings)
#if not options.freeze:
loss_classify = loss(y_pred, y)
loss_center = center_loss(embeddings, feature_center[y])
epoch_loss_raw[0] += (loss_classify + loss_center +
loss_triplet).item()
epoch_loss_raw[1] += loss_classify.item()
epoch_loss_raw[2] += loss_center.item()
epoch_loss_raw[3] += loss_triplet.item()
writer.add_scalars(main_tag='Raw',
tag_scalar_dict={
'classify': loss_classify.item(),
'center': loss_center.item(),
'triplet': loss_triplet.item()
},
global_step=step)
# vis gradient flow
if (1 + i) % 100 == 0:
writer.add_figure('Raw_grad_flow',
plot_grad_flow_v2(net.named_parameters()),
global_step=(step + 1) // 100)
# Update Feature Center
embed = embeddings.detach()
for idx in range(len(y)):
feature_center[y[idx]] += beta * (embed[idx] -
feature_center[y[idx]])
# metrics: top-1, top-3, top-5 error
with torch.no_grad():
epoch_acc[0] += accuracy(y_pred, y,
topk=(1, 3, 5)).astype(np.float)
##################################
# Attention Cropping
##################################
with torch.no_grad():
crop_images = []
for batch_index in range(attention_map.size(0)):
theta = torch.max(
attention_map[batch_index]) * np.random.uniform(0.4, 0.6)
crop_mask = attention_map[batch_index] > theta
nonzero_indices = torch.nonzero(crop_mask[0, ...])
height_min = torch.clamp(nonzero_indices[:, 0].min().float() /
attention_map.size(2) - 0.01,
min=0) * X.size(2)
height_max = torch.clamp(nonzero_indices[:, 0].max().float() /
attention_map.size(2) + 0.01,
max=1) * X.size(2)
width_min = torch.clamp(nonzero_indices[:, 1].min().float() /
attention_map.size(3) - 0.01,
min=0) * X.size(3)
width_max = torch.clamp(nonzero_indices[:, 1].max().float() /
attention_map.size(3) + 0.01,
max=1) * X.size(3)
if (height_min == height_max) or (width_min == width_max):
crop_images.append(X)
else:
crop_images.append(
rescale_central(X[batch_index:batch_index + 1, ...], [
int(height_min.item()),
int(height_max.item()),
int(width_min.item()),
int(width_max.item())
], crop_size[0]))
crop_images = torch.cat(crop_images, dim=0)
# save crop images to ./images/
if save_crop_image:
for idx in range(crop_images.size(0)):
img = ImageFolderWithName.tensor2img(
dataset.invTrans(crop_images[idx]))[0]
img = img.astype(np.uint8)
fn = fns[idx].split('/')
dataset.save_image(img,
'./images/' + fn[-2] + '/' + fn[-1])
# crop images forward
y_pred, embeddings_cropped, _ = net(crop_images)
# loss
loss_triplet_crop = loss_metric(embeddings_cropped)
loss_triplet += loss_triplet_crop
embedding_loss = center_loss(embeddings_cropped, embed)
loss_crop_classify = loss(y_pred, y)
loss_classify += loss_crop_classify
epoch_loss_crop += np.array([
(loss_crop_classify + loss_triplet_crop).item(),
loss_crop_classify.item(),
loss_triplet_crop.item(),
embedding_loss.item()
])
if (1 + i) % 100 == 0:
writer.add_figure('Cropped_grad_flow',
plot_grad_flow_v2(net.named_parameters()),
global_step=(step + 1) // 100)
writer.add_image('Crop Image',
dataset.invTrans(crop_images[0, ...]),
global_step=(step + 1) // 100)
# metrics: top-1, top-3, top-5 error
with torch.no_grad():
epoch_acc[1] += accuracy(y_pred, y,
topk=(1, 3, 5)).astype(np.float)
##################################
# Attention Dropping
##################################
with torch.no_grad():
drop_mask = F.upsample_bilinear(attention_map,
size=(X.size(2), X.size(3)))
for batch_index in range(attention_map.size(0)):
theta = torch.max(
attention_map[batch_index]) * np.random.uniform(0.4, 0.6)
drop_mask[batch_index] = drop_mask[batch_index] < theta
drop_images = X * drop_mask.float()
if save_crop_image:
for idx in range(crop_images.size(0)):
img = ImageFolderWithName.tensor2img(
dataset.invTrans(drop_images[idx]))[0]
img = img.astype(np.uint8)
fn = fns[idx].split('/')
dataset.save_image(
img, './images/' + fn[-2] + '/' +
fn[-1].replace('.', '_drop.'))
if (1 + i) % 100 == 0:
writer.add_image('Drop_Image',
dataset.invTrans(drop_images[0, ...]),
global_step=(step + 1) // 100)
# drop images forward
y_pred, _, _ = net(drop_images)
# loss
loss_drop_classify = loss(y_pred, y)
loss_classify += 0.5 * loss_drop_classify
epoch_loss_drop[0] += loss_drop_classify.item()
# backward
batch_loss = loss_classify / 3 + loss_triplet / 2 + loss_center + embedding_loss
optimizer.zero_grad()
batch_loss.backward()
optimizer.step()
# metrics: top-1, top-3, top-5 error
with torch.no_grad():
epoch_acc[2] += accuracy(y_pred, y,
topk=(1, 3, 5)).astype(np.float)
# end of this batch
batches += 1
step += 1
batch_end = time.time()
if (i + 1) % verbose == 0:
logging.info(
'\tBatch %d: (Raw) Loss %.3f (%.3f, %.3f, %.3f, %.3f), Accuracy: (%.2f, %.2f, %.2f), (Crop) Loss %.3f (%.3f, %.3f, %.3f), Accuracy: (%.2f, %.2f, %.2f), (Drop) Loss %.3f, Accuracy: (%.2f, %.2f, %.2f), Time %3.2f'
%
(i + 1, epoch_loss_raw[0] / batches, epoch_loss_raw[1] /
batches, epoch_loss_raw[2] / batches, epoch_loss_raw[3] /
batches, epoch_loss_raw[4] / batches, epoch_acc[0, 0] /
batches, epoch_acc[0, 1] / batches, epoch_acc[0, 2] / batches,
epoch_loss_crop[0] / batches, epoch_loss_crop[1] / batches,
epoch_loss_crop[2] / batches, epoch_loss_crop[3] / batches,
epoch_acc[1, 0] / batches, epoch_acc[1, 1] / batches,
epoch_acc[1, 2] / batches, epoch_loss_drop[0] / batches,
epoch_acc[2, 0] / batches, epoch_acc[2, 1] / batches,
epoch_acc[2, 2] / batches, batch_end - batch_start))
# save checkpoint model
if epoch % save_freq == 0:
state_dict = net.state_dict()
for key in state_dict.keys():
state_dict[key] = state_dict[key].cpu()
torch.save(
{
'epoch': epoch,
'save_dir': save_dir,
'state_dict': state_dict,
'feature_center': feature_center.cpu()
}, os.path.join(save_dir, '%03d.ckpt' % (epoch + 1)))
# end of this epoch
end_time = time.time()
# metrics for average
epoch_loss_raw /= batches
epoch_loss_crop /= batches
epoch_loss_drop /= batches
epoch_acc /= batches
# show information for this epoch
logging.info(
'Train: (Raw) Loss %.4f, Accuracy: (%.2f, %.2f, %.2f), (Crop) Loss %.4f, Accuracy: (%.2f, %.2f, %.2f), (Drop) Loss %.4f, Accuracy: (%.2f, %.2f, %.2f), Time %3.2f'
%
(epoch_loss_raw[0], epoch_acc[0, 0], epoch_acc[0, 1], epoch_acc[0, 2],
epoch_loss_crop[0], epoch_acc[1, 0], epoch_acc[1, 1], epoch_acc[1, 2],
epoch_loss_drop[0], epoch_acc[2, 0], epoch_acc[2, 1], epoch_acc[2, 2],
end_time - start_time))
def test(opt):
lib.print_model_settings(locals().copy())
""" dataset preparation """
if not opt.data_filtering_off:
print(
'Filtering the images containing characters which are not in opt.character'
)
print(
'Filtering the images whose label is longer than opt.batch_max_length'
)
# see https://github.com/clovaai/deep-text-recognition-benchmark/blob/6593928855fb7abb999a99f428b3e4477d4ae356/dataset.py#L130
log = open(os.path.join(opt.exp_dir, opt.exp_name, 'log_dataset.txt'), 'a')
AlignCollate_valid = AlignPairCollate(imgH=opt.imgH,
imgW=opt.imgW,
keep_ratio_with_pad=opt.PAD)
valid_dataset, valid_dataset_log = hierarchical_dataset(
root=opt.valid_data, opt=opt)
valid_loader = torch.utils.data.DataLoader(
valid_dataset,
batch_size=opt.batch_size,
shuffle=
False, # 'True' to check training progress with validation function.
num_workers=int(opt.workers),
collate_fn=AlignCollate_valid,
pin_memory=True,
drop_last=True)
log.write(valid_dataset_log)
print('-' * 80)
log.write('-' * 80 + '\n')
log.close()
if 'Attn' in opt.Prediction:
converter = AttnLabelConverter(opt.character)
else:
converter = CTCLabelConverter(opt.character)
opt.num_class = len(converter.character)
g_ema = styleGANGen(opt.size,
opt.latent,
opt.n_mlp,
opt.num_class,
channel_multiplier=opt.channel_multiplier)
g_ema = torch.nn.DataParallel(g_ema).to(device)
g_ema.eval()
print('model input parameters', opt.imgH, opt.imgW, opt.input_channel,
opt.output_channel, opt.hidden_size, opt.num_class,
opt.batch_max_length)
## Loading pre-trained files
if opt.modelFolderFlag:
if len(
glob.glob(
os.path.join(opt.exp_dir, opt.exp_name,
"iter_*_synth.pth"))) > 0:
opt.saved_synth_model = glob.glob(
os.path.join(opt.exp_dir, opt.exp_name,
"iter_*_synth.pth"))[-1]
if opt.saved_synth_model != '' and opt.saved_synth_model != 'None':
print(f'loading pretrained synth model from {opt.saved_synth_model}')
checkpoint = torch.load(opt.saved_synth_model)
g_ema.load_state_dict(checkpoint['g_ema'], strict=False)
# pdb.set_trace()
if opt.truncation < 1:
with torch.no_grad():
mean_latent = g_ema.module.mean_latent_content(opt.truncation_mean)
else:
mean_latent = None
cntr = 0
for i, (image_input_tensors, image_gt_tensors, labels_1,
labels_2) in enumerate(valid_loader):
print(i, len(valid_loader))
image_input_tensors = image_input_tensors.to(device)
image_gt_tensors = image_gt_tensors.to(device)
batch_size = image_input_tensors.size(0)
text_1, length_1 = converter.encode(
labels_1, batch_max_length=opt.batch_max_length)
text_2, length_2 = converter.encode(
labels_2, batch_max_length=opt.batch_max_length)
#forward pass from style and word generator
if opt.fixedStyleBatch:
fixstyle = []
# pdb.set_trace()
style = mixing_noise(1, opt.latent, opt.mixing, device)
fixstyle.append(style[0].repeat(opt.batch_size, 1))
if len(style) > 1:
fixstyle.append(style[1].repeat(opt.batch_size, 1))
style = fixstyle
else:
style = mixing_noise(opt.batch_size, opt.latent, opt.mixing,
device)
if 'CTC' in opt.Prediction:
images_recon_2, _ = g_ema(style,
text_2,
input_is_latent=opt.input_latent,
inject_index=5,
truncation=opt.truncation,
truncation_latent=mean_latent,
randomize_noise=False)
else:
images_recon_2, _ = g_ema(style,
text_2[:, 1:-1],
input_is_latent=opt.input_latent,
inject_index=5,
truncation=opt.truncation,
truncation_latent=mean_latent,
randomize_noise=False)
# os.makedirs(os.path.join(opt.valDir,str(iteration)), exist_ok=True)
for trImgCntr in range(batch_size):
try:
save_image(
tensor2im(image_input_tensors[trImgCntr].detach()),
os.path.join(
opt.valDir,
str(cntr) + '_' + str(trImgCntr) + '_sInput_' +
labels_1[trImgCntr] + '.png'))
save_image(
tensor2im(image_gt_tensors[trImgCntr].detach()),
os.path.join(
opt.valDir,
str(cntr) + '_' + str(trImgCntr) + '_csGT_' +
labels_2[trImgCntr] + '.png'))
save_image(
tensor2im(images_recon_2[trImgCntr].detach()),
os.path.join(
opt.valDir,
str(cntr) + '_' + str(trImgCntr) + '_csRecon_' +
labels_2[trImgCntr] + '.png'))
except:
print('Warning while saving training image')
cntr += 1
def train(opt):
lib.print_model_settings(locals().copy())
if 'Attn' in opt.Prediction:
converter = AttnLabelConverter(opt.character)
else:
converter = CTCLabelConverter(opt.character)
opt.classes = converter.character
""" dataset preparation """
if not opt.data_filtering_off:
print(
'Filtering the images containing characters which are not in opt.character'
)
print(
'Filtering the images whose label is longer than opt.batch_max_length'
)
# see https://github.com/clovaai/deep-text-recognition-benchmark/blob/6593928855fb7abb999a99f428b3e4477d4ae356/dataset.py#L130
log = open(os.path.join(opt.exp_dir, opt.exp_name, 'log_dataset.txt'), 'a')
AlignCollate_valid = AlignPHOCCollate(imgH=opt.imgH,
imgW=opt.imgW,
keep_ratio_with_pad=opt.PAD)
train_dataset = LmdbStylePHOCDataset(root=opt.train_data, opt=opt)
train_loader = torch.utils.data.DataLoader(
train_dataset,
batch_size=opt.batch_size *
2, #*2 to sample different images from training encoder and discriminator real images
shuffle=
True, # 'True' to check training progress with validation function.
num_workers=int(opt.workers),
collate_fn=AlignCollate_valid,
pin_memory=True,
drop_last=True)
print('-' * 80)
valid_dataset = LmdbStylePHOCDataset(root=opt.valid_data, opt=opt)
valid_loader = torch.utils.data.DataLoader(
valid_dataset,
batch_size=opt.batch_size *
2, #*2 to sample different images from training encoder and discriminator real images
shuffle=
False, # 'True' to check training progress with validation function.
num_workers=int(opt.workers),
collate_fn=AlignCollate_valid,
pin_memory=True,
drop_last=True)
print('-' * 80)
log.write('-' * 80 + '\n')
log.close()
phoc_dataset = phoc_gen(opt)
phoc_loader = torch.utils.data.DataLoader(phoc_dataset,
batch_size=opt.batch_size,
shuffle=True,
num_workers=int(opt.workers),
pin_memory=True,
drop_last=True)
opt.num_class = len(converter.character)
if opt.zAlone:
genModel = styleGANGen(opt.size,
opt.latent,
opt.latent,
opt.n_mlp,
channel_multiplier=opt.channel_multiplier)
g_ema = styleGANGen(opt.size,
opt.latent,
opt.latent,
opt.n_mlp,
channel_multiplier=opt.channel_multiplier)
else:
genModel = styleGANGen(opt.size,
opt.latent + phoc_dataset.phoc_size,
opt.latent,
opt.n_mlp,
channel_multiplier=opt.channel_multiplier)
g_ema = styleGANGen(opt.size,
opt.latent + phoc_dataset.phoc_size,
opt.latent,
opt.n_mlp,
channel_multiplier=opt.channel_multiplier)
disEncModel = styleGANDis(opt.size,
channel_multiplier=opt.channel_multiplier,
input_dim=opt.input_channel,
code_s_dim=phoc_dataset.phoc_size)
accumulate(g_ema, genModel, 0)
uCriterion = torch.nn.MSELoss()
sCriterion = torch.nn.MSELoss()
genModel = torch.nn.DataParallel(genModel).to(device)
g_ema = torch.nn.DataParallel(g_ema).to(device)
genModel.train()
g_ema.eval()
disEncModel = torch.nn.DataParallel(disEncModel).to(device)
disEncModel.train()
g_reg_ratio = opt.g_reg_every / (opt.g_reg_every + 1)
d_reg_ratio = opt.d_reg_every / (opt.d_reg_every + 1)
optimizer = optim.Adam(
genModel.parameters(),
lr=opt.lr * g_reg_ratio,
betas=(0**g_reg_ratio, 0.99**g_reg_ratio),
)
dis_optimizer = optim.Adam(
disEncModel.parameters(),
lr=opt.lr * d_reg_ratio,
betas=(0**d_reg_ratio, 0.99**d_reg_ratio),
)
## Loading pre-trained files
if opt.modelFolderFlag:
if len(
glob.glob(
os.path.join(opt.exp_dir, opt.exp_name,
"iter_*_synth.pth"))) > 0:
opt.saved_synth_model = glob.glob(
os.path.join(opt.exp_dir, opt.exp_name,
"iter_*_synth.pth"))[-1]
# if opt.saved_ocr_model !='' and opt.saved_ocr_model !='None':
# print(f'loading pretrained ocr model from {opt.saved_ocr_model}')
# checkpoint = torch.load(opt.saved_ocr_model)
# ocrModel.load_state_dict(checkpoint)
# if opt.saved_gen_model !='' and opt.saved_gen_model !='None':
# print(f'loading pretrained gen model from {opt.saved_gen_model}')
# checkpoint = torch.load(opt.saved_gen_model, map_location=lambda storage, loc: storage)
# genModel.module.load_state_dict(checkpoint['g'])
# g_ema.module.load_state_dict(checkpoint['g_ema'])
if opt.saved_synth_model != '' and opt.saved_synth_model != 'None':
print(f'loading pretrained synth model from {opt.saved_synth_model}')
checkpoint = torch.load(opt.saved_synth_model)
# styleModel.load_state_dict(checkpoint['styleModel'])
# mixModel.load_state_dict(checkpoint['mixModel'])
genModel.load_state_dict(checkpoint['genModel'])
g_ema.load_state_dict(checkpoint['g_ema'])
disEncModel.load_state_dict(checkpoint['disEncModel'])
optimizer.load_state_dict(checkpoint["optimizer"])
dis_optimizer.load_state_dict(checkpoint["dis_optimizer"])
# if opt.imgReconLoss == 'l1':
# recCriterion = torch.nn.L1Loss()
# elif opt.imgReconLoss == 'ssim':
# recCriterion = ssim
# elif opt.imgReconLoss == 'ms-ssim':
# recCriterion = msssim
# loss averager
loss_avg_dis = Averager()
loss_avg_gen = Averager()
loss_avg_unsup = Averager()
loss_avg_sup = Averager()
log_r1_val = Averager()
log_avg_path_loss_val = Averager()
log_avg_mean_path_length_avg = Averager()
log_ada_aug_p = Averager()
""" final options """
with open(os.path.join(opt.exp_dir, opt.exp_name, 'opt.txt'),
'a') as opt_file:
opt_log = '------------ Options -------------\n'
args = vars(opt)
for k, v in args.items():
opt_log += f'{str(k)}: {str(v)}\n'
opt_log += '---------------------------------------\n'
print(opt_log)
opt_file.write(opt_log)
""" start training """
start_iter = 0
if opt.saved_synth_model != '' and opt.saved_synth_model != 'None':
try:
start_iter = int(
opt.saved_synth_model.split('_')[-2].split('.')[0])
print(f'continue to train, start_iter: {start_iter}')
except:
pass
#get schedulers
scheduler = get_scheduler(optimizer, opt)
dis_scheduler = get_scheduler(dis_optimizer, opt)
start_time = time.time()
iteration = start_iter
cntr = 0
mean_path_length = 0
d_loss_val = 0
r1_loss = torch.tensor(0.0, device=device)
g_loss_val = 0
path_loss = torch.tensor(0.0, device=device)
path_lengths = torch.tensor(0.0, device=device)
mean_path_length_avg = 0
# loss_dict = {}
accum = 0.5**(32 / (10 * 1000))
ada_augment = torch.tensor([0.0, 0.0], device=device)
ada_aug_p = opt.augment_p if opt.augment_p > 0 else 0.0
ada_aug_step = opt.ada_target / opt.ada_length
r_t_stat = 0
# sample_z = torch.randn(opt.n_sample, opt.latent, device=device)
while (True):
# print(cntr)
# train part
if opt.lr_policy != "None":
scheduler.step()
dis_scheduler.step()
image_input_tensors, _, labels_1, _, phoc_1, _ = iter(
train_loader).next()
z_code, z_labels = iter(phoc_loader).next()
image_input_tensors = image_input_tensors.to(device)
gt_image_tensors = image_input_tensors[:opt.batch_size]
real_image_tensors = image_input_tensors[opt.batch_size:]
phoc_1 = phoc_1.to(device)
gt_phoc_tensors = phoc_1[:opt.batch_size]
labels_1 = labels_1[:opt.batch_size]
z_code = z_code.to(device)
requires_grad(genModel, False)
# requires_grad(styleModel, False)
# requires_grad(mixModel, False)
requires_grad(disEncModel, True)
text_1, length_1 = converter.encode(
labels_1, batch_max_length=opt.batch_max_length)
style = mixing_noise(z_code, opt.batch_size, opt.latent, opt.mixing,
device)
if opt.zAlone:
#to validate orig style gan results
newstyle = []
newstyle.append(style[0][:, :opt.latent])
if len(style) > 1:
newstyle.append(style[1][:, :opt.latent])
style = newstyle
fake_img, _ = genModel(style, input_is_latent=opt.input_latent)
#unsupervised code prediction on generated image
u_pred_code = disEncModel(fake_img, mode='enc')
uCost = uCriterion(u_pred_code, z_code)
#supervised code prediction on gt image
s_pred_code = disEncModel(gt_image_tensors, mode='enc')
sCost = uCriterion(s_pred_code, gt_phoc_tensors)
#Domain discriminator
fake_pred = disEncModel(fake_img)
real_pred = disEncModel(real_image_tensors)
disCost = d_logistic_loss(real_pred, fake_pred)
dis_enc_cost = disCost + opt.gamma_e * uCost + opt.beta * sCost
loss_avg_dis.add(disCost)
loss_avg_sup.add(opt.beta * sCost)
loss_avg_unsup.add(opt.gamma_e * uCost)
disEncModel.zero_grad()
dis_enc_cost.backward()
dis_optimizer.step()
d_regularize = cntr % opt.d_reg_every == 0
if d_regularize:
real_image_tensors.requires_grad = True
real_pred = disEncModel(real_image_tensors)
r1_loss = d_r1_loss(real_pred, real_image_tensors)
disEncModel.zero_grad()
(opt.r1 / 2 * r1_loss * opt.d_reg_every +
0 * real_pred[0]).backward()
dis_optimizer.step()
# loss_dict["r1"] = r1_loss
# [Word Generator] update
image_input_tensors, _, labels_1, _, phoc_1, _ = iter(
train_loader).next()
z_code, z_labels = iter(phoc_loader).next()
image_input_tensors = image_input_tensors.to(device)
gt_image_tensors = image_input_tensors[:opt.batch_size]
real_image_tensors = image_input_tensors[opt.batch_size:]
phoc_1 = phoc_1.to(device)
gt_phoc_tensors = phoc_1[:opt.batch_size]
labels_1 = labels_1[:opt.batch_size]
z_code = z_code.to(device)
requires_grad(genModel, True)
requires_grad(disEncModel, False)
text_1, length_1 = converter.encode(
labels_1, batch_max_length=opt.batch_max_length)
style = mixing_noise(z_code, opt.batch_size, opt.latent, opt.mixing,
device)
if opt.zAlone:
#to validate orig style gan results
newstyle = []
newstyle.append(style[0][:, :opt.latent])
if len(style) > 1:
newstyle.append(style[1][:, :opt.latent])
style = newstyle
fake_img, _ = genModel(style, input_is_latent=opt.input_latent)
#unsupervised code prediction on generated image
u_pred_code = disEncModel(fake_img, mode='enc')
uCost = uCriterion(u_pred_code, z_code)
fake_pred = disEncModel(fake_img)
disGenCost = g_nonsaturating_loss(fake_pred)
gen_enc_cost = disGenCost + opt.gamma_g * uCost
loss_avg_gen.add(disGenCost)
loss_avg_unsup.add(opt.gamma_g * uCost)
# loss_dict["g"] = disGenCost
genModel.zero_grad()
disEncModel.zero_grad()
gen_enc_cost.backward()
optimizer.step()
g_regularize = cntr % opt.g_reg_every == 0
if g_regularize:
image_input_tensors, _, labels_1, _, phoc_1, _ = iter(
train_loader).next()
z_code, z_labels = iter(phoc_loader).next()
image_input_tensors = image_input_tensors.to(device)
path_batch_size = max(1, opt.batch_size // opt.path_batch_shrink)
gt_image_tensors = image_input_tensors[:path_batch_size]
phoc_1 = phoc_1.to(device)
gt_phoc_tensors = phoc_1[:path_batch_size]
labels_1 = labels_1[:path_batch_size]
z_code = z_code.to(device)
z_code = z_code[:path_batch_size]
z_labels = z_labels[:path_batch_size]
text_1, length_1 = converter.encode(
labels_1, batch_max_length=opt.batch_max_length)
style = mixing_noise(z_code, path_batch_size, opt.latent,
opt.mixing, device)
if opt.zAlone:
#to validate orig style gan results
newstyle = []
newstyle.append(style[0][:, :opt.latent])
if len(style) > 1:
newstyle.append(style[1][:, :opt.latent])
style = newstyle
fake_img, grad = genModel(style,
return_latents=True,
g_path_regularize=True,
mean_path_length=mean_path_length)
decay = 0.01
path_lengths = torch.sqrt(grad.pow(2).sum(2).mean(1))
mean_path_length_orig = mean_path_length + decay * (
path_lengths.mean() - mean_path_length)
path_loss = (path_lengths - mean_path_length_orig).pow(2).mean()
mean_path_length = mean_path_length_orig.detach().item()
# path_loss, mean_path_length, path_lengths = g_path_regularize(
# images_recon_2, latents, mean_path_length
# )
genModel.zero_grad()
weighted_path_loss = opt.path_regularize * opt.g_reg_every * path_loss
if opt.path_batch_shrink:
weighted_path_loss += 0 * fake_img[0, 0, 0, 0]
weighted_path_loss.backward()
optimizer.step()
# mean_path_length_avg = (
# reduce_sum(mean_path_length).item() / get_world_size()
# )
#commented above for multi-gpu , non-distributed setting
mean_path_length_avg = mean_path_length
accumulate(g_ema, genModel, accum)
log_r1_val.add(r1_loss)
log_avg_path_loss_val.add(path_loss)
log_avg_mean_path_length_avg.add(torch.tensor(mean_path_length_avg))
log_ada_aug_p.add(torch.tensor(ada_aug_p))
if get_rank() == 0:
if wandb and opt.wandb:
wandb.log({
"Generator": g_loss_val,
"Discriminator": d_loss_val,
"Augment": ada_aug_p,
"Rt": r_t_stat,
"R1": r1_val,
"Path Length Regularization": path_loss_val,
"Mean Path Length": mean_path_length,
"Real Score": real_score_val,
"Fake Score": fake_score_val,
"Path Length": path_length_val,
})
# validation part
if (
iteration + 1
) % opt.valInterval == 0 or iteration == 0: # To see training progress, we also conduct validation when 'iteration == 0'
#generate paired content with similar style
z_code_1, z_labels_1 = iter(phoc_loader).next()
z_code_2, z_labels_2 = iter(phoc_loader).next()
z_code_1 = z_code_1.to(device)
z_code_2 = z_code_2.to(device)
style_1 = mixing_noise(z_code_1, opt.batch_size, opt.latent,
opt.mixing, device)
style_2 = []
style_2.append(
torch.cat((style_1[0][:, :opt.latent], z_code_2), dim=1))
if len(style_1) > 1:
style_2.append(
torch.cat((style_1[1][:, :opt.latent], z_code_2), dim=1))
if opt.zAlone:
#to validate orig style gan results
newstyle = []
newstyle.append(style_1[0][:, :opt.latent])
if len(style_1) > 1:
newstyle.append(style_1[1][:, :opt.latent])
style_1 = newstyle
style_2 = newstyle
fake_img_1, _ = g_ema(style_1, input_is_latent=opt.input_latent)
fake_img_2, _ = g_ema(style_2, input_is_latent=opt.input_latent)
os.makedirs(os.path.join(opt.trainDir, str(iteration)),
exist_ok=True)
for trImgCntr in range(opt.batch_size):
try:
save_image(
tensor2im(fake_img_1[trImgCntr].detach()),
os.path.join(
opt.trainDir, str(iteration),
str(trImgCntr) + '_pair1_' +
z_labels_1[trImgCntr] + '.png'))
save_image(
tensor2im(fake_img_2[trImgCntr].detach()),
os.path.join(
opt.trainDir, str(iteration),
str(trImgCntr) + '_pair2_' +
z_labels_2[trImgCntr] + '.png'))
except:
print('Warning while saving training image')
elapsed_time = time.time() - start_time
# for log
with open(os.path.join(opt.exp_dir, opt.exp_name, 'log_train.txt'),
'a') as log:
# training loss and validation loss
loss_log = f'[{iteration+1}/{opt.num_iter}] \
Train Dis loss: {loss_avg_dis.val():0.5f}, Train Gen loss: {loss_avg_gen.val():0.5f},\
Train UnSup loss: {loss_avg_unsup.val():0.5f}, Train Sup loss: {loss_avg_sup.val():0.5f}, \
Train R1-val loss: {log_r1_val.val():0.5f}, Train avg-path-loss: {log_avg_path_loss_val.val():0.5f}, \
Train mean-path-length loss: {log_avg_mean_path_length_avg.val():0.5f}, Train ada-aug-p: {log_ada_aug_p.val():0.5f}, \
Elapsed_time: {elapsed_time:0.5f}'
#plotting
lib.plot.plot(os.path.join(opt.plotDir, 'Train-Dis-Loss'),
loss_avg_dis.val().item())
lib.plot.plot(os.path.join(opt.plotDir, 'Train-Gen-Loss'),
loss_avg_gen.val().item())
lib.plot.plot(os.path.join(opt.plotDir, 'Train-UnSup-Loss'),
loss_avg_unsup.val().item())
lib.plot.plot(os.path.join(opt.plotDir, 'Train-Sup-Loss'),
loss_avg_sup.val().item())
lib.plot.plot(os.path.join(opt.plotDir, 'Train-r1_val'),
log_r1_val.val().item())
lib.plot.plot(os.path.join(opt.plotDir, 'Train-path_loss_val'),
log_avg_path_loss_val.val().item())
lib.plot.plot(
os.path.join(opt.plotDir, 'Train-mean_path_length_avg'),
log_avg_mean_path_length_avg.val().item())
lib.plot.plot(os.path.join(opt.plotDir, 'Train-ada_aug_p'),
log_ada_aug_p.val().item())
print(loss_log)
loss_avg_dis.reset()
loss_avg_gen.reset()
loss_avg_unsup.reset()
loss_avg_sup.reset()
log_r1_val.reset()
log_avg_path_loss_val.reset()
log_avg_mean_path_length_avg.reset()
log_ada_aug_p.reset()
lib.plot.flush()
lib.plot.tick()
# save model per 1e+5 iter.
if (iteration) % 1e+4 == 0:
torch.save(
{
'genModel': genModel.state_dict(),
'g_ema': g_ema.state_dict(),
'disEncModel': disEncModel.state_dict(),
'optimizer': optimizer.state_dict(),
'dis_optimizer': dis_optimizer.state_dict()
},
os.path.join(opt.exp_dir, opt.exp_name,
'iter_' + str(iteration + 1) + '_synth.pth'))
if (iteration + 1) == opt.num_iter:
print('end the training')
sys.exit()
iteration += 1
cntr += 1
def main():
img = ds.image2(append_index=True)
img = ds.format_square(img)
ds.save_image('image2.tif', img)
def save_final(folder):
dataset.clean(folder)
final1 = get_random_seeds(1)
final2 = get_random_seeds(1)
dataset.save_image(folder, "crossover.AA",
model.generator.generate(final1).numpy()[0])
dataset.save_image(folder, "crossover.BB",
model.generator.generate(final2).numpy()[0])
dataset.save_image(folder, "crossover.AB",
model.generator.generate(final1, final2, 3).numpy()[0])
dataset.save_image(folder, "crossover.BA",
model.generator.generate(final2, final1, 3).numpy()[0])
for i in range(9):
s = i / 4 - 1
f = i / 8
dataset.save_image(folder, "A.fade" + str(i),
model.generator.generate(s * final1).numpy()[0])
dataset.save_image(folder, "B.fade" + str(i),
model.generator.generate(s * final2).numpy()[0])
dataset.save_image(
folder, "transition" + str(i),
model.generator.generate(final1 * (1 - f) + final2 * f).numpy()[0])
for i in range(5):
dataset.save_image(folder, "A.noise" + str(i),
model.generator.generate(final1).numpy()[0])
dataset.save_image(folder, "B.noise" + str(i),
model.generator.generate(final2).numpy()[0])
dataset.save_image(
folder, "0.noise" + str(i),
model.generator.generate(np.zeros(
(1, model.input_size))).numpy()[0])
plot_confidence()
gen_iterations.append(g_iterations)
disc_iterations.append(d_iterations)
if epoch % output_freq == 0:
print("Finished epoch", epoch, "iteration", iterations)
print(
" generator: ", g_iterations,
"iterations, loss: %.6f and false positive rate: %.3f" %
(g_loss, g_acc))
print(
" discriminator:", d_iterations,
"iterations, loss: %.6f and accuracy: %.3f" % (d_loss, d_acc))
# 3. Output
if epoch % save_freq == 0:
test_image = model.generator.generate(test_seed)
dataset.save_image(resolution, epoch, test_image.numpy()[0])
epoch += 1
iterations += 1
if model.discriminator.resolution == max_resolution:
save_final("final")
save_model()
break
else:
save_final("final" + str(model.discriminator.resolution))
print("Adding resolution: ",
model.discriminator.resolution * 2,
"x",
model.discriminator.resolution * 2,
sep="")
model.add_resolution(model.generator, model.discriminator)
gan = tf.keras.Model(
def train_sp_data(noise_data, graph_name, class_name, epochs):
if class_name.find('sp') != -1:
cls = 'sp'
else:
cls = 'gaussian'
# 训练
with tf.Session() as sess:
merged = tf.summary.merge_all()
saver = tf.train.Saver()
lr = 1e-3
sess.run(tf.global_variables_initializer())
writer = tf.summary.FileWriter('./' + graph_name,
sess.graph) #再graph文件夹下使用默认图
for e in range(epochs): #每一个epoch
summary, train_loss, _ = sess.run(
[merged, cost, opt],
feed_dict={
inputs_: noise_data,
targets_: images_data,
learning_rate: lr * (0.997)**e
})
#每个epoch展示一下loss
print('Epoch : ', e, ' Training Cost : ', train_loss,
' Learning Rate : ',
lr * (0.997)**e)
saver.save(sess, './encode_model/')
writer.add_summary(summary, e)
#展示去噪之前和去噪之后的照片
fig, axes = plt.subplots(nrows=2,
ncols=11,
sharex=True,
sharey=True,
figsize=(20, 4))
in_imgs = np.copy(noise_data) #原图
#去噪后重塑的图片
reconstructed = sess.run(decoded,
feed_dict={
inputs_: in_imgs,
targets_: images_data
})
for is_noise_or_recon, (images, row) in enumerate(
zip([in_imgs, reconstructed], axes)):
for index, (img, ax) in enumerate(zip(images, row)):
if is_noise_or_recon == 0: #说明是原噪声图
psnr_score = tf.image.psnr(noise_data[index],
images_data[index],
max_val=1)
psnr_score = sess.run(psnr_score)
ssim1 = ssim(noise_data[index],
images_data[index],
multichannel=True)
ax.set_title(
str(round(ssim1, 2)) + ' ' + str(round(psnr_score, 2)))
else: #说明是还原图
psnr_score = tf.image.psnr(reconstructed[index],
images_data[index],
max_val=1)
psnr_score = sess.run(psnr_score)
ssim1 = ssim(reconstructed[index],
images_data[index],
multichannel=True)
ax.set_title(
str(round(ssim1, 2)) + ' ' + str(round(psnr_score, 2)))
ax.imshow(img.reshape((256, 256)), cmap='Greys_r')
#使用opencv保存的图片必须是0到255的uint8
img_save = np.array(img.reshape((256, 256)) * 255.0,
dtype=np.uint8)
if cls == 'sp': #说明是对椒盐噪声进行去噪
dataset.save_image(img_save, 'Image_outDenoise\\',
str(index + 1) + '.jpg')
else: #说明是对高斯噪声进行去噪
dataset.save_image(img_save,
'Image_Gaussian_outDenoising\\',
str(index + 1) + '.jpg')
ax.get_xaxis().set_visible(False)
ax.get_yaxis().set_visible(False)
fig.tight_layout(pad=0.01)
plt.show()
#pan = torch.Tensor(pan)
#D1 = mtc.D_lambda_torch(img, blur).numpy()
#D2 = mtc.D_s_torch(img, blur, pan, pan_l).numpy()
D_lambda.append(D1)
D_s.append(D2)
QNR.append((1 - D1) * (1 - D2))
i = i + 1
x = x + blk // 2
y = y + blk // 2
out = out / cnt
if a.save == 1:
save_image("{}/{}.tif".format(a.save_path, a.model),
out.transpose(2, 0, 1), 4)
if a.ref == 1:
print("SAM, CC, sCC, ERGAS, Q4, UIQC, SSIM, MPSNR")
print ("%.4lf+-%.4lf, %.4lf+-%.4lf, %.4lf+-%.4lf, %.4lf+-%.4lf, %.4lf+-%.4lf, %.4lf+-%.4lf, %.4lf+-%.4lf, %.4lf+-%.4lf " % ( \
np.mean(SAM), np.std(SAM), np.mean(CC), np.std(CC), np.mean(SCC), np.std(SCC), np.mean(ERGAS), np.std(ERGAS), \
np.mean(Q4), np.std(Q4), np.mean(UIQC), np.std(UIQC), np.mean(SSIM), np.std(SSIM), np.mean(MPSNR), np.std(MPSNR) \
) )
elif a.ref == 0:
print("D_lambda, D_s, QNR")
print ("%.4lf+-%.4lf, %.4lf+-%.4lf, %.4lf+-%.4lf " % ( \
np.mean(D_lambda), np.std(D_lambda), np.mean(D_s), np.std(D_s), np.mean(QNR), np.std(QNR) \
) )
print("done")
def train(opt):
plotDir = os.path.join(opt.exp_dir,opt.exp_name,'plots')
if not os.path.exists(plotDir):
os.makedirs(plotDir)
lib.print_model_settings(locals().copy())
""" dataset preparation """
if not opt.data_filtering_off:
print('Filtering the images containing characters which are not in opt.character')
print('Filtering the images whose label is longer than opt.batch_max_length')
# see https://github.com/clovaai/deep-text-recognition-benchmark/blob/6593928855fb7abb999a99f428b3e4477d4ae356/dataset.py#L130
opt.select_data = opt.select_data.split('-')
opt.batch_ratio = opt.batch_ratio.split('-')
#considering the real images for discriminator
opt.batch_size = opt.batch_size*2
train_dataset = Batch_Balanced_Dataset(opt)
log = open(os.path.join(opt.exp_dir,opt.exp_name,'log_dataset.txt'), 'a')
AlignCollate_valid = AlignCollate(imgH=opt.imgH, imgW=opt.imgW, keep_ratio_with_pad=opt.PAD)
valid_dataset, valid_dataset_log = hierarchical_dataset(root=opt.valid_data, opt=opt)
valid_loader = torch.utils.data.DataLoader(
valid_dataset, batch_size=opt.batch_size,
shuffle=False, # 'True' to check training progress with validation function.
num_workers=int(opt.workers),
collate_fn=AlignCollate_valid, pin_memory=True)
log.write(valid_dataset_log)
print('-' * 80)
log.write('-' * 80 + '\n')
log.close()
""" model configuration """
if 'CTC' in opt.Prediction:
converter = CTCLabelConverter(opt.character)
else:
converter = AttnLabelConverter(opt.character)
opt.num_class = len(converter.character)
if opt.rgb:
opt.input_channel = 3
model = AdaINGen(opt)
ocrModel = Model(opt)
disModel = MsImageDisV1(opt)
print('model input parameters', opt.imgH, opt.imgW, opt.num_fiducial, opt.input_channel, opt.output_channel,
opt.hidden_size, opt.num_class, opt.batch_max_length, opt.Transformation, opt.FeatureExtraction,
opt.SequenceModeling, opt.Prediction)
# weight initialization
for currModel in [model, ocrModel, disModel]:
for name, param in currModel.named_parameters():
if 'localization_fc2' in name:
print(f'Skip {name} as it is already initialized')
continue
try:
if 'bias' in name:
init.constant_(param, 0.0)
elif 'weight' in name:
init.kaiming_normal_(param)
except Exception as e: # for batchnorm.
if 'weight' in name:
param.data.fill_(1)
continue
# data parallel for multi-GPU
ocrModel = torch.nn.DataParallel(ocrModel).to(device)
if not opt.ocrFixed:
ocrModel.train()
else:
ocrModel.module.Transformation.eval()
ocrModel.module.FeatureExtraction.eval()
ocrModel.module.AdaptiveAvgPool.eval()
# ocrModel.module.SequenceModeling.eval()
ocrModel.module.Prediction.eval()
model = torch.nn.DataParallel(model).to(device)
model.train()
disModel = torch.nn.DataParallel(disModel).to(device)
disModel.train()
if opt.modelFolderFlag:
if len(glob.glob(os.path.join(opt.exp_dir,opt.exp_name,"iter_*_synth.pth")))>0:
opt.saved_synth_model = glob.glob(os.path.join(opt.exp_dir,opt.exp_name,"iter_*_synth.pth"))[-1]
if len(glob.glob(os.path.join(opt.exp_dir,opt.exp_name,"iter_*_dis.pth")))>0:
opt.saved_dis_model = glob.glob(os.path.join(opt.exp_dir,opt.exp_name,"iter_*_dis.pth"))[-1]
#loading pre-trained model
if opt.saved_ocr_model != '' and opt.saved_ocr_model != 'None':
print(f'loading pretrained ocr model from {opt.saved_ocr_model}')
if opt.FT:
ocrModel.load_state_dict(torch.load(opt.saved_ocr_model), strict=False)
else:
ocrModel.load_state_dict(torch.load(opt.saved_ocr_model))
print("OCRModel:")
print(ocrModel)
if opt.saved_synth_model != '' and opt.saved_synth_model != 'None':
print(f'loading pretrained synth model from {opt.saved_synth_model}')
if opt.FT:
model.load_state_dict(torch.load(opt.saved_synth_model), strict=False)
else:
model.load_state_dict(torch.load(opt.saved_synth_model))
print("SynthModel:")
print(model)
if opt.saved_dis_model != '' and opt.saved_dis_model != 'None':
print(f'loading pretrained discriminator model from {opt.saved_dis_model}')
if opt.FT:
disModel.load_state_dict(torch.load(opt.saved_dis_model), strict=False)
else:
disModel.load_state_dict(torch.load(opt.saved_dis_model))
print("DisModel:")
print(disModel)
""" setup loss """
if 'CTC' in opt.Prediction:
ocrCriterion = torch.nn.CTCLoss(zero_infinity=True).to(device)
else:
ocrCriterion = torch.nn.CrossEntropyLoss(ignore_index=0).to(device) # ignore [GO] token = ignore index 0
recCriterion = torch.nn.L1Loss()
styleRecCriterion = torch.nn.L1Loss()
# loss averager
loss_avg_ocr = Averager()
loss_avg = Averager()
loss_avg_dis = Averager()
loss_avg_ocrRecon_1 = Averager()
loss_avg_ocrRecon_2 = Averager()
loss_avg_gen = Averager()
loss_avg_imgRecon = Averager()
loss_avg_styRecon = Averager()
##---------------------------------------##
# filter that only require gradient decent
filtered_parameters = []
params_num = []
for p in filter(lambda p: p.requires_grad, model.parameters()):
filtered_parameters.append(p)
params_num.append(np.prod(p.size()))
print('Trainable params num : ', sum(params_num))
# [print(name, p.numel()) for name, p in filter(lambda p: p[1].requires_grad, model.named_parameters())]
# setup optimizer
if opt.optim=='adam':
optimizer = optim.Adam(filtered_parameters, lr=opt.lr, betas=(opt.beta1, opt.beta2), weight_decay=opt.weight_decay)
else:
optimizer = optim.Adadelta(filtered_parameters, lr=opt.lr, rho=opt.rho, eps=opt.eps, weight_decay=opt.weight_decay)
print("SynthOptimizer:")
print(optimizer)
#filter parameters for OCR training
ocr_filtered_parameters = []
ocr_params_num = []
for p in filter(lambda p: p.requires_grad, ocrModel.parameters()):
ocr_filtered_parameters.append(p)
ocr_params_num.append(np.prod(p.size()))
print('OCR Trainable params num : ', sum(ocr_params_num))
# setup optimizer
if opt.optim=='adam':
ocr_optimizer = optim.Adam(ocr_filtered_parameters, lr=opt.lr, betas=(opt.beta1, opt.beta2), weight_decay=opt.weight_decay)
else:
ocr_optimizer = optim.Adadelta(ocr_filtered_parameters, lr=opt.lr, rho=opt.rho, eps=opt.eps, weight_decay=opt.weight_decay)
print("OCROptimizer:")
print(ocr_optimizer)
#filter parameters for OCR training
dis_filtered_parameters = []
dis_params_num = []
for p in filter(lambda p: p.requires_grad, disModel.parameters()):
dis_filtered_parameters.append(p)
dis_params_num.append(np.prod(p.size()))
print('Dis Trainable params num : ', sum(dis_params_num))
# setup optimizer
if opt.optim=='adam':
dis_optimizer = optim.Adam(dis_filtered_parameters, lr=opt.lr, betas=(opt.beta1, opt.beta2), weight_decay=opt.weight_decay)
else:
dis_optimizer = optim.Adadelta(dis_filtered_parameters, lr=opt.lr, rho=opt.rho, eps=opt.eps, weight_decay=opt.weight_decay)
print("DisOptimizer:")
print(dis_optimizer)
##---------------------------------------##
""" final options """
with open(os.path.join(opt.exp_dir,opt.exp_name,'opt.txt'), 'a') as opt_file:
opt_log = '------------ Options -------------\n'
args = vars(opt)
for k, v in args.items():
opt_log += f'{str(k)}: {str(v)}\n'
opt_log += '---------------------------------------\n'
print(opt_log)
opt_file.write(opt_log)
""" start training """
start_iter = 0
if opt.saved_synth_model != '' and opt.saved_synth_model != 'None':
try:
start_iter = int(opt.saved_synth_model.split('_')[-2].split('.')[0])
print(f'continue to train, start_iter: {start_iter}')
except:
pass
#get schedulers
scheduler = get_scheduler(optimizer,opt)
ocr_scheduler = get_scheduler(ocr_optimizer,opt)
dis_scheduler = get_scheduler(dis_optimizer,opt)
start_time = time.time()
best_accuracy = -1
best_norm_ED = -1
best_accuracy_ocr = -1
best_norm_ED_ocr = -1
iteration = start_iter
cntr=0
while(True):
# train part
if opt.lr_policy !="None":
scheduler.step()
ocr_scheduler.step()
dis_scheduler.step()
image_tensors_all, labels_1_all, labels_2_all = train_dataset.get_batch()
# ## comment
# pdb.set_trace()
# for imgCntr in range(image_tensors.shape[0]):
# save_image(tensor2im(image_tensors[imgCntr]),'temp/'+str(imgCntr)+'.png')
# pdb.set_trace()
# ###
# print(cntr)
cntr+=1
disCnt = int(image_tensors_all.size(0)/2)
image_tensors, image_tensors_real, labels_gt, labels_2 = image_tensors_all[:disCnt], image_tensors_all[disCnt:disCnt+disCnt], labels_1_all[:disCnt], labels_2_all[:disCnt]
image = image_tensors.to(device)
image_real = image_tensors_real.to(device)
batch_size = image.size(0)
##-----------------------------------##
#generate text(labels) from ocr.forward
if opt.ocrFixed:
# ocrModel.eval()
length_for_pred = torch.IntTensor([opt.batch_max_length] * batch_size).to(device)
text_for_pred = torch.LongTensor(batch_size, opt.batch_max_length + 1).fill_(0).to(device)
if 'CTC' in opt.Prediction:
preds = ocrModel(image, text_for_pred)
preds = preds[:, :text_for_loss.shape[1] - 1, :]
preds_size = torch.IntTensor([preds.size(1)] * batch_size)
_, preds_index = preds.max(2)
labels_1 = converter.decode(preds_index.data, preds_size.data)
else:
preds = ocrModel(image, text_for_pred, is_train=False)
_, preds_index = preds.max(2)
labels_1 = converter.decode(preds_index, length_for_pred)
for idx, pred in enumerate(labels_1):
pred_EOS = pred.find('[s]')
labels_1[idx] = pred[:pred_EOS] # prune after "end of sentence" token ([s])
# ocrModel.train()
else:
labels_1 = labels_gt
##-----------------------------------##
text_1, length_1 = converter.encode(labels_1, batch_max_length=opt.batch_max_length)
text_2, length_2 = converter.encode(labels_2, batch_max_length=opt.batch_max_length)
#forward pass from style and word generator
images_recon_1, images_recon_2, style = model(image, text_1, text_2)
if 'CTC' in opt.Prediction:
if not opt.ocrFixed:
#ocr training with orig image
preds_ocr = ocrModel(image, text_1)
preds_size_ocr = torch.IntTensor([preds_ocr.size(1)] * batch_size)
preds_ocr = preds_ocr.log_softmax(2).permute(1, 0, 2)
ocrCost_train = ocrCriterion(preds_ocr, text_1, preds_size_ocr, length_1)
#content loss for reconstructed images
preds_1 = ocrModel(images_recon_1, text_1)
preds_size_1 = torch.IntTensor([preds_1.size(1)] * batch_size)
preds_1 = preds_1.log_softmax(2).permute(1, 0, 2)
preds_2 = ocrModel(images_recon_2, text_2)
preds_size_2 = torch.IntTensor([preds_2.size(1)] * batch_size)
preds_2 = preds_2.log_softmax(2).permute(1, 0, 2)
ocrCost_1 = ocrCriterion(preds_1, text_1, preds_size_1, length_1)
ocrCost_2 = ocrCriterion(preds_2, text_2, preds_size_2, length_2)
# ocrCost = 0.5*( ocrCost_1 + ocrCost_2 )
else:
if not opt.ocrFixed:
#ocr training with orig image
preds_ocr = ocrModel(image, text_1[:, :-1]) # align with Attention.forward
target_ocr = text_1[:, 1:] # without [GO] Symbol
ocrCost_train = ocrCriterion(preds_ocr.view(-1, preds_ocr.shape[-1]), target_ocr.contiguous().view(-1))
#content loss for reconstructed images
preds_1 = ocrModel(images_recon_1, text_1[:, :-1], is_train=False) # align with Attention.forward
target_1 = text_1[:, 1:] # without [GO] Symbol
preds_2 = ocrModel(images_recon_2, text_2[:, :-1], is_train=False) # align with Attention.forward
target_2 = text_2[:, 1:] # without [GO] Symbol
ocrCost_1 = ocrCriterion(preds_1.view(-1, preds_1.shape[-1]), target_1.contiguous().view(-1))
ocrCost_2 = ocrCriterion(preds_2.view(-1, preds_2.shape[-1]), target_2.contiguous().view(-1))
# ocrCost = 0.5*(ocrCost_1+ocrCost_2)
if not opt.ocrFixed:
#training OCR
ocrModel.zero_grad()
ocrCost_train.backward()
# torch.nn.utils.clip_grad_norm_(ocrModel.parameters(), opt.grad_clip) # gradient clipping with 5 (Default)
ocr_optimizer.step()
#if ocr is fixed; ignore this loss
loss_avg_ocr.add(ocrCost_train)
else:
loss_avg_ocr.add(torch.tensor(0.0))
#Domain discriminator: Dis update
disModel.zero_grad()
disCost = opt.disWeight*0.5*(disModel.module.calc_dis_loss(images_recon_1.detach(), image_real) + disModel.module.calc_dis_loss(images_recon_2.detach(), image))
disCost.backward()
# torch.nn.utils.clip_grad_norm_(disModel.parameters(), opt.grad_clip) # gradient clipping with 5 (Default)
dis_optimizer.step()
loss_avg_dis.add(disCost)
# #[Style Encoder] + [Word Generator] update
#Adversarial loss
disGenCost = 0.5*(disModel.module.calc_gen_loss(images_recon_1)+disModel.module.calc_gen_loss(images_recon_2))
#Input reconstruction loss
recCost = recCriterion(images_recon_1,image)
#Pair style reconstruction loss
if opt.styleReconWeight == 0.0:
styleRecCost = torch.tensor(0.0)
else:
if opt.styleDetach:
styleRecCost = styleRecCriterion(model(images_recon_2, None, None, styleFlag=True), style.detach())
else:
styleRecCost = styleRecCriterion(model(images_recon_2, None, None, styleFlag=True), style)
#OCR Content cost
ocrCost = 0.5*(ocrCost_1+ocrCost_2)
cost = opt.ocrWeight*ocrCost + opt.reconWeight*recCost + opt.disWeight*disGenCost + opt.styleReconWeight*styleRecCost
model.zero_grad()
ocrModel.zero_grad()
disModel.zero_grad()
cost.backward()
# torch.nn.utils.clip_grad_norm_(model.parameters(), opt.grad_clip) # gradient clipping with 5 (Default)
optimizer.step()
loss_avg.add(cost)
#Individual losses
loss_avg_ocrRecon_1.add(opt.ocrWeight*0.5*ocrCost_1)
loss_avg_ocrRecon_2.add(opt.ocrWeight*0.5*ocrCost_2)
loss_avg_gen.add(opt.disWeight*disGenCost)
loss_avg_imgRecon.add(opt.reconWeight*recCost)
loss_avg_styRecon.add(opt.styleReconWeight*styleRecCost)
# validation part
if (iteration + 1) % opt.valInterval == 0 or iteration == 0: # To see training progress, we also conduct validation when 'iteration == 0'
#Save training images
os.makedirs(os.path.join(opt.exp_dir,opt.exp_name,'trainImages',str(iteration)), exist_ok=True)
for trImgCntr in range(batch_size):
try:
save_image(tensor2im(image[trImgCntr].detach()),os.path.join(opt.exp_dir,opt.exp_name,'trainImages',str(iteration),str(trImgCntr)+'_input_'+labels_gt[trImgCntr]+'.png'))
save_image(tensor2im(images_recon_1[trImgCntr].detach()),os.path.join(opt.exp_dir,opt.exp_name,'trainImages',str(iteration),str(trImgCntr)+'_recon_'+labels_1[trImgCntr]+'.png'))
save_image(tensor2im(images_recon_2[trImgCntr].detach()),os.path.join(opt.exp_dir,opt.exp_name,'trainImages',str(iteration),str(trImgCntr)+'_pair_'+labels_2[trImgCntr]+'.png'))
except:
print('Warning while saving training image')
elapsed_time = time.time() - start_time
# for log
with open(os.path.join(opt.exp_dir,opt.exp_name,'log_train.txt'), 'a') as log:
model.eval()
ocrModel.module.Transformation.eval()
ocrModel.module.FeatureExtraction.eval()
ocrModel.module.AdaptiveAvgPool.eval()
ocrModel.module.SequenceModeling.eval()
ocrModel.module.Prediction.eval()
disModel.eval()
with torch.no_grad():
valid_loss, current_accuracy, current_norm_ED, preds, confidence_score, labels, infer_time, length_of_data = validation_synth_lrw_res(
iteration, model, ocrModel, disModel, recCriterion, styleRecCriterion, ocrCriterion, valid_loader, converter, opt)
model.train()
if not opt.ocrFixed:
ocrModel.train()
else:
# ocrModel.module.Transformation.eval()
# ocrModel.module.FeatureExtraction.eval()
# ocrModel.module.AdaptiveAvgPool.eval()
ocrModel.module.SequenceModeling.train()
# ocrModel.module.Prediction.eval()
disModel.train()
# training loss and validation loss
loss_log = f'[{iteration+1}/{opt.num_iter}] Train OCR loss: {loss_avg_ocr.val():0.5f}, Train Synth loss: {loss_avg.val():0.5f}, Train Dis loss: {loss_avg_dis.val():0.5f}, Valid OCR loss: {valid_loss[0]:0.5f}, Valid Synth loss: {valid_loss[1]:0.5f}, Valid Dis loss: {valid_loss[2]:0.5f}, Elapsed_time: {elapsed_time:0.5f}'
current_model_log_ocr = f'{"Current_accuracy_OCR":17s}: {current_accuracy[0]:0.3f}, {"Current_norm_ED_OCR":17s}: {current_norm_ED[0]:0.2f}'
current_model_log_1 = f'{"Current_accuracy_recon":17s}: {current_accuracy[1]:0.3f}, {"Current_norm_ED_recon":17s}: {current_norm_ED[1]:0.2f}'
current_model_log_2 = f'{"Current_accuracy_pair":17s}: {current_accuracy[2]:0.3f}, {"Current_norm_ED_pair":17s}: {current_norm_ED[2]:0.2f}'
#plotting
lib.plot.plot(os.path.join(plotDir,'Train-OCR-Loss'), loss_avg_ocr.val().item())
lib.plot.plot(os.path.join(plotDir,'Train-Synth-Loss'), loss_avg.val().item())
lib.plot.plot(os.path.join(plotDir,'Train-Dis-Loss'), loss_avg_dis.val().item())
lib.plot.plot(os.path.join(plotDir,'Train-OCR-Recon1-Loss'), loss_avg_ocrRecon_1.val().item())
lib.plot.plot(os.path.join(plotDir,'Train-OCR-Recon2-Loss'), loss_avg_ocrRecon_2.val().item())
lib.plot.plot(os.path.join(plotDir,'Train-Gen-Loss'), loss_avg_gen.val().item())
lib.plot.plot(os.path.join(plotDir,'Train-ImgRecon1-Loss'), loss_avg_imgRecon.val().item())
lib.plot.plot(os.path.join(plotDir,'Train-StyRecon2-Loss'), loss_avg_styRecon.val().item())
lib.plot.plot(os.path.join(plotDir,'Valid-OCR-Loss'), valid_loss[0].item())
lib.plot.plot(os.path.join(plotDir,'Valid-Synth-Loss'), valid_loss[1].item())
lib.plot.plot(os.path.join(plotDir,'Valid-Dis-Loss'), valid_loss[2].item())
lib.plot.plot(os.path.join(plotDir,'Valid-OCR-Recon1-Loss'), valid_loss[3].item())
lib.plot.plot(os.path.join(plotDir,'Valid-OCR-Recon2-Loss'), valid_loss[4].item())
lib.plot.plot(os.path.join(plotDir,'Valid-Gen-Loss'), valid_loss[5].item())
lib.plot.plot(os.path.join(plotDir,'Valid-ImgRecon1-Loss'), valid_loss[6].item())
lib.plot.plot(os.path.join(plotDir,'Valid-StyRecon2-Loss'), valid_loss[7].item())
lib.plot.plot(os.path.join(plotDir,'Orig-OCR-WordAccuracy'), current_accuracy[0])
lib.plot.plot(os.path.join(plotDir,'Recon-OCR-WordAccuracy'), current_accuracy[1])
lib.plot.plot(os.path.join(plotDir,'Pair-OCR-WordAccuracy'), current_accuracy[2])
lib.plot.plot(os.path.join(plotDir,'Orig-OCR-CharAccuracy'), current_norm_ED[0])
lib.plot.plot(os.path.join(plotDir,'Recon-OCR-CharAccuracy'), current_norm_ED[1])
lib.plot.plot(os.path.join(plotDir,'Pair-OCR-CharAccuracy'), current_norm_ED[2])
# keep best accuracy model (on valid dataset)
if current_accuracy[1] > best_accuracy:
best_accuracy = current_accuracy[1]
torch.save(model.state_dict(), os.path.join(opt.exp_dir,opt.exp_name,'best_accuracy.pth'))
torch.save(disModel.state_dict(), os.path.join(opt.exp_dir,opt.exp_name,'best_accuracy_dis.pth'))
if current_norm_ED[1] > best_norm_ED:
best_norm_ED = current_norm_ED[1]
torch.save(model.state_dict(), os.path.join(opt.exp_dir,opt.exp_name,'best_norm_ED.pth'))
torch.save(disModel.state_dict(), os.path.join(opt.exp_dir,opt.exp_name,'best_norm_ED_dis.pth'))
best_model_log = f'{"Best_accuracy_Recon":17s}: {best_accuracy:0.3f}, {"Best_norm_ED_Recon":17s}: {best_norm_ED:0.2f}'
# keep best accuracy model (on valid dataset)
if current_accuracy[0] > best_accuracy_ocr:
best_accuracy_ocr = current_accuracy[0]
if not opt.ocrFixed:
torch.save(ocrModel.state_dict(), os.path.join(opt.exp_dir,opt.exp_name,'best_accuracy_ocr.pth'))
if current_norm_ED[0] > best_norm_ED_ocr:
best_norm_ED_ocr = current_norm_ED[0]
if not opt.ocrFixed:
torch.save(ocrModel.state_dict(), os.path.join(opt.exp_dir,opt.exp_name,'best_norm_ED_ocr.pth'))
best_model_log_ocr = f'{"Best_accuracy_ocr":17s}: {best_accuracy_ocr:0.3f}, {"Best_norm_ED_ocr":17s}: {best_norm_ED_ocr:0.2f}'
loss_model_log = f'{loss_log}\n{current_model_log_ocr}\n{current_model_log_1}\n{current_model_log_2}\n{best_model_log_ocr}\n{best_model_log}'
print(loss_model_log)
log.write(loss_model_log + '\n')
# show some predicted results
dashed_line = '-' * 80
head = f'{"Ground Truth":32s} | {"Prediction":25s} | Confidence Score & T/F'
predicted_result_log = f'{dashed_line}\n{head}\n{dashed_line}\n'
for gt_ocr, pred_ocr, confidence_ocr, gt_1, pred_1, confidence_1, gt_2, pred_2, confidence_2 in zip(labels[0][:5], preds[0][:5], confidence_score[0][:5], labels[1][:5], preds[1][:5], confidence_score[1][:5], labels[2][:5], preds[2][:5], confidence_score[2][:5]):
if 'Attn' in opt.Prediction:
# gt_ocr = gt_ocr[:gt_ocr.find('[s]')]
pred_ocr = pred_ocr[:pred_ocr.find('[s]')]
# gt_1 = gt_1[:gt_1.find('[s]')]
pred_1 = pred_1[:pred_1.find('[s]')]
# gt_2 = gt_2[:gt_2.find('[s]')]
pred_2 = pred_2[:pred_2.find('[s]')]
predicted_result_log += f'{"ocr"}: {gt_ocr:27s} | {pred_ocr:25s} | {confidence_ocr:0.4f}\t{str(pred_ocr == gt_ocr)}\n'
predicted_result_log += f'{"recon"}: {gt_1:25s} | {pred_1:25s} | {confidence_1:0.4f}\t{str(pred_1 == gt_1)}\n'
predicted_result_log += f'{"pair"}: {gt_2:26s} | {pred_2:25s} | {confidence_2:0.4f}\t{str(pred_2 == gt_2)}\n'
predicted_result_log += f'{dashed_line}'
print(predicted_result_log)
log.write(predicted_result_log + '\n')
loss_avg_ocr.reset()
loss_avg.reset()
loss_avg_dis.reset()
loss_avg_ocrRecon_1.reset()
loss_avg_ocrRecon_2.reset()
loss_avg_gen.reset()
loss_avg_imgRecon.reset()
loss_avg_styRecon.reset()
lib.plot.flush()
lib.plot.tick()
# save model per 1e+5 iter.
if (iteration) % 1e+5 == 0:
torch.save(
model.state_dict(), os.path.join(opt.exp_dir,opt.exp_name,'iter_'+str(iteration+1)+'_synth.pth'))
if not opt.ocrFixed:
torch.save(
ocrModel.state_dict(), os.path.join(opt.exp_dir,opt.exp_name,'iter_'+str(iteration+1)+'_ocr.pth'))
torch.save(
disModel.state_dict(), os.path.join(opt.exp_dir,opt.exp_name,'iter_'+str(iteration+1)+'_dis.pth'))
if (iteration + 1) == opt.num_iter:
print('end the training')
sys.exit()
iteration += 1
imgMul_o_u = upsample(imgMul_o) # h * w * 4
imgPan_o = rawPan # h * w
imgPan_o_d = upsample(downsample(imgPan_o)) # h * w
imgMul = imgMul.transpose(2, 0, 1)
imgLR = imgLR.transpose(2, 0, 1)
imgLR_u = imgLR_u.transpose(2, 0, 1)
imgMul_o = imgMul_o.transpose(2, 0, 1)
imgMul_o_u = imgMul_o_u.transpose(2, 0, 1)
if i == 0:
a = h // 4
b = w // 4
a = a // 50 * 50
b = b // 50 * 50
save_image(newMul, imgMul[:, :a * 4, :b * 4], 4)
save_image(newLR, imgLR[:, :a, :b], 4)
save_image(newLR_u, imgLR_u[:, :a * 4, :b * 4], 4)
save_image(newPan, imgPan[:a * 4, :b * 4], 1)
save_image(newPan_d, imgPan_d[:a * 4, :b * 4], 1)
# a = h # keep raw size
# b = w
# a = a // 50 * 50
# b = b // 50 * 50
save_image(newMul_o, imgMul_o[:, :a, :b], 4)
save_image(newMul_o_u, imgMul_o_u[:, :a * 4, :b * 4], 4)
save_image(newPan_o, imgPan_o[:a * 4, :b * 4], 1)
save_image(newPan_o_d, imgPan_o_d[:a * 4, :b * 4], 1)
else:
save_image(newMul, imgMul, 4)
