def train(opt):
""" dataset preparation """
opt.select_data = opt.select_data.split('-')
opt.batch_ratio = opt.batch_ratio.split('-')
train_dataset = Batch_Balanced_Dataset(opt)
AlignCollate_valid = AlignCollate(imgH=opt.imgH,
imgW=opt.imgW,
keep_ratio_with_pad=opt.PAD)
valid_dataset = hierarchical_dataset(root=opt.valid_data, opt=opt)
valid_loader = torch.utils.data.DataLoader(
valid_dataset,
batch_size=opt.batch_size,
# 'True' to check training progress with validation function.
shuffle=True,
num_workers=int(opt.workers),
collate_fn=AlignCollate_valid,
pin_memory=True)
print('-' * 80)
""" model configuration """
if 'Transformer' in opt.SequenceModeling:
converter = TransformerLabelConverter(opt.character)
elif '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 = Model(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 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
""" setup loss """
if 'Transformer' in opt.SequenceModeling:
criterion = transformer_loss
elif 'CTC' in opt.Prediction:
criterion = torch.nn.CTCLoss(zero_infinity=True).cuda()
else:
# ignore [GO] token = ignore index 0
criterion = torch.nn.CrossEntropyLoss(ignore_index=0).cuda()
# loss averager
loss_avg = 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))
elif 'Transformer' in opt.SequenceModeling and opt.use_scheduled_optim:
optimizer = optim.Adam(filtered_parameters,
betas=(0.9, 0.98),
eps=1e-09)
optimizer_schedule = ScheduledOptim(optimizer, opt.d_model,
opt.n_warmup_steps)
else:
optimizer = optim.Adadelta(filtered_parameters,
lr=opt.lr,
rho=opt.rho,
eps=opt.eps)
print("Optimizer:")
print(optimizer)
""" final options """
# print(opt)
with open(f'./saved_models/{opt.experiment_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
start_time = time.time()
best_accuracy = -1
best_norm_ED = 1e+6
pickle.load = partial(pickle.load, encoding="latin1")
pickle.Unpickler = partial(pickle.Unpickler, encoding="latin1")
if opt.load_weights != '' and check_isfile(opt.load_weights):
# load pretrained weights but ignore layers that don't match in size
checkpoint = torch.load(opt.load_weights, pickle_module=pickle)
if type(checkpoint) == dict:
pretrain_dict = checkpoint['state_dict']
else:
pretrain_dict = checkpoint
model_dict = model.state_dict()
pretrain_dict = {
k: v
for k, v in pretrain_dict.items()
if k in model_dict and model_dict[k].size() == v.size()
}
model_dict.update(pretrain_dict)
model.load_state_dict(model_dict)
print("Loaded pretrained weights from '{}'".format(opt.load_weights))
del checkpoint
torch.cuda.empty_cache()
if opt.continue_model != '':
print(f'loading pretrained model from {opt.continue_model}')
checkpoint = torch.load(opt.continue_model)
print(checkpoint.keys())
model.load_state_dict(checkpoint['state_dict'])
start_iter = checkpoint['step'] + 1
print('continue to train start_iter: ', start_iter)
if 'optimizer' in checkpoint.keys():
optimizer.load_state_dict(checkpoint['optimizer'])
for state in optimizer.state.values():
for k, v in state.items():
if isinstance(v, torch.Tensor):
state[k] = v.cuda()
if 'best_accuracy' in checkpoint.keys():
best_accuracy = checkpoint['best_accuracy']
if 'best_norm_ED' in checkpoint.keys():
best_norm_ED = checkpoint['best_norm_ED']
del checkpoint
torch.cuda.empty_cache()
# data parallel for multi-GPU
model = torch.nn.DataParallel(model).cuda()
model.train()
print("Model size:", count_num_param(model), 'M')
if 'Transformer' in opt.SequenceModeling and opt.use_scheduled_optim:
optimizer_schedule.n_current_steps = start_iter
for i in tqdm(range(start_iter, opt.num_iter)):
for p in model.parameters():
p.requires_grad = True
cpu_images, cpu_texts = train_dataset.get_batch()
image = cpu_images.cuda()
if 'Transformer' in opt.SequenceModeling:
text, length, text_pos = converter.encode(cpu_texts,
opt.batch_max_length)
elif 'CTC' in opt.Prediction:
text, length = converter.encode(cpu_texts)
else:
text, length = converter.encode(cpu_texts, opt.batch_max_length)
batch_size = image.size(0)
if 'Transformer' in opt.SequenceModeling:
preds = model(image, text, tgt_pos=text_pos)
target = text[:, 1:] # without <s> Symbol
cost = criterion(preds.view(-1, preds.shape[-1]),
target.contiguous().view(-1))
elif 'CTC' in opt.Prediction:
preds = model(image, text).log_softmax(2)
preds_size = torch.IntTensor([preds.size(1)] * batch_size)
preds = preds.permute(1, 0, 2) # to use CTCLoss format
cost = criterion(preds, text, preds_size, length)
else:
preds = model(image, text)
target = text[:, 1:] # without [GO] Symbol
cost = criterion(preds.view(-1, preds.shape[-1]),
target.contiguous().view(-1))
model.zero_grad()
cost.backward()
if 'Transformer' in opt.SequenceModeling and opt.use_scheduled_optim:
optimizer_schedule.step_and_update_lr()
elif 'Transformer' in opt.SequenceModeling:
optimizer.step()
else:
# gradient clipping with 5 (Default)
torch.nn.utils.clip_grad_norm_(model.parameters(), opt.grad_clip)
optimizer.step()
loss_avg.add(cost)
# validation part
if i > 0 and (i + 1) % opt.valInterval == 0:
elapsed_time = time.time() - start_time
print(
f'[{i+1}/{opt.num_iter}] Loss: {loss_avg.val():0.5f} elapsed_time: {elapsed_time:0.5f}'
)
# for log
with open(f'./saved_models/{opt.experiment_name}/log_train.txt',
'a') as log:
log.write(
f'[{i+1}/{opt.num_iter}] Loss: {loss_avg.val():0.5f} elapsed_time: {elapsed_time:0.5f}\n'
)
loss_avg.reset()
model.eval()
with torch.no_grad():
valid_loss, current_accuracy, current_norm_ED, preds, gts, infer_time, length_of_data = validation(
model, criterion, valid_loader, converter, opt)
model.train()
for pred, gt in zip(preds[:5], gts[:5]):
if 'Transformer' in opt.SequenceModeling:
pred = pred[:pred.find('</s>')]
gt = gt[:gt.find('</s>')]
elif 'Attn' in opt.Prediction:
pred = pred[:pred.find('[s]')]
gt = gt[:gt.find('[s]')]
print(f'{pred:20s}, gt: {gt:20s}, {str(pred == gt)}')
log.write(
f'{pred:20s}, gt: {gt:20s}, {str(pred == gt)}\n')
valid_log = f'[{i+1}/{opt.num_iter}] valid loss: {valid_loss:0.5f}'
valid_log += f' accuracy: {current_accuracy:0.3f}, norm_ED: {current_norm_ED:0.2f}'
print(valid_log)
log.write(valid_log + '\n')
# keep best accuracy model
if current_accuracy > best_accuracy:
best_accuracy = current_accuracy
state_dict = model.module.state_dict()
save_checkpoint(
{
'best_accuracy': best_accuracy,
'state_dict': state_dict,
}, False,
f'./saved_models/{opt.experiment_name}/best_accuracy.pth'
)
if current_norm_ED < best_norm_ED:
best_norm_ED = current_norm_ED
state_dict = model.module.state_dict()
save_checkpoint(
{
'best_norm_ED': best_norm_ED,
'state_dict': state_dict,
}, False,
f'./saved_models/{opt.experiment_name}/best_norm_ED.pth'
)
# torch.save(
# model.state_dict(), f'./saved_models/{opt.experiment_name}/best_norm_ED.pth')
best_model_log = f'best_accuracy: {best_accuracy:0.3f}, best_norm_ED: {best_norm_ED:0.2f}'
print(best_model_log)
log.write(best_model_log + '\n')
# save model per 1000 iter.
if (i + 1) % 1000 == 0:
state_dict = model.module.state_dict()
optimizer_state_dict = optimizer.state_dict()
save_checkpoint(
{
'state_dict': state_dict,
'optimizer': optimizer_state_dict,
'step': i,
'best_accuracy': best_accuracy,
'best_norm_ED': best_norm_ED,
}, False,
f'./saved_models/{opt.experiment_name}/iter_{i+1}.pth')
def train(opt):
lib.print_model_settings(locals().copy())
# train_transform = transforms.Compose([
# # transforms.RandomResizedCrop(input_size),
# transforms.Resize((opt.imgH, opt.imgW)),
# # transforms.RandomHorizontalFlip(),
# transforms.ToTensor(),
# transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
# ])
# val_transform = transforms.Compose([
# transforms.Resize((opt.imgH, opt.imgW)),
# # transforms.CenterCrop(input_size),
# transforms.ToTensor(),
# transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
# ])
AlignFontCollateObj = AlignFontCollate(imgH=opt.imgH,
imgW=opt.imgW,
keep_ratio_with_pad=opt.PAD)
train_dataset = fontDataset(imgDir=opt.train_img_dir,
annFile=opt.train_ann_file,
transform=None,
numClasses=opt.numClasses)
train_loader = torch.utils.data.DataLoader(
train_dataset,
batch_size=opt.batch_size,
shuffle=
False, # 'True' to check training progress with validation function.
sampler=data_sampler(train_dataset,
shuffle=True,
distributed=opt.distributed),
num_workers=int(opt.workers),
collate_fn=AlignFontCollateObj,
pin_memory=True,
drop_last=False)
# numClasses = len(train_dataset.Idx2F)
numClasses = np.unique(train_dataset.fontIdx).size
train_loader = sample_data(train_loader)
print('-' * 80)
numTrainSamples = len(train_dataset)
# valid_dataset = LmdbStyleDataset(root=opt.valid_data, opt=opt)
valid_dataset = fontDataset(imgDir=opt.train_img_dir,
annFile=opt.val_ann_file,
transform=None,
F2Idx=train_dataset.F2Idx,
Idx2F=train_dataset.Idx2F,
numClasses=opt.numClasses)
valid_loader = torch.utils.data.DataLoader(
valid_dataset,
batch_size=opt.batch_size,
shuffle=
False, # 'True' to check training progress with validation function.
sampler=data_sampler(valid_dataset,
shuffle=False,
distributed=opt.distributed),
num_workers=int(opt.workers),
collate_fn=AlignFontCollateObj,
pin_memory=True,
drop_last=False)
numTestSamples = len(valid_dataset)
print('numClasses', numClasses)
print('numTrainSamples', numTrainSamples)
print('numTestSamples', numTestSamples)
vggFontModel = VGGFontModel(models.vgg19(pretrained=opt.preTrained),
numClasses).to(device)
for name, param in vggFontModel.classifier.named_parameters():
try:
if 'bias' in name:
init.constant_(param, 0.0)
elif 'weight' in name:
init.kaiming_normal_(param)
except Exception as e: # for batchnorm.
print('Exception in weight init' + name)
if 'weight' in name:
param.data.fill_(1)
continue
if opt.optim == "sgd":
print('SGD optimizer')
optimizer = optim.SGD(vggFontModel.parameters(),
lr=opt.lr,
momentum=0.9)
elif opt.optim == "adam":
print('Adam optimizer')
optimizer = optim.Adam(vggFontModel.parameters(), lr=opt.lr)
#get schedulers
scheduler = get_scheduler(optimizer, opt)
criterion = torch.nn.CrossEntropyLoss()
if opt.modelFolderFlag:
if len(
glob.glob(
os.path.join(opt.exp_dir, opt.exp_name,
"iter_*_vggFont.pth"))) > 0:
opt.saved_font_model = glob.glob(
os.path.join(opt.exp_dir, opt.exp_name,
"iter_*_vggFont.pth"))[-1]
## Loading pre-trained files
if opt.saved_font_model != '' and opt.saved_font_model != 'None':
print(f'loading pretrained synth model from {opt.saved_font_model}')
checkpoint = torch.load(opt.saved_font_model,
map_location=lambda storage, loc: storage)
vggFontModel.load_state_dict(checkpoint['vggFontModel'])
optimizer.load_state_dict(checkpoint["optimizer"])
scheduler.load_state_dict(checkpoint["scheduler"])
# print('Model Initialization')
#
# print('Loaded checkpoint')
if opt.distributed:
vggFontModel = torch.nn.parallel.DistributedDataParallel(
vggFontModel,
device_ids=[opt.local_rank],
output_device=opt.local_rank,
broadcast_buffers=False,
find_unused_parameters=True)
vggFontModel.train()
# print('Loaded distributed')
if opt.distributed:
vggFontModel_module = vggFontModel.module
else:
vggFontModel_module = vggFontModel
# print('Loading module')
# loss averager
loss_train = Averager()
loss_val = Averager()
train_acc = Averager()
val_acc = Averager()
train_acc_5 = Averager()
val_acc_5 = 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_font_model != '' and opt.saved_font_model != 'None':
try:
start_iter = int(opt.saved_font_model.split('_')[-2].split('.')[0])
print(f'continue to train, start_iter: {start_iter}')
except:
pass
iteration = start_iter
cntr = 0
# trainCorrect=0
# tCntr=0
while (True):
# print(cntr)
# train part
start_time = time.time()
if not opt.testFlag:
image_input_tensors, labels_gt = next(train_loader)
image_input_tensors = image_input_tensors.to(device)
labels_gt = labels_gt.view(-1).to(device)
preds = vggFontModel(image_input_tensors)
loss = criterion(preds, labels_gt)
vggFontModel.zero_grad()
loss.backward()
optimizer.step()
# _, preds_max = preds.max(dim=1)
# trainCorrect += (preds_max == labels_gt).sum()
# tCntr+=preds_max.shape[0]
acc1, acc5 = getNumCorrect(preds,
labels_gt,
topk=(1, min(numClasses, 5)))
train_acc.addScalar(acc1, preds.shape[0])
train_acc_5.addScalar(acc5, preds.shape[0])
loss_train.add(loss)
if opt.lr_policy != "None":
scheduler.step()
# print
if get_rank() == 0:
if (
iteration + 1
) % opt.valInterval == 0 or iteration == 0 or opt.testFlag: # To see training progress, we also conduct validation when 'iteration == 0'
#validation
# iCntr=torch.tensor(0.0).to(device)
# valCorrect=torch.tensor(0.0).to(device)
vggFontModel.eval()
print('Inside val', iteration)
for vCntr, (image_input_tensors,
labels_gt) in enumerate(valid_loader):
# print('vCntr--',vCntr)
if opt.debugFlag and vCntr > 2:
break
with torch.no_grad():
image_input_tensors = image_input_tensors.to(device)
labels_gt = labels_gt.view(-1).to(device)
preds = vggFontModel(image_input_tensors)
loss = criterion(preds, labels_gt)
loss_val.add(loss)
# _, preds_max = preds.max(dim=1)
# valCorrect += (preds_max == labels_gt).sum()
# iCntr+=preds_max.shape[0]
acc1, acc5 = getNumCorrect(preds,
labels_gt,
topk=(1, min(numClasses,
5)))
val_acc.addScalar(acc1, preds.shape[0])
val_acc_5.addScalar(acc5, preds.shape[0])
vggFontModel.train()
elapsed_time = time.time() - start_time
#DO HERE
with open(
os.path.join(opt.exp_dir, opt.exp_name,
'log_train.txt'), 'a') as log:
# print('COUNT-------',val_acc_5.n_count)
# training loss and validation loss
loss_log = f'[{iteration+1}/{opt.num_iter}] \
Train loss: {loss_train.val():0.5f}, Val loss: {loss_val.val():0.5f}, \
Train Top-1 Acc: {train_acc.val()*100:0.5f}, Train Top-5 Acc: {train_acc_5.val()*100:0.5f}, \
Val Top-1 Acc: {val_acc.val()*100:0.5f}, Val Top-5 Acc: {val_acc_5.val()*100:0.5f}, \
Elapsed_time: {elapsed_time:0.5f}'
#plotting
lib.plot.plot(os.path.join(opt.plotDir, 'Train-Loss'),
loss_train.val().item())
lib.plot.plot(os.path.join(opt.plotDir, 'Val-Loss'),
loss_val.val().item())
lib.plot.plot(os.path.join(opt.plotDir, 'Train-Top-1-Acc'),
train_acc.val() * 100)
lib.plot.plot(os.path.join(opt.plotDir, 'Train-Top-5-Acc'),
train_acc_5.val() * 100)
lib.plot.plot(os.path.join(opt.plotDir, 'Val-Top-1-Acc'),
val_acc.val() * 100)
lib.plot.plot(os.path.join(opt.plotDir, 'Val-Top-5-Acc'),
val_acc_5.val() * 100)
print(loss_log)
log.write(loss_log + "\n")
loss_train.reset()
loss_val.reset()
train_acc.reset()
val_acc.reset()
train_acc_5.reset()
val_acc_5.reset()
# trainCorrect=0
# tCntr=0
lib.plot.flush()
# save model per 30000 iter.
if (iteration) % 15000 == 0:
torch.save(
{
'vggFontModel': vggFontModel_module.state_dict(),
'optimizer': optimizer.state_dict(),
'scheduler': scheduler.state_dict()
},
os.path.join(opt.exp_dir, opt.exp_name, 'iter_' +
str(iteration + 1) + '_vggFont.pth'))
lib.plot.tick()
if (iteration + 1) == opt.num_iter:
print('end the training')
sys.exit()
iteration += 1
cntr += 1
def train(opt):
os.makedirs(opt.log, exist_ok=True)
writer = SummaryWriter(opt.log)
""" 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'
)
opt.select_data = opt.select_data.split('-')
opt.batch_ratio = opt.batch_ratio.split('-')
train_dataset = Batch_Balanced_Dataset(opt)
log = open(f'./saved_models/{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 """
ctc_converter = CTCLabelConverter(opt.character)
attn_converter = AttnLabelConverter(opt.character)
opt.num_class = len(attn_converter.character)
if opt.rgb:
opt.input_channel = 3
model = Model(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)
# 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
# data parallel for multi-GPU
model = torch.nn.DataParallel(model).to(device)
model.train()
if opt.saved_model != '':
print(f'loading pretrained model from {opt.saved_model}')
if opt.FT:
model.load_state_dict(torch.load(opt.saved_model), strict=False)
else:
model.load_state_dict(torch.load(opt.saved_model))
""" setup loss """
loss_avg = Averager()
ctc_loss = torch.nn.CTCLoss(zero_infinity=True).to(device)
attn_loss = torch.nn.CrossEntropyLoss(ignore_index=0).to(device)
# 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))
# 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("Optimizer:")
""" final options """
# print(opt)
with open(f'./saved_models/{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_model != '':
try:
start_iter = int(opt.saved_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
iteration = start_iter
pbar = tqdm(range(opt.num_iter))
for iteration in pbar:
# train part
image_tensors, labels = train_dataset.get_batch()
image = image_tensors.to(device)
ctc_text, ctc_length = ctc_converter.encode(
labels, batch_max_length=opt.batch_max_length)
attn_text, attn_length = attn_converter.encode(
labels, batch_max_length=opt.batch_max_length)
batch_size = image.size(0)
preds, refiner = model(image, attn_text[:, :-1])
refiner_size = torch.IntTensor([refiner.size(1)] * batch_size)
refiner = refiner.log_softmax(2).permute(1, 0, 2)
refiner_loss = ctc_loss(refiner, ctc_text, refiner_size, ctc_length)
total_loss = opt.lambda_ctc * refiner_loss
target = attn_text[:, 1:] # without [GO] Symbol
for pred in preds:
total_loss += opt.lambda_attn * attn_loss(
pred.view(-1, pred.shape[-1]),
target.contiguous().view(-1))
model.zero_grad()
total_loss.backward()
torch.nn.utils.clip_grad_norm_(
model.parameters(),
opt.grad_clip) # gradient clipping with 5 (Default)
optimizer.step()
loss_avg.add(total_loss)
if loss_avg.val() <= 0.6:
opt.grad_clip = 2
if loss_avg.val() <= 0.3:
opt.grad_clip = 1
preds = (p.cpu() for p in preds)
refiner = refiner.cpu()
image = image.cpu()
torch.cuda.empty_cache()
writer.add_scalar('train_loss', loss_avg.val(), iteration)
pbar.set_description('Iteration {0}/{1}, AvgLoss {2}'.format(
iteration, opt.num_iter, loss_avg.val()))
# validation part
if (iteration + 1) % opt.valInterval == 0 or iteration == 0:
elapsed_time = time.time() - start_time
# for log
with open(f'./saved_models/{opt.exp_name}/log_train.txt',
'a') as log:
model.eval()
with torch.no_grad():
valid_loss, current_accuracy, current_norm_ED, preds, confidence_score, labels, infer_time, length_of_data = validation(
model, attn_loss, valid_loader, attn_converter, opt)
model.train()
# training loss and validation loss
loss_log = f'[{iteration+1}/{opt.num_iter}] Train loss: {loss_avg.val():0.5f}, Valid loss: {valid_loss:0.5f}, Elapsed_time: {elapsed_time:0.5f}'
writer.add_scalar('Val_loss', valid_loss)
pbar.set_description(loss_log)
loss_avg.reset()
current_model_log = f'{"Current_accuracy":17s}: {current_accuracy:0.3f}, {"Current_norm_ED":17s}: {current_norm_ED:0.2f}'
# keep best accuracy model (on valid dataset)
if current_accuracy > best_accuracy:
best_accuracy = current_accuracy
torch.save(
model.state_dict(),
f'./saved_models/{opt.exp_name}/best_accuracy_{str(best_accuracy)}.pth'
)
if current_norm_ED > best_norm_ED:
best_norm_ED = current_norm_ED
torch.save(
model.state_dict(),
f'./saved_models/{opt.exp_name}/best_norm_ED.pth')
best_model_log = f'{"Best_accuracy":17s}: {best_accuracy:0.3f}, {"Best_norm_ED":17s}: {best_norm_ED:0.2f}'
loss_model_log = f'{loss_log}\n{current_model_log}\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":25s} | {"Prediction":25s} | Confidence Score & T/F'
predicted_result_log = f'{dashed_line}\n{head}\n{dashed_line}\n'
for gt, pred, confidence in zip(labels[:5], preds[:5],
confidence_score[:5]):
if 'Attn' or 'Transformer' in opt.Prediction:
gt = gt[:gt.find('[s]')]
pred = pred[:pred.find('[s]')]
predicted_result_log += f'{gt:25s} | {pred:25s} | {confidence:0.4f}\t{str(pred == gt)}\n'
predicted_result_log += f'{dashed_line}'
log.write(predicted_result_log + '\n')
# save model per 1e+3 iter.
if (iteration + 1) % 1e+3 == 0:
torch.save(model.state_dict(),
f'./saved_models/{opt.exp_name}/SCATTER_STR.pth')
if (iteration + 1) == opt.num_iter:
print('end the training')
sys.exit()
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():
""" dataset preparation """
train_dataset_lmdb = LmdbDataset(cfg.lmdb_trainset_dir_name)
val_dataset_lmdb = LmdbDataset(cfg.lmdb_valset_dir_name)
train_loader = torch.utils.data.DataLoader(
train_dataset_lmdb, batch_size=cfg.batch_size,
collate_fn=data_collate,
shuffle=True,
num_workers=int(cfg.workers),
pin_memory=True)
valid_loader = torch.utils.data.DataLoader(
val_dataset_lmdb, batch_size=cfg.batch_size,
collate_fn=data_collate,
shuffle=True, # 'True' to check training progress with validation function.
num_workers=int(cfg.workers),
pin_memory=True)
# --------------------训练过程---------------------------------
model = advancedEAST()
if int(cfg.train_task_id[-3:]) != 256:
id_num = cfg.train_task_id[-3:]
idx_dic = {'384': 256, '512': 384, '640': 512, '736': 640}
model.load_state_dict(torch.load('./saved_model/3T{}_best_loss.pth'.format(idx_dic[id_num])))
elif os.path.exists('./saved_model/3T{}_best_loss.pth'.format(cfg.train_task_id)):
model.load_state_dict(torch.load('./saved_model/3T{}_best_loss.pth'.format(cfg.train_task_id)))
model = model.to(device)
optimizer = optim.Adam(model.parameters(), lr=cfg.lr, weight_decay=cfg.decay)
loss_func = quad_loss
train_Loss_list = []
val_Loss_list = []
'''start training'''
start_iter = 0
if cfg.saved_model != '':
try:
start_iter = int(cfg.saved_model.split('_')[-1].split('.')[0])
print('continue to train, start_iter: {}'.format(start_iter))
except Exception as e:
print(e)
pass
start_time = time.time()
best_mF1_score = 0
i = start_iter
step_num = 0
start_time = time.time()
loss_avg = Averager()
val_loss_avg = Averager()
eval_p_r_f = eval_pre_rec_f1()
while(True):
model.train()
# train part
# training-----------------------------
for image_tensors, labels, gt_xy_list in train_loader:
step_num += 1
batch_x = image_tensors.to(device).float()
batch_y = labels.to(device).float() # float64转float32
out = model(batch_x)
loss = loss_func(batch_y, out)
optimizer.zero_grad()
loss.backward()
optimizer.step()
loss_avg.add(loss)
train_Loss_list.append(loss_avg.val())
if i == 5 or (i + 1) % 10 == 0:
eval_p_r_f.add(out, gt_xy_list) # 非常耗时！！！
# save model per 100 epochs.
if (i + 1) % 1e+2 == 0:
torch.save(model.state_dict(), './saved_models/{}/{}_iter_{}.pth'.format(cfg.train_task_id, cfg.train_task_id, step_num+1))
print('Epoch:[{}/{}] Training Loss: {:.3f}'.format(i + 1, cfg.epoch_num, train_Loss_list[-1].item()))
loss_avg.reset()
if i == 5 or (i + 1) % 10 == 0:
mPre, mRec, mF1_score = eval_p_r_f.val()
print('Training meanPrecision:{:.2f}% meanRecall:{:.2f}% meanF1-score:{:.2f}%'.format(mPre, mRec, mF1_score))
eval_p_r_f.reset()
# evaluation--------------------------------
if (i + 1) % cfg.valInterval == 0:
elapsed_time = time.time() - start_time
print('Elapsed time:{}s'.format(round(elapsed_time)))
model.eval()
for image_tensors, labels, gt_xy_list in valid_loader:
batch_x = image_tensors.to(device)
batch_y = labels.to(device).float() # float64转float32
out = model(batch_x)
loss = loss_func(batch_y, out)
val_loss_avg.add(loss)
val_Loss_list.append(val_loss_avg.val())
eval_p_r_f.add(out, gt_xy_list)
mPre, mRec, mF1_score = eval_p_r_f.val()
print('validation meanPrecision:{:.2f}% meanRecall:{:.2f}% meanF1-score:{:.2f}%'.format(mPre, mRec, mF1_score))
eval_p_r_f.reset()
if mF1_score > best_mF1_score: # 记录最佳模型
best_mF1_score = mF1_score
torch.save(model.state_dict(), './saved_models/{}/{}_best_mF1_score_{:.3f}.pth'.format(cfg.train_task_id, cfg.train_task_id, mF1_score))
torch.save(model.state_dict(), './saved_model/{}_best_mF1_score.pth'.format(cfg.train_task_id))
print('Validation loss:{:.3f}'.format(val_loss_avg.val().item()))
val_loss_avg.reset()
if i == cfg.epoch_num:
torch.save(model.state_dict(), './saved_models/{}/{}_iter_{}.pth'.format(cfg.train_task_id, cfg.train_task_id, i+1))
print('End the training')
break
i += 1
sys.exit()
def train(self, model, dataloader, train_loader, valid_loader):
if not os.path.exists(os.path.join(self.save_path, self.args.name)):
os.makedirs(os.path.join(self.save_path, self.args.name))
else:
if not self.args.delete:
raise SyntaxError(
f'{os.path.join(self.save_path, self.args.name)} is exist.'
)
save_folder = os.path.join(self.save_path, self.args.name)
classes = model.classes
model = torch.nn.DataParallel(model).to(self.device)
model.to(self.device)
criterion = torch.nn.CrossEntropyLoss(ignore_index=0).to(self.device)
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))
# optimizer & scheduler
optimizer = optim.Adam(filtered_parameters,
lr=self.lr,
betas=(0.9, 0.999))
scheduler = optim.lr_scheduler.CosineAnnealingLR(optimizer,
eta_min=1e-5,
T_max=self.epochs)
best_acc = 0
taken_time = time()
for epoch in range(self.epochs):
t_loss_avg = Averager()
v_loss_avg = Averager()
t_calc = ScoreCalc()
v_calc = ScoreCalc()
model.train()
word_target = None
word_preds = None
with tqdm(train_loader, unit="batch") as tepoch:
for batch, batch_sampler in enumerate(tepoch):
tepoch.set_description(
f"Epoch {epoch+1} / Batch {batch+1}")
img = batch_sampler[0].to(self.device)
text = batch_sampler[1][0].to(self.device)
length = batch_sampler[1][1]
if (self.args.choose_model == "ASTER"):
preds = model(img, text[:, :-1],
max(length).cpu().numpy())
else:
preds = model(img, text[:, :-1],
max(length).cpu().numpy())
target = text[:, 1:]
t_cost = criterion(
preds.contiguous().view(-1, preds.shape[-1]),
target.contiguous().view(-1))
model.zero_grad()
t_cost.backward()
torch.nn.utils.clip_grad_norm_(
model.parameters(),
5) # gradient clipping with 5 (Default)
optimizer.step()
scheduler.step()
t_loss_avg.add(t_cost)
self.batch_size = len(text)
pred_max = torch.argmax(
F.softmax(preds, dim=2).view(self.batch_size, -1,
classes), 2)
t_calc.add(
target,
F.softmax(preds, dim=2).view(self.batch_size, -1,
classes), length)
#print(dataloader.dataset.converter.decode(target,length),dataloader.dataset.converter.decode(pred_max,length))
if batch % (300) == 0:
word_target = dataloader.dataset.converter.decode(
target, length)[0]
word_preds = dataloader.dataset.converter.decode(
pred_max, length)[0]
tepoch.set_postfix(loss=t_loss_avg.val().item(),acc=t_calc.val().item(),\
preds=word_preds,target=word_target)
del batch_sampler, pred_max, img, text, length
model.eval()
with tqdm(valid_loader, unit="batch") as vepoch:
for batch, batch_sampler in enumerate(vepoch):
vepoch.set_description(
f"Epoch {epoch+1} / Batch {batch+1}")
with torch.no_grad():
img = batch_sampler[0].to(self.device)
text = batch_sampler[1][0].to(self.device)
length = batch_sampler[1][1].to(self.device)
preds = model(img, text[:, :-1],
max(length).cpu().numpy())
target = text[:, 1:]
v_cost = criterion(
preds.contiguous().view(-1, preds.shape[-1]),
target.contiguous().view(-1))
torch.nn.utils.clip_grad_norm_(
model.parameters(),
5) # gradient clipping with 5 (Default)
v_loss_avg.add(v_cost)
batch_size = len(text)
pred_max = torch.argmax(
F.softmax(preds,
dim=2).view(batch_size, -1, classes), 2)
v_calc.add(
target,
F.softmax(preds,
dim=2).view(batch_size, -1, classes),
length)
vepoch.set_postfix(loss=v_loss_avg.val().item(),
acc=v_calc.val().item())
del batch_sampler, v_cost, pred_max, img, text, length
if not os.path.exists(os.path.join(save_folder, self.args.name)):
os.makedirs(os.path.join(save_folder, self.args.name))
#save_plt(xs,os.path.join(save_folder,name),0,epoch)
log = dict()
log['epoch'] = epoch + 1
log['t_loss'] = t_loss_avg.val().item()
log['t_acc'] = t_calc.val().item()
log['v_loss'] = v_loss_avg.val().item()
log['v_acc'] = v_calc.val().item()
log['time'] = time() - taken_time
with open(os.path.join(save_folder, f'{self.args.name}.log'),
'a') as f:
json.dump(log, f, indent=2)
best_loss = t_loss_avg.val().item()
if best_acc < v_calc.val().item():
best_acc = v_calc.val().item()
torch.save(model.state_dict(),
os.path.join(save_folder, f'{self.args.name}.pth'))
def validation(model, criterion, evaluation_loader, converter, opt):
""" validation or evaluation """
n_correct = 0
norm_ED = 0
length_of_data = 0
infer_time = 0
valid_loss_avg = Averager()
for i, (image_tensors, labels) in enumerate(evaluation_loader):
batch_size = image_tensors.size(0)
length_of_data = length_of_data + batch_size
image = image_tensors.to(device)
# For max length prediction
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)
text_for_loss, length_for_loss = converter.encode(
labels, batch_max_length=opt.batch_max_length)
start_time = time.time()
preds = model(image, text_for_pred,
is_train=False) # tensor torch.Size([1, 26, 1024])
forward_time = time.time() - start_time
preds = preds[:, :text_for_loss.shape[1] - 1, :]
target = text_for_loss[:, 1:] # without [GO] Symbol
cost = criterion(preds.contiguous().view(-1, preds.shape[-1]),
target.contiguous().view(-1))
# select max probabilty (greedy decoding) then decode index to character
_, preds_index = preds.max(2)
preds_str = converter.decode(preds_index, length_for_pred)
labels = converter.decode(text_for_loss[:, 1:], length_for_loss)
infer_time += forward_time
valid_loss_avg.add(cost)
# calculate accuracy & confidence score
preds_prob = F.softmax(preds, dim=2)
preds_max_prob, _ = preds_prob.max(dim=2)
confidence_score_list = []
for gt, pred, pred_max_prob in zip(labels, preds_str, preds_max_prob):
gt = gt[:gt.find('[s]')]
pred_EOS = pred.find('[s]')
pred = pred[:pred_EOS] # prune after "end of sentence" token ([s])
pred_max_prob = pred_max_prob[:pred_EOS]
# To evaluate 'case sensitive model' with alphanumeric and case insensitve setting.
if opt.sensitive and opt.data_filtering_off:
pred = pred.lower()
gt = gt.lower()
alphanumeric_case_insensitve = '0123456789abcdefghijklmnopqrstuvwxyz'
out_of_alphanumeric_case_insensitve = f'[^{alphanumeric_case_insensitve}]'
pred = re.sub(out_of_alphanumeric_case_insensitve, '', pred)
gt = re.sub(out_of_alphanumeric_case_insensitve, '', gt)
if pred == gt:
n_correct += 1
'''
(old version) ICDAR2017 DOST Normalized Edit Distance https://rrc.cvc.uab.es/?ch=7&com=tasks
"For each word we calculate the normalized edit distance to the length of the ground truth transcription."
if len(gt) == 0:
norm_ED += 1
else:
norm_ED += edit_distance(pred, gt) / len(gt)
'''
# ICDAR2019 Normalized Edit Distance
if len(gt) == 0 or len(pred) == 0:
norm_ED += 0
elif len(gt) > len(pred):
norm_ED += 1 - edit_distance(pred, gt) / len(gt)
else:
norm_ED += 1 - edit_distance(pred, gt) / len(pred)
# calculate confidence score (= multiply of pred_max_prob)
try:
confidence_score = pred_max_prob.cumprod(dim=0)[-1]
except:
confidence_score = 0 # for empty pred case, when prune after "end of sentence" token ([s])
confidence_score_list.append(confidence_score)
# print(pred, gt, pred==gt, confidence_score)
accuracy = n_correct / float(length_of_data) * 100
norm_ED = norm_ED / float(
length_of_data) # ICDAR2019 Normalized Edit Distance
return valid_loss_avg.val(
), accuracy, norm_ED, preds_str, confidence_score_list, labels, infer_time, length_of_data
def main(args):
device = torch.device(args.device)
ensure_path(args.save_path)
data = Data(args.dataset, args.n_batches, args.train_way, args.test_way, args.shot, args.query)
train_loader = data.train_loader
val_loader = data.valid_loader
model = Convnet(x_dim=2).to(device)
optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
lr_scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=20, gamma=0.5)
def save_model(name):
torch.save(model.state_dict(), osp.join(args.save_path, name + '.pth'))
trlog = dict(
args=vars(args),
train_loss=[],
val_loss=[],
train_acc=[],
val_acc=[],
max_acc=0.0,
)
timer = Timer()
for epoch in range(1, args.max_epoch + 1):
lr_scheduler.step()
model.train()
tl = Averager()
ta = Averager()
for i, batch in enumerate(train_loader, 1):
data, _ = [_.to(device) for _ in batch]
data = data.reshape(-1, 2, 105, 105)
p = args.shot * args.train_way
embedded = model(data)
embedded_shot, embedded_query = embedded[:p], embedded[p:]
proto = embedded_shot.reshape(args.shot, args.train_way, -1).mean(dim=0)
label = torch.arange(args.train_way).repeat(args.query).to(device)
logits = euclidean_metric(embedded_query, proto)
loss = F.cross_entropy(logits, label)
acc = count_acc(logits, label)
print('epoch {}, train {}/{}, loss={:.4f} acc={:.4f}'
.format(epoch, i, len(train_loader), loss.item(), acc))
tl.add(loss.item())
ta.add(acc)
optimizer.zero_grad()
loss.backward()
optimizer.step()
tl = tl.item()
ta = ta.item()
model.eval()
vl = Averager()
va = Averager()
for i, batch in enumerate(val_loader, 1):
data, _ = [_.cuda() for _ in batch]
data = data.reshape(-1, 2, 105, 105)
p = args.shot * args.test_way
data_shot, data_query = data[:p], data[p:]
proto = model(data_shot)
proto = proto.reshape(args.shot, args.test_way, -1).mean(dim=0)
label = torch.arange(args.test_way).repeat(args.query).to(device)
logits = euclidean_metric(model(data_query), proto)
loss = F.cross_entropy(logits, label)
acc = count_acc(logits, label)
vl.add(loss.item())
va.add(acc)
vl = vl.item()
va = va.item()
print('epoch {}, val, loss={:.4f} acc={:.4f}'.format(epoch, vl, va))
if va > trlog['max_acc']:
trlog['max_acc'] = va
save_model('max-acc')
trlog['train_loss'].append(tl)
trlog['train_acc'].append(ta)
trlog['val_loss'].append(vl)
trlog['val_acc'].append(va)
torch.save(trlog, osp.join(args.save_path, 'trlog'))
save_model('epoch-last')
if epoch % args.save_epoch == 0:
save_model('epoch-{}'.format(epoch))
print('ETA:{}/{}'.format(timer.measure(), timer.measure(epoch / args.max_epoch)))
def validation(model, criterion, evaluation_loader, converter, opt):
""" validation or evaluation """
n_correct = 0
norm_ED = 0
length_of_data = 0
infer_time = 0
valid_loss_avg = Averager()
for i, (image_tensors, labels) in enumerate(evaluation_loader):
batch_size = image_tensors.size(0)
length_of_data = length_of_data + batch_size
image = image_tensors.to(device)
# For max length prediction
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)
text_for_loss, length_for_loss = converter.encode(
labels, batch_max_length=opt.batch_max_length)
start_time = time.time()
if 'CTC' in opt.Prediction:
preds = model(image, text_for_pred)
forward_time = time.time() - start_time
# Calculate evaluation loss for CTC deocder.
preds_size = torch.IntTensor([preds.size(1)] * batch_size)
# permute 'preds' to use CTCloss format
if opt.baiduCTC:
cost = criterion(preds.permute(1, 0, 2), text_for_loss,
preds_size, length_for_loss) / batch_size
else:
cost = criterion(
preds.log_softmax(2).permute(1, 0, 2), text_for_loss,
preds_size, length_for_loss)
# Select max probabilty (greedy decoding) then decode index to character
if opt.baiduCTC:
_, preds_index = preds.max(2)
preds_index = preds_index.view(-1)
else:
_, preds_index = preds.max(2)
preds_str = converter.decode(preds_index.data, preds_size.data)
else:
preds = model(image, text_for_pred, is_train=False)
forward_time = time.time() - start_time
preds = preds[:, :text_for_loss.shape[1] - 1, :]
target = text_for_loss[:, 1:] # without [GO] Symbol
cost = criterion(preds.contiguous().view(-1, preds.shape[-1]),
target.contiguous().view(-1))
# select max probabilty (greedy decoding) then decode index to character
_, preds_index = preds.max(2)
preds_str = converter.decode(preds_index, length_for_pred)
labels = converter.decode(text_for_loss[:, 1:], length_for_loss)
infer_time += forward_time
valid_loss_avg.add(cost)
# calculate accuracy & confidence score
preds_prob = F.softmax(preds, dim=2)
preds_max_prob, _ = preds_prob.max(dim=2)
confidence_score_list = []
for gt, pred, pred_max_prob in zip(labels, preds_str, preds_max_prob):
if 'Attn' in opt.Prediction:
gt = gt[:gt.find('[s]')]
pred_EOS = pred.find('[s]')
pred = pred[:
pred_EOS] # prune after "end of sentence" token ([s])
pred_max_prob = pred_max_prob[:pred_EOS]
# To evaluate 'case sensitive model' with alphanumeric and case insensitve setting.
if opt.sensitive and opt.data_filtering_off:
pred = pred.lower()
gt = gt.lower()
alphanumeric_case_insensitve = ' 0123456789가각간갇갈감갑값갓강갖같갚갛개객걀걔거걱건걷걸검겁것겉게겨격겪견결겹경곁계고곡곤곧골곰곱곳공과관광괜괴굉교구국군굳굴굵굶굽궁권귀귓규균귤그극근글긁금급긋긍기긴길김깅깊까깍깎깐깔깜깝깡깥깨꺼꺾껌껍껏껑께껴꼬꼭꼴꼼꼽꽂꽃꽉꽤꾸꾼꿀꿈뀌끄끈끊끌끓끔끗끝끼낌나낙낚난날낡남납낫낭낮낯낱낳내냄냇냉냐냥너넉넌널넓넘넣네넥넷녀녁년념녕노녹논놀놈농높놓놔뇌뇨누눈눕뉘뉴늄느늑는늘늙능늦늬니닐님다닥닦단닫달닭닮담답닷당닿대댁댐댓더덕던덜덟덤덥덧덩덮데델도독돈돌돕돗동돼되된두둑둘둠둡둥뒤뒷드득든듣들듬듭듯등디딩딪따딱딴딸땀땅때땜떠떡떤떨떻떼또똑뚜뚫뚱뛰뜨뜩뜯뜰뜻띄라락란람랍랑랗래랜램랫략량러럭런럴럼럽럿렁렇레렉렌려력련렬렵령례로록론롬롭롯료루룩룹룻뤄류륙률륭르른름릇릎리릭린림립릿링마막만많말맑맘맙맛망맞맡맣매맥맨맵맺머먹먼멀멈멋멍멎메멘멩며면멸명몇모목몬몰몸몹못몽묘무묵묶문묻물뭄뭇뭐뭘뭣므미민믿밀밉밌및밑바박밖반받발밝밟밤밥방밭배백뱀뱃뱉버번벌범법벗베벤벨벼벽변별볍병볕보복볶본볼봄봇봉뵈뵙부북분불붉붐붓붕붙뷰브븐블비빌빔빗빚빛빠빡빨빵빼뺏뺨뻐뻔뻗뼈뼉뽑뿌뿐쁘쁨사삭산살삶삼삿상새색샌생샤서석섞선설섬섭섯성세섹센셈셋셔션소속손솔솜솟송솥쇄쇠쇼수숙순숟술숨숫숭숲쉬쉰쉽슈스슨슬슴습슷승시식신싣실싫심십싯싱싶싸싹싼쌀쌍쌓써썩썰썹쎄쏘쏟쑤쓰쓴쓸씀씌씨씩씬씹씻아악안앉않알앓암압앗앙앞애액앨야약얀얄얇양얕얗얘어억언얹얻얼엄업없엇엉엊엌엎에엔엘여역연열엷염엽엿영옆예옛오옥온올옮옳옷옹와완왕왜왠외왼요욕용우욱운울움웃웅워원월웨웬위윗유육율으윽은을음응의이익인일읽잃임입잇있잊잎자작잔잖잘잠잡잣장잦재쟁쟤저적전절젊점접젓정젖제젠젯져조족존졸좀좁종좋좌죄주죽준줄줌줍중쥐즈즉즌즐즘증지직진질짐집짓징짙짚짜짝짧째쨌쩌쩍쩐쩔쩜쪽쫓쭈쭉찌찍찢차착찬찮찰참찻창찾채책챔챙처척천철첩첫청체쳐초촉촌촛총촬최추축춘출춤춥춧충취츠측츰층치칙친칠침칫칭카칸칼캄캐캠커컨컬컴컵컷케켓켜코콘콜콤콩쾌쿄쿠퀴크큰클큼키킬타탁탄탈탑탓탕태택탤터턱턴털텅테텍텔템토톤톨톱통퇴투툴툼퉁튀튜트특튼튿틀틈티틱팀팅파팎판팔팝패팩팬퍼퍽페펜펴편펼평폐포폭폰표푸푹풀품풍퓨프플픔피픽필핏핑하학한할함합항해핵핸햄햇행향허헌험헤헬혀현혈협형혜호혹혼홀홈홉홍화확환활황회획횟횡효후훈훌훔훨휘휴흉흐흑흔흘흙흡흥흩희흰히힘.?'
out_of_alphanumeric_case_insensitve = f'[^{alphanumeric_case_insensitve}]'
pred = re.sub(out_of_alphanumeric_case_insensitve, '', pred)
gt = re.sub(out_of_alphanumeric_case_insensitve, '', gt)
if pred == gt:
n_correct += 1
'''
(old version) ICDAR2017 DOST Normalized Edit Distance https://rrc.cvc.uab.es/?ch=7&com=tasks
"For each word we calculate the normalized edit distance to the length of the ground truth transcription."
if len(gt) == 0:
norm_ED += 1
else:
norm_ED += edit_distance(pred, gt) / len(gt)
'''
# ICDAR2019 Normalized Edit Distance
if len(gt) == 0 or len(pred) == 0:
norm_ED += 0
elif len(gt) > len(pred):
norm_ED += 1 - edit_distance(pred, gt) / len(gt)
else:
norm_ED += 1 - edit_distance(pred, gt) / len(pred)
# calculate confidence score (= multiply of pred_max_prob)
try:
confidence_score = pred_max_prob.cumprod(dim=0)[-1]
except:
confidence_score = 0 # for empty pred case, when prune after "end of sentence" token ([s])
confidence_score_list.append(confidence_score)
# print(pred, gt, pred==gt, confidence_score)
accuracy = n_correct / float(length_of_data) * 100
norm_ED = norm_ED / float(
length_of_data) # ICDAR2019 Normalized Edit Distance
return valid_loss_avg.val(
), accuracy, norm_ED, preds_str, confidence_score_list, labels, infer_time, length_of_data
def train(opt, log):
"""dataset preparation"""
# train dataset. for convenience
if opt.select_data == "label":
select_data = [
"1.SVT",
"2.IIIT",
"3.IC13",
"4.IC15",
"5.COCO",
"6.RCTW17",
"7.Uber",
"8.ArT",
"9.LSVT",
"10.MLT19",
"11.ReCTS",
]
elif opt.select_data == "synth":
select_data = ["MJ", "ST"]
elif opt.select_data == "synth_SA":
select_data = ["MJ", "ST", "SA"]
opt.batch_ratio = "0.4-0.4-0.2" # same ratio with SCATTER paper.
elif opt.select_data == "mix":
select_data = [
"1.SVT",
"2.IIIT",
"3.IC13",
"4.IC15",
"5.COCO",
"6.RCTW17",
"7.Uber",
"8.ArT",
"9.LSVT",
"10.MLT19",
"11.ReCTS",
"MJ",
"ST",
]
elif opt.select_data == "mix_SA":
select_data = [
"1.SVT",
"2.IIIT",
"3.IC13",
"4.IC15",
"5.COCO",
"6.RCTW17",
"7.Uber",
"8.ArT",
"9.LSVT",
"10.MLT19",
"11.ReCTS",
"MJ",
"ST",
"SA",
]
else:
select_data = opt.select_data.split("-")
# set batch_ratio for each data.
if opt.batch_ratio:
batch_ratio = opt.batch_ratio.split("-")
else:
batch_ratio = [round(1 / len(select_data), 3)] * len(select_data)
train_loader = Batch_Balanced_Dataset(opt, opt.train_data, select_data,
batch_ratio, log)
if opt.semi != "None":
select_data_unlabel = ["U1.Book32", "U2.TextVQA", "U3.STVQA"]
batch_ratio_unlabel = [round(1 / len(select_data_unlabel), 3)
] * len(select_data_unlabel)
dataset_root_unlabel = "data_CVPR2021/training/unlabel/"
train_loader_unlabel_semi = Batch_Balanced_Dataset(
opt,
dataset_root_unlabel,
select_data_unlabel,
batch_ratio_unlabel,
log,
learn_type="semi",
)
AlignCollate_valid = AlignCollate(opt, mode="test")
valid_dataset, valid_dataset_log = hierarchical_dataset(
root=opt.valid_data, opt=opt, mode="test")
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=False,
)
log.write(valid_dataset_log)
print("-" * 80)
log.write("-" * 80 + "\n")
""" model configuration """
if "CTC" in opt.Prediction:
converter = CTCLabelConverter(opt.character)
else:
converter = AttnLabelConverter(opt.character)
opt.sos_token_index = converter.dict["[SOS]"]
opt.eos_token_index = converter.dict["[EOS]"]
opt.num_class = len(converter.character)
model = Model(opt)
# 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
# data parallel for multi-GPU
model = torch.nn.DataParallel(model).to(device)
model.train()
if opt.saved_model != "":
fine_tuning_log = f"### loading pretrained model from {opt.saved_model}\n"
if "MoCo" in opt.saved_model or "MoCo" in opt.self_pre:
pretrained_state_dict_qk = torch.load(opt.saved_model)
pretrained_state_dict = {}
for name in pretrained_state_dict_qk:
if "encoder_q" in name:
rename = name.replace("encoder_q.", "")
pretrained_state_dict[rename] = pretrained_state_dict_qk[
name]
else:
pretrained_state_dict = torch.load(opt.saved_model)
for name, param in model.named_parameters():
try:
param.data.copy_(pretrained_state_dict[name].data
) # load from pretrained model
if opt.FT == "freeze":
param.requires_grad = False # Freeze
fine_tuning_log += f"pretrained layer (freezed): {name}\n"
else:
fine_tuning_log += f"pretrained layer: {name}\n"
except:
fine_tuning_log += f"non-pretrained layer: {name}\n"
print(fine_tuning_log)
log.write(fine_tuning_log + "\n")
# print("Model:")
# print(model)
log.write(repr(model) + "\n")
""" setup loss """
if "CTC" in opt.Prediction:
criterion = torch.nn.CTCLoss(zero_infinity=True).to(device)
else:
# ignore [PAD] token
criterion = torch.nn.CrossEntropyLoss(
ignore_index=converter.dict["[PAD]"]).to(device)
if "Pseudo" in opt.semi:
criterion_SemiSL = PseudoLabelLoss(opt, converter, criterion)
elif "MeanT" in opt.semi:
criterion_SemiSL = MeanTeacherLoss(opt, student_for_init_teacher=model)
# loss averager
train_loss_avg = Averager()
semi_loss_avg = Averager() # semi supervised loss avg
# filter that only require gradient descent
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(f"Trainable params num: {sum(params_num)}")
log.write(f"Trainable params num: {sum(params_num)}\n")
# [print(name, p.numel()) for name, p in filter(lambda p: p[1].requires_grad, model.named_parameters())]
# setup optimizer
if opt.optimizer == "sgd":
optimizer = torch.optim.SGD(
filtered_parameters,
lr=opt.lr,
momentum=opt.sgd_momentum,
weight_decay=opt.sgd_weight_decay,
)
elif opt.optimizer == "adadelta":
optimizer = torch.optim.Adadelta(filtered_parameters,
lr=opt.lr,
rho=opt.rho,
eps=opt.eps)
elif opt.optimizer == "adam":
optimizer = torch.optim.Adam(filtered_parameters, lr=opt.lr)
print("Optimizer:")
print(optimizer)
log.write(repr(optimizer) + "\n")
if "super" in opt.schedule:
if opt.optimizer == "sgd":
cycle_momentum = True
else:
cycle_momentum = False
scheduler = torch.optim.lr_scheduler.OneCycleLR(
optimizer,
max_lr=opt.lr,
cycle_momentum=cycle_momentum,
div_factor=20,
final_div_factor=1000,
total_steps=opt.num_iter,
)
print("Scheduler:")
print(scheduler)
log.write(repr(scheduler) + "\n")
""" final options """
# print(opt)
opt_log = "------------ Options -------------\n"
args = vars(opt)
for k, v in args.items():
if str(k) == "character" and len(str(v)) > 500:
opt_log += f"{str(k)}: So many characters to show all: number of characters: {len(str(v))}\n"
else:
opt_log += f"{str(k)}: {str(v)}\n"
opt_log += "---------------------------------------\n"
print(opt_log)
log.write(opt_log)
log.close()
""" start training """
start_iter = 0
if opt.saved_model != "":
try:
start_iter = int(opt.saved_model.split("_")[-1].split(".")[0])
print(f"continue to train, start_iter: {start_iter}")
except:
pass
start_time = time.time()
best_score = -1
# training loop
for iteration in tqdm(
range(start_iter + 1, opt.num_iter + 1),
total=opt.num_iter,
position=0,
leave=True,
):
if "MeanT" in opt.semi:
image_tensors, image_tensors_ema, labels = train_loader.get_batch_ema(
)
else:
image_tensors, labels = train_loader.get_batch()
image = image_tensors.to(device)
labels_index, labels_length = converter.encode(
labels, batch_max_length=opt.batch_max_length)
batch_size = image.size(0)
# default recognition loss part
if "CTC" in opt.Prediction:
preds = model(image)
preds_size = torch.IntTensor([preds.size(1)] * batch_size)
preds_log_softmax = preds.log_softmax(2).permute(1, 0, 2)
loss = criterion(preds_log_softmax, labels_index, preds_size,
labels_length)
else:
preds = model(image,
labels_index[:, :-1]) # align with Attention.forward
target = labels_index[:, 1:] # without [SOS] Symbol
loss = criterion(preds.view(-1, preds.shape[-1]),
target.contiguous().view(-1))
# semi supervised part (SemiSL)
if "Pseudo" in opt.semi:
image_unlabel, _ = train_loader_unlabel_semi.get_batch_two_images()
image_unlabel = image_unlabel.to(device)
loss_SemiSL = criterion_SemiSL(image_unlabel, model)
loss = loss + loss_SemiSL
semi_loss_avg.add(loss_SemiSL)
elif "MeanT" in opt.semi:
(
image_tensors_unlabel,
image_tensors_unlabel_ema,
) = train_loader_unlabel_semi.get_batch_two_images()
image_unlabel = image_tensors_unlabel.to(device)
student_input = torch.cat([image, image_unlabel], dim=0)
image_ema = image_tensors_ema.to(device)
image_unlabel_ema = image_tensors_unlabel_ema.to(device)
teacher_input = torch.cat([image_ema, image_unlabel_ema], dim=0)
loss_SemiSL = criterion_SemiSL(
student_input=student_input,
student_logit=preds,
student=model,
teacher_input=teacher_input,
iteration=iteration,
)
loss = loss + loss_SemiSL
semi_loss_avg.add(loss_SemiSL)
model.zero_grad()
loss.backward()
torch.nn.utils.clip_grad_norm_(
model.parameters(),
opt.grad_clip) # gradient clipping with 5 (Default)
optimizer.step()
train_loss_avg.add(loss)
if "super" in opt.schedule:
scheduler.step()
else:
adjust_learning_rate(optimizer, iteration, opt)
# validation part.
# To see training progress, we also conduct validation when 'iteration == 1'
if iteration % opt.val_interval == 0 or iteration == 1:
# for validation log
with open(f"./saved_models/{opt.exp_name}/log_train.txt",
"a") as log:
model.eval()
with torch.no_grad():
(
valid_loss,
current_score,
preds,
confidence_score,
labels,
infer_time,
length_of_data,
) = validation(model, criterion, valid_loader, converter,
opt)
model.train()
# keep best score (accuracy or norm ED) model on valid dataset
# Do not use this on test datasets. It would be an unfair comparison
# (training should be done without referring test set).
if current_score > best_score:
best_score = current_score
torch.save(
model.state_dict(),
f"./saved_models/{opt.exp_name}/best_score.pth",
)
# validation log: loss, lr, score (accuracy or norm ED), time.
lr = optimizer.param_groups[0]["lr"]
elapsed_time = time.time() - start_time
valid_log = f"\n[{iteration}/{opt.num_iter}] Train_loss: {train_loss_avg.val():0.5f}, Valid_loss: {valid_loss:0.5f}"
valid_log += f", Semi_loss: {semi_loss_avg.val():0.5f}\n"
valid_log += f'{"Current_score":17s}: {current_score:0.2f}, Current_lr: {lr:0.7f}\n'
valid_log += f'{"Best_score":17s}: {best_score:0.2f}, Infer_time: {infer_time:0.1f}, Elapsed_time: {elapsed_time:0.1f}'
# show some predicted results
dashed_line = "-" * 80
head = f'{"Ground Truth":25s} | {"Prediction":25s} | Confidence Score & T/F'
predicted_result_log = f"{dashed_line}\n{head}\n{dashed_line}\n"
for gt, pred, confidence in zip(labels[:5], preds[:5],
confidence_score[:5]):
if "Attn" in opt.Prediction:
gt = gt[:gt.find("[EOS]")]
pred = pred[:pred.find("[EOS]")]
predicted_result_log += f"{gt:25s} | {pred:25s} | {confidence:0.4f}\t{str(pred == gt)}\n"
predicted_result_log += f"{dashed_line}"
valid_log = f"{valid_log}\n{predicted_result_log}"
print(valid_log)
log.write(valid_log + "\n")
opt.writer.add_scalar("train/train_loss",
float(f"{train_loss_avg.val():0.5f}"),
iteration)
opt.writer.add_scalar("train/semi_loss",
float(f"{semi_loss_avg.val():0.5f}"),
iteration)
opt.writer.add_scalar("train/lr", float(f"{lr:0.7f}"),
iteration)
opt.writer.add_scalar("train/elapsed_time",
float(f"{elapsed_time:0.1f}"), iteration)
opt.writer.add_scalar("valid/valid_loss",
float(f"{valid_loss:0.5f}"), iteration)
opt.writer.add_scalar("valid/current_score",
float(f"{current_score:0.2f}"),
iteration)
opt.writer.add_scalar("valid/best_score",
float(f"{best_score:0.2f}"), iteration)
train_loss_avg.reset()
semi_loss_avg.reset()
""" Evaluation at the end of training """
print("Start evaluation on benchmark testset")
""" keep evaluation model and result logs """
os.makedirs(f"./result/{opt.exp_name}", exist_ok=True)
os.makedirs(f"./evaluation_log", exist_ok=True)
saved_best_model = f"./saved_models/{opt.exp_name}/best_score.pth"
# os.system(f'cp {saved_best_model} ./result/{opt.exp_name}/')
model.load_state_dict(torch.load(f"{saved_best_model}"))
opt.eval_type = "benchmark"
model.eval()
with torch.no_grad():
total_accuracy, eval_data_list, accuracy_list = benchmark_all_eval(
model, criterion, converter, opt)
opt.writer.add_scalar("test/total_accuracy",
float(f"{total_accuracy:0.2f}"), iteration)
for eval_data, accuracy in zip(eval_data_list, accuracy_list):
accuracy = float(accuracy)
opt.writer.add_scalar(f"test/{eval_data}", float(f"{accuracy:0.2f}"),
iteration)
print(
f'finished the experiment: {opt.exp_name}, "CUDA_VISIBLE_DEVICES" was {opt.CUDA_VISIBLE_DEVICES}'
)
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('-')
train_dataset = Batch_Balanced_Dataset(opt)
log = open(f'./saved_models/{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 + ' ')
log.close()
""" model configuration """
# if 'CTC' in opt.Prediction:
if opt.baiduCTC:
CTC_converter = CTCLabelConverterForBaiduWarpctc(opt.character)
else:
CTC_converter = CTCLabelConverter(opt.character)
# else:
Attn_converter = AttnLabelConverter(opt.character)
opt.num_class_ctc = len(CTC_converter.character)
opt.num_class_attn = len(Attn_converter.character)
if opt.rgb:
opt.input_channel = 3
model = Model(opt)
print('model input parameters', opt.imgH, opt.imgW, opt.num_fiducial,
opt.input_channel, opt.output_channel, opt.hidden_size,
opt.num_class_ctc, opt.num_class_attn, opt.batch_max_length,
opt.Transformation, opt.FeatureExtraction, opt.SequenceModeling,
opt.Prediction)
# weight initialization
for name, param in model.named_parameters():
# print(name)
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()
print("Model:")
print(model)
# print(summary(model, (1, opt.imgH, opt.imgW,1)))
""" setup loss """
if opt.baiduCTC:
# need to install warpctc. see our guideline.
if opt.label_smooth:
criterion_major_path = SmoothCTCLoss(num_classes=opt.num_class_ctc,
weight=0.05)
else:
criterion_major_path = CTCLoss()
#criterion_major_path = CTCLoss(average_frames=False, reduction="mean", blank=0)
else:
criterion_major_path = torch.nn.CTCLoss(zero_infinity=True).to(device)
# else:
# criterion = torch.nn.CrossEntropyLoss(ignore_index=0).to(device) # ignore [GO] token = ignore index 0
# loss averager
#criterion_major_path = torch.nn.CTCLoss(zero_infinity=True).to(device)
criterion_guide_path = torch.nn.CrossEntropyLoss(ignore_index=0).to(device)
loss_avg_major_path = Averager()
loss_avg_guide_path = Averager()
# filter that only require gradient decent
guide_parameters = []
major_parameters = []
guide_model_part_names = [
"Transformation", "FeatureExtraction", "SequenceModeling_Attn",
"Attention"
]
major_model_part_names = ["SequenceModeling_CTC", "CTC"]
for name, param in model.named_parameters():
if param.requires_grad:
if name.split(".")[1] in guide_model_part_names:
guide_parameters.append(param)
elif name.split(".")[1] in major_model_part_names:
major_parameters.append(param)
# print(name)
# [print(name, p.numel()) for name, p in filter(lambda p: p[1].requires_grad, model.named_parameters())]
if opt.continue_training:
guide_parameters = []
# setup optimizer
if opt.adam:
optimizer = optim.Adam(filtered_parameters,
lr=opt.lr,
betas=(opt.beta1, 0.999))
else:
optimizer_ctc = AdamW(major_parameters, lr=opt.lr)
if not opt.continue_training:
optimizer_attn = AdamW(guide_parameters, lr=opt.lr)
scheduler_ctc = get_linear_schedule_with_warmup(
optimizer_ctc, num_warmup_steps=10000, num_training_steps=opt.num_iter)
scheduler_attn = get_linear_schedule_with_warmup(
optimizer_attn,
num_warmup_steps=10000,
num_training_steps=opt.num_iter)
start_iter = 0
if opt.saved_model != '' and (not opt.continue_training):
print(f'loading pretrained model from {opt.saved_model}')
checkpoint = torch.load(opt.saved_model)
start_iter = checkpoint['start_iter'] + 1
if not opt.adam:
optimizer_ctc.load_state_dict(
checkpoint['optimizer_ctc_state_dict'])
if not opt.continue_training:
optimizer_attn.load_state_dict(
checkpoint['optimizer_attn_state_dict'])
scheduler_ctc.load_state_dict(
checkpoint['scheduler_ctc_state_dict'])
scheduler_attn.load_state_dict(
checkpoint['scheduler_attn_state_dict'])
print(scheduler_ctc.get_lr())
print(scheduler_attn.get_lr())
if opt.FT:
model.load_state_dict(checkpoint['model_state_dict'], strict=False)
else:
model.load_state_dict(checkpoint['model_state_dict'])
if opt.continue_training:
model.load_state_dict(torch.load(opt.saved_model))
# print("Optimizer:")
# print(optimizer)
#
scheduler_ctc = get_linear_schedule_with_warmup(
optimizer_ctc,
num_warmup_steps=10000,
num_training_steps=opt.num_iter,
last_epoch=start_iter - 1)
scheduler_attn = get_linear_schedule_with_warmup(
optimizer_attn,
num_warmup_steps=10000,
num_training_steps=opt.num_iter,
last_epoch=start_iter - 1)
""" final options """
# print(opt)
with open(f'./saved_models/{opt.exp_name}/opt.txt', 'a') as opt_file:
opt_log = '------------ Options ------------- '
args = vars(opt)
for k, v in args.items():
opt_log += f'{str(k)}: {str(v)} '
opt_log += '--------------------------------------- '
print(opt_log)
opt_file.write(opt_log)
""" start training """
start_time = time.time()
best_accuracy = -1
best_norm_ED = -1
iteration = start_iter - 1
if opt.continue_training:
start_iter = 0
while (True):
# train part
image_tensors, labels = train_dataset.get_batch()
iteration += 1
if iteration < start_iter:
continue
image = image_tensors.to(device)
# print(image.size())
text_attn, length_attn = Attn_converter.encode(
labels, batch_max_length=opt.batch_max_length)
#print("1")
text_ctc, length_ctc = CTC_converter.encode(
labels, batch_max_length=opt.batch_max_length)
#print("2")
#if iteration == start_iter :
# writer.add_graph(model, (image, text_attn))
batch_size = image.size(0)
preds_major, preds_guide = model(image, text_attn[:, :-1])
#print("10")
preds_size = torch.IntTensor([preds_major.size(1)] * batch_size)
if opt.baiduCTC:
preds_major = preds_major.permute(1, 0, 2) # to use CTCLoss format
if opt.label_smooth:
cost_ctc = criterion_major_path(preds_major, text_ctc,
preds_size, length_ctc,
batch_size)
else:
cost_ctc = criterion_major_path(
preds_major, text_ctc, preds_size, length_ctc) / batch_size
else:
preds_major = preds_major.log_softmax(2).permute(1, 0, 2)
cost_ctc = criterion_major_path(preds_major, text_ctc, preds_size,
length_ctc)
#print("3")
# preds = model(image, text[:, :-1]) # align with Attention.forward
target = text_attn[:, 1:] # without [GO] Symbol
if not opt.continue_training:
cost_attn = criterion_guide_path(
preds_guide.view(-1, preds_guide.shape[-1]),
target.contiguous().view(-1))
optimizer_attn.zero_grad()
cost_attn.backward(retain_graph=True)
torch.nn.utils.clip_grad_norm_(
guide_parameters,
opt.grad_clip) # gradient clipping with 5 (Default)
optimizer_attn.step()
optimizer_ctc.zero_grad()
cost_ctc.backward()
torch.nn.utils.clip_grad_norm_(
major_parameters,
opt.grad_clip) # gradient clipping with 5 (Default)
optimizer_ctc.step()
scheduler_ctc.step()
scheduler_attn.step()
#print("4")
loss_avg_major_path.add(cost_ctc)
if not opt.continue_training:
loss_avg_guide_path.add(cost_attn)
if (iteration + 1) % 100 == 0:
writer.add_scalar("Loss/train_ctc", loss_avg_major_path.val(),
(iteration + 1) // 100)
loss_avg_major_path.reset()
if not opt.continue_training:
writer.add_scalar("Loss/train_attn", loss_avg_guide_path.val(),
(iteration + 1) // 100)
loss_avg_guide_path.reset()
# validation part
if (
iteration + 1
) % opt.valInterval == 0: #or iteration == 0: # To see training progress, we also conduct validation when 'iteration == 0'
elapsed_time = time.time() - start_time
# for log
with open(f'./saved_models/{opt.exp_name}/log_train.txt',
'a') as log:
model.eval()
with torch.no_grad():
valid_loss, current_accuracy, current_norm_ED, preds, confidence_score, labels, infer_time, length_of_data = validation(
model, criterion_major_path, valid_loader,
CTC_converter, opt)
model.train()
writer.add_scalar("Loss/valid", valid_loss,
(iteration + 1) // opt.valInterval)
writer.add_scalar("Metrics/accuracy", current_accuracy,
(iteration + 1) // opt.valInterval)
writer.add_scalar("Metrics/norm_ED", current_norm_ED,
(iteration + 1) // opt.valInterval)
# loss_log = f'[{iteration+1}/{opt.num_iter}] Train loss: {train_loss:0.5f}, Valid loss: {valid_loss:0.5f}, Elapsed_time: {elapsed_time:0.5f}'
# loss_avg.reset()
current_model_log = f'{"Current_accuracy":17s}: {current_accuracy:0.3f}, {"Current_norm_ED":17s}: {current_norm_ED:0.2f}'
# training loss and validation loss
if not opt.continue_training:
loss_log = f'[{iteration+1}/{opt.num_iter}] Train loss ctc: {loss_avg_major_path.val():0.5f}, Train loss attn: {loss_avg_guide_path.val():0.5f}, Valid loss: {valid_loss:0.5f}, Elapsed_time: {elapsed_time:0.5f}'
else:
loss_log = f'[{iteration+1}/{opt.num_iter}] Train loss ctc: {loss_avg_major_path.val():0.5f}, Valid loss: {valid_loss:0.5f}, Elapsed_time: {elapsed_time:0.5f}'
loss_avg_major_path.reset()
if not opt.continue_training:
loss_avg_guide_path.reset()
current_model_log = f'{"Current_accuracy":17s}: {current_accuracy:0.3f}, {"Current_norm_ED":17s}: {current_norm_ED:0.2f}'
# keep best accuracy model (on valid dataset)
if current_accuracy > best_accuracy:
best_accuracy = current_accuracy
torch.save(model.state_dict(),
f'{fol_ckpt}/best_accuracy.pth')
if current_norm_ED > best_norm_ED:
best_norm_ED = current_norm_ED
torch.save(model.state_dict(),
f'{fol_ckpt}/best_norm_ED.pth')
best_model_log = f'{"Best_accuracy":17s}: {best_accuracy:0.3f}, {"Best_norm_ED":17s}: {best_norm_ED:0.2f}'
loss_model_log = f'{loss_log} {current_model_log} {best_model_log}'
print(loss_model_log)
log.write(loss_model_log + ' ')
# show some predicted results
dashed_line = '-' * 80
head = f'{"Ground Truth":25s} | {"Prediction":25s} | Confidence Score & T/F'
predicted_result_log = f'{dashed_line} {head} {dashed_line} '
for gt, pred, confidence in zip(labels[:5], preds[:5],
confidence_score[:5]):
# if 'Attn' in opt.Prediction:
# gt = gt[:gt.find('[s]')]
# pred = pred[:pred.find('[s]')]
predicted_result_log += f'{gt:25s} | {pred:25s} | {confidence:0.4f} {str(pred == gt)} '
predicted_result_log += f'{dashed_line}'
print(predicted_result_log)
log.write(predicted_result_log + ' ')
# save model per 1e+5 iter.
if (iteration + 1) % 1e+3 == 0 and (not opt.continue_training):
# print(scheduler_ctc.get_lr())
# print(scheduler_attn.get_lr())
torch.save(
{
'model_state_dict': model.state_dict(),
'optimizer_attn_state_dict': optimizer_attn.state_dict(),
'optimizer_ctc_state_dict': optimizer_ctc.state_dict(),
'start_iter': iteration,
'scheduler_ctc_state_dict': scheduler_ctc.state_dict(),
'scheduler_attn_state_dict': scheduler_attn.state_dict(),
}, f'{fol_ckpt}/current_model.pth')
if (iteration + 1) == opt.num_iter:
print('end the training')
sys.exit()
def do_train_pass(train_batches_shot,
train_batches_query,
train_batches_labels,
shot,
way,
query,
expressions,
train,
test,
id_to_token=None,
id_to_tag=None,
tag_to_id=None,
test_cls=None):
model, optimizer = expressions
llog, alog = Averager(), Averager()
for i, (batch_shot, batch_query, batch_labels) in enumerate(
zip(train_batches_shot, train_batches_query, train_batches_labels),
1):
flog = Aggregate_F()
flog_old = Aggregate_F()
data_token_shot = [x for _, _, x, _, _, _, _ in batch_shot]
data_sentence_shot = [
sent[sent_id] for sent, _, _, _, _, sent_id, _ in batch_shot
]
data_sentence_labels_shot = [
label[sent_id] for _, label, _, _, _, sent_id, _ in batch_shot
]
data_sentence_bert_shot = [
bert_emb[sent_id]
for _, _, _, bert_emb, _, sent_id, _ in batch_shot
]
data_sentence_labels_shot = [[
int(batch_labels[token]) for token in sent
] for sent in data_sentence_labels_shot]
#print(data_sentence_labels_shot)
(data_sentence_shot, data_sentence_labels_shot,
data_sentence_bert_shot,
sentence_shot_lens) = pad_query_sentences(data_sentence_shot,
data_sentence_labels_shot,
data_sentence_bert_shot,
MAX_SENT_LEN,
PAD_CLS=way + 1)
#data_sentence_labels_shot = [[int(batch_labels[token]) for token in sent] for sent in data_sentence_labels_shot]
#print(data_sentence_labels_shot)
#exit()
proto = model(data_sentence_shot,
data_token_shot,
data_sentence_bert_shot,
sentence_shot_lens,
shot=True)
sorted_batch_labels = sorted(batch_labels.items(),
key=lambda kv: (kv[1], kv[0]))
###print(sorted_batch_labels)
#start with the zero!!
zero_indices = np.argwhere(data_sentence_labels_shot == 0)
###print(zero_indices)
###print(zero_indices.size())
old_proto = proto[zero_indices[0], zero_indices[1]]
###print(proto.size())
###print(old_proto.size())
new_proto = model.return_attn()(old_proto)
#print(model.return_attn())
#print(new_proto.size())
weights = F.softmax(new_proto, dim=0)
#print(weights)
#print(weights.size())
###print(weights.size())
new_proto = (weights * old_proto).sum(dim=0, keepdim=True)
#print(new_proto)
###print(new_proto.size())
# exit()
#print(sorted_batch_labels)
#Aggregate all the values of the same label and then take mean!!
for (key, val) in sorted_batch_labels:
if val != 0:
val_indices = np.argwhere(data_sentence_labels_shot == val)
#print(val)
#print(val_indices.size())
new_proto = torch.cat([
new_proto, proto[val_indices[0], val_indices[1]].mean(
dim=0, keepdim=True)
],
dim=0)
###print(new_proto.size())
###exit()
###proto = proto.reshape(shot, way-1, -1).mean(dim=0)
###dim_size = proto.size()[1]
###proto = torch.cat([torch.zeros(1, dim_size).to(device), proto])
data_token_query = [x for _, _, x, _, _, _, _ in batch_query]
data_sentence_query = [
sent[sent_id] for sent, _, _, _, _, sent_id, _ in batch_query
]
data_sentence_labels_query = [
label[sent_id] for _, label, _, _, _, sent_id, _ in batch_query
]
data_sentence_bert_query = [
bert_emb[sent_id]
for _, _, _, bert_emb, _, sent_id, _ in batch_query
]
'''zero_indices = np.argwhere(np.array(batch_labels) == 0)
nonzero_indices_part = np.argwhere(np.array(batch_labels) > 0)
nonzero_indices = []
for i in range(int(len(zero_indices)/ float(query))):
nonzero_indices += [ind[0] for ind in nonzero_indices_part]'''
#zero_indices = np.argwhere(np.array(batch_labels) == 0)
#batch_labels = [batch_labels[ind] for ind in nonzero_indices] + [batch_labels[ind[0]] for ind in zero_indices]
#data_token_query = [data_token_query[ind] for ind in nonzero_indices] + [data_token_query[ind[0]] for ind in zero_indices]
#data_sentence_query = [data_sentence_query[ind] for ind in nonzero_indices] + [data_sentence_query[ind[0]] for ind in zero_indices]
##batch_labels = [label-1 for label in batch_labels]
'''count6 = 0
count7 = 0
for sent_label in data_sentence_labels_query:
for token in sent_label:
if token == 6 :
count6+=1
if token == 7:
count7+=1
print("Count of 6\t"+str(count6)+"\tCount of 7\t"+str(count7))'''
data_sentence_labels_query = [[
int(batch_labels[token]) for token in sent
] for sent in data_sentence_labels_query]
#print(batch_labels)
#print(np.argwhere(data_sentence_labels_query == np.array(batch_labels).any()))
#labels = torch.LongTensor(np.array(batch_labels)).to(device)
(data_sentence_query, labels, data_sentence_bert_query,
sentence_query_lens) = pad_query_sentences(data_sentence_query,
data_sentence_labels_query,
data_sentence_bert_query,
MAX_SENT_LEN,
PAD_CLS=way + 1)
query_matrix = model(data_sentence_query, data_token_query,
data_sentence_bert_query, sentence_query_lens)
for vec, index, sentence in zip(query_matrix, data_token_query,
data_sentence_query):
if vec.sum() == 0.:
print(index[0])
for token in sentence:
token = token.to('cpu').item()
if token == "__PAD__":
continue
print(token)
#print(id_to_token)
print(id_to_token[int(token)])
print("Finally-------------------------")
print(id_to_token[sentence[index[0]].to('cpu').item()])
print("---")
logits = euclidean_metric(query_matrix, new_proto)
#print("Training_logits\t")
#print(logits)
#print(logits.size())
logits[:, :, 0] = model.return_0class()
###print(logits)
softmax_scores = F.softmax(logits, dim=2)
#print(softmax_scores)
#labels = torch.LongTensor(np.array(data_sentence_labels_query)).to(device)
logits_t = logits.transpose(2, 1)
#print(logits_t.size())
#print(labels.size())
#exit()
loss_function = torch.nn.CrossEntropyLoss(ignore_index=way + 1)
loss = loss_function(logits_t, labels)
#loss = F.cross_entropy(logits_t, labels)
llog.add(loss.item())
correct, total_preds, total_gold, confidence = count_F(
softmax_scores,
labels,
batch_labels.values(),
train=True,
PAD_CLS=way + 1,
id_to_tag=id_to_tag)
flog.add(correct, total_preds, total_gold)
item1, item2, item3 = flog.item()
print("Correct")
print(item1)
print("Predicted")
print(item2)
print("Gold")
print(item3)
f_score1 = flog.f_score()
print(f_score1)
print("Token-level accuracy----")
correct, total_preds, total_gold, confidence = count_F_old(
softmax_scores,
labels,
batch_labels.values(),
train=True,
PAD_CLS=way + 1)
flog_old.add(correct, total_preds, total_gold)
item1, item2, item3 = flog_old.item()
print("Correct")
print(item1)
print("Predicted")
print(item2)
print("Gold")
print(item3)
f_score1_old = flog_old.f_score()
print(f_score1_old)
#exit()
if train:
optimizer.zero_grad()
loss.backward()
optimizer.step()
return llog, flog
def train(opt):
print(opt.local_rank)
opt.device = torch.device('cuda:{}'.format(opt.local_rank))
device = opt.device
""" dataset preparation """
train_dataset = Batch_Balanced_Dataset(opt)
valid_loader = train_dataset.getValDataloader()
print('-' * 80)
""" model configuration """
if 'CTC' == opt.Prediction:
converter = CTCLabelConverter(opt.character, opt)
elif 'Attn' == opt.Prediction:
converter = AttnLabelConverter(opt.character, opt)
elif 'CTC_Attn' == opt.Prediction:
converter = CTCLabelConverter(opt.character, opt), AttnLabelConverter(
opt.character, opt)
opt.num_class = len(opt.character)
if opt.rgb:
opt.input_channel = 3
model = Model(opt)
model.to(opt.device)
print(model)
print('model input parameters', opt.rgb, 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 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
""" setup loss """
if 'CTC' == opt.Prediction:
criterion = torch.nn.CTCLoss(zero_infinity=True).to(device)
elif 'Attn' == opt.Prediction:
criterion = torch.nn.CrossEntropyLoss(
ignore_index=0).to(device), torch.nn.MSELoss(
reduction="sum").to(device)
# ignore [GO] token = ignore index 0
elif 'CTC_Attn' == opt.Prediction:
criterion = torch.nn.CTCLoss(
zero_infinity=True).to(device), torch.nn.CrossEntropyLoss(
ignore_index=0).to(device), torch.nn.MSELoss(
reduction='sum').to(device)
# loss averager
loss_avg = 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()))
if opt.local_rank == 0:
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.sgd:
optimizer = optim.SGD(filtered_parameters,
lr=opt.lr,
momentum=0.9,
weight_decay=opt.weight_decay)
elif 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)
if opt.local_rank == 0:
print("Optimizer:")
print(optimizer)
if opt.sync_bn:
model = apex.parallel.convert_syncbn_model(model)
if opt.amp > 1:
model, optimizer = amp.initialize(model,
optimizer,
opt_level="O" + str(opt.amp),
keep_batchnorm_fp32=True,
loss_scale="dynamic")
else:
model, optimizer = amp.initialize(model,
optimizer,
opt_level="O" + str(opt.amp))
# data parallel for multi-GPU
model = DDP(model)
if opt.continue_model != '':
print(f'loading pretrained model from {opt.continue_model}')
try:
model.load_state_dict(
torch.load(opt.continue_model,
map_location=torch.device(
'cuda', torch.cuda.current_device())))
except:
traceback.print_exc()
print(f'COPYING pretrained model from {opt.continue_model}')
pretrained_dict = torch.load(opt.continue_model,
map_location=torch.device(
'cuda',
torch.cuda.current_device()))
model_dict = model.state_dict()
pretrained_dict2 = dict()
for k, v in pretrained_dict.items():
if opt.Prediction == 'Attn':
if 'module.Prediction_attn.' in k:
k = k.replace('module.Prediction_attn.',
'module.Prediction.')
if k in model_dict and model_dict[k].shape == v.shape:
pretrained_dict2[k] = v
model_dict.update(pretrained_dict2)
model.load_state_dict(model_dict)
model.train()
""" final options """
with open(f'./saved_models/{opt.experiment_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'
opt_log += str(model)
print(opt_log)
opt_file.write(opt_log)
""" start training """
start_iter = 0
start_time = time.time()
best_accuracy = -1
best_norm_ED = 1e+6
i = start_iter
ct = opt.batch_mul
model.zero_grad()
dist.barrier()
while (True):
# train part
start = time.time()
image, labels, pos = train_dataset.sync_get_batch()
end = time.time()
data_t = end - start
start = time.time()
batch_size = image.size(0)
if 'CTC' == opt.Prediction:
text, length = converter.encode(
labels, batch_max_length=opt.batch_max_length)
preds = model(image, text).log_softmax(2)
preds_size = torch.IntTensor([preds.size(1)] *
batch_size).to(device)
preds = preds.permute(1, 0, 2) # to use CTCLoss format
# To avoid ctc_loss issue, disabled cudnn for the computation of the ctc_loss
# https://github.com/jpuigcerver/PyLaia/issues/16
torch.backends.cudnn.enabled = False
cost = criterion(preds, text, preds_size, length)
torch.backends.cudnn.enabled = True
elif 'Attn' == opt.Prediction:
text, length = converter.encode(
labels, batch_max_length=opt.batch_max_length)
preds = model(image, text[:, :-1]) # align with Attention.forward
preds_attn = preds[0]
preds_alpha = preds[1]
target = text[:, 1:] # without [GO] Symbol
cost = criterion[0](preds_attn.view(-1, preds_attn.shape[-1]),
target.contiguous().view(-1))
if opt.posreg_w > 0.001:
cost_pos = alpha_loss(preds_alpha, pos, opt, criterion[1])
print('attn_cost = ', cost, 'pos_cost = ',
cost_pos * opt.posreg_w)
cost += opt.posreg_w * cost_pos
else:
print('attn_cost = ', cost_attn)
elif 'CTC_Attn' == opt.Prediction:
text_ctc, length_ctc = converter[0].encode(
labels, batch_max_length=opt.batch_max_length)
text_attn, length_attn = converter[1].encode(
labels, batch_max_length=opt.batch_max_length)
""" ctc prediction and loss """
#should input text_attn here
preds = model(image, text_attn[:, :-1])
preds_ctc = preds[0].log_softmax(2)
preds_ctc_size = torch.IntTensor([preds_ctc.size(1)] *
batch_size).to(device)
preds_ctc = preds_ctc.permute(1, 0, 2) # to use CTCLoss format
# To avoid ctc_loss issue, disabled cudnn for the computation of the ctc_loss
# https://github.com/jpuigcerver/PyLaia/issues/16
torch.backends.cudnn.enabled = False
cost_ctc = criterion[0](preds_ctc, text_ctc, preds_ctc_size,
length_ctc)
torch.backends.cudnn.enabled = True
""" attention prediction and loss """
preds_attn = preds[1][0] # align with Attention.forward
preds_alpha = preds[1][1]
target = text_attn[:, 1:] # without [GO] Symbol
cost_attn = criterion[1](preds_attn.view(-1, preds_attn.shape[-1]),
target.contiguous().view(-1))
cost = opt.ctc_attn_loss_ratio * cost_ctc + (
1 - opt.ctc_attn_loss_ratio) * cost_attn
if opt.posreg_w > 0.001:
cost_pos = alpha_loss(preds_alpha, pos, opt, criterion[2])
cost += opt.posreg_w * cost_pos
cost_ctc = reduce_tensor(cost_ctc)
cost_attn = reduce_tensor(cost_attn)
cost_pos = reduce_tensor(cost_pos)
if opt.local_rank == 0:
print('ctc_cost = ', cost_ctc, 'attn_cost = ', cost_attn,
'pos_cost = ', cost_pos * opt.posreg_w)
else:
cost_ctc = reduce_tensor(cost_ctc)
cost_attn = reduce_tensor(cost_attn)
if opt.local_rank == 0:
print('ctc_cost = ', cost_ctc, 'attn_cost = ', cost_attn)
cost /= opt.batch_mul
if opt.amp > 0:
with amp.scale_loss(cost, optimizer) as scaled_loss:
scaled_loss.backward()
else:
cost.backward()
""" https://github.com/davidlmorton/learning-rate-schedules/blob/master/increasing_batch_size_without_increasing_memory.ipynb """
ct -= 1
if ct == 0:
if opt.amp > 0:
torch.nn.utils.clip_grad_norm_(amp.master_params(optimizer),
opt.grad_clip)
else:
torch.nn.utils.clip_grad_norm_(
model.parameters(),
opt.grad_clip) # gradient clipping with 5 (Default)
optimizer.step()
model.zero_grad()
ct = opt.batch_mul
else:
continue
train_t = time.time() - start
cost = reduce_tensor(cost)
loss_avg.add(cost)
if opt.local_rank == 0:
print('iter', i, 'loss =', cost, ', data_t=', data_t, ',train_t=',
train_t, ', batchsz=', opt.batch_mul * opt.batch_size)
sys.stdout.flush()
# validation part
if (i > 0 and i % opt.valInterval == 0) or (i == 0 and
opt.continue_model != ''):
elapsed_time = time.time() - start_time
print(
f'[{i}/{opt.num_iter}] Loss: {loss_avg.val():0.5f} elapsed_time: {elapsed_time:0.5f}'
)
# for log
with open(f'./saved_models/{opt.experiment_name}/log_train.txt',
'a') as log:
log.write(
f'[{i}/{opt.num_iter}] Loss: {loss_avg.val():0.5f} elapsed_time: {elapsed_time:0.5f}\n'
)
loss_avg.reset()
model.eval()
with torch.no_grad():
if 'CTC_Attn' in opt.Prediction:
# we only count for attention accuracy, because ctc is used to help attention
valid_loss, current_accuracy_ctc, current_accuracy, current_norm_ED_ctc, current_norm_ED, preds, labels, infer_time, length_of_data = validation(
model, criterion[1], valid_loader, converter[1],
opt, converter[0])
elif 'Attn' in opt.Prediction:
valid_loss, current_accuracy, current_norm_ED, preds, labels, infer_time, length_of_data = validation(
model, criterion[0], valid_loader, converter, opt)
else:
valid_loss, current_accuracy, current_norm_ED, preds, labels, infer_time, length_of_data = validation(
model, criterion, valid_loader, converter, opt)
model.train()
for pred, gt in zip(preds[:10], labels[:10]):
if 'Attn' in opt.Prediction:
pred = pred[:pred.find('[s]')]
gt = gt[:gt.find('[s]')]
print(f'{pred:20s}, gt: {gt:20s}, {str(pred == gt)}')
log.write(
f'{pred:20s}, gt: {gt:20s}, {str(pred == gt)}\n')
valid_log = f'[{i}/{opt.num_iter}] valid loss: {valid_loss:0.5f}'
valid_log += f' accuracy: {current_accuracy:0.3f}, norm_ED: {current_norm_ED:0.2f}'
if 'CTC_Attn' in opt.Prediction:
valid_log += f' ctc_accuracy: {current_accuracy_ctc:0.3f}, ctc_norm_ED: {current_norm_ED_ctc:0.2f}'
current_accuracy = max(current_accuracy,
current_accuracy_ctc)
current_norm_ED = min(current_norm_ED, current_norm_ED_ctc)
if opt.local_rank == 0:
print(valid_log)
log.write(valid_log + '\n')
# keep best accuracy model
if current_accuracy > best_accuracy:
best_accuracy = current_accuracy
torch.save(
model.state_dict(),
f'./saved_models/{opt.experiment_name}/best_accuracy.pth'
)
torch.save(
model,
f'./saved_models/{opt.experiment_name}/best_accuracy.model'
)
if current_norm_ED < best_norm_ED:
best_norm_ED = current_norm_ED
torch.save(
model.state_dict(),
f'./saved_models/{opt.experiment_name}/best_norm_ED.pth'
)
torch.save(
model,
f'./saved_models/{opt.experiment_name}/best_norm_ED.model'
)
best_model_log = f'best_accuracy: {best_accuracy:0.3f}, best_norm_ED: {best_norm_ED:0.2f}'
print(best_model_log)
log.write(best_model_log + '\n')
# save model per iter.
if (i + 1) % opt.save_interval == 0 and opt.local_rank == 0:
torch.save(model.state_dict(),
f'./saved_models/{opt.experiment_name}/iter_{i+1}.pth')
if i == opt.num_iter:
print('end the training')
sys.exit()
if opt.prof_iter > 0 and i > opt.prof_iter:
sys.exit()
i += 1
def do_pass(batches,
counters,
shot,
way,
query,
expressions,
train,
test,
id_to_token=None,
id_to_tag=None,
test_cls=None):
model, optimizer = expressions
llog, alog = Averager(), Averager()
if test:
output_file = open("./output.txt" + str(test_cls), 'w')
for i, (batch, counter) in enumerate(zip(batches, counters), 1):
#print("Batch number\t"+str(i))
data_token = [x for _, x, _, _ in batch]
data_sentence = [sent for sent, _, _, _ in batch]
data_label = [label for _, _, label, _ in batch]
p = shot * way
#print(len(data_token))
#print(p)
#print(shot)
#print(way)
data_token_shot, data_token_query = data_token[:p], data_token[p:]
data_sentence_shot, data_sentence_query = data_sentence[:
p], data_sentence[
p:]
counter_token, counter_query = counter[:p], counter[p:]
(data_sentence_shot,
sentence_shot_lens), (data_sentence_query,
query_shot_lens) = pad_sentences(
data_sentence_shot,
MAX_SENT_LEN), pad_sentences(
data_sentence_query, MAX_SENT_LEN)
proto = model(data_sentence_shot, data_token_shot, sentence_shot_lens)
proto = proto.reshape(shot, way, -1).mean(dim=0)
####label = torch.arange(way).repeat(query)
if not train:
#print(len(data_token))
#print(p)
#print(way)
query = int((len(data_token) - p) / way)
#print(query)
#exit()
label = torch.arange(way).repeat(query)
label = label.type(torch.LongTensor).to(device)
logits = euclidean_metric(
model(data_sentence_query, data_token_query, query_shot_lens),
proto)
#print(list(model.parameters()))
#print(model.return_0class())
#print(logits.size())
logits[:, 0] = model.return_0class()
#print(logits.size())
#print(label.size())
#print(len(counter_query))
#print(counter_query)
#print("---")
loss = F.cross_entropy(logits, label)
acc = count_acc(logits, label, counter_query)
llog.add(loss.item())
alog.add(acc)
if train:
optimizer.zero_grad()
loss.backward()
optimizer.step()
if test:
#print the outputs to a file
save_dev_output(output_file, logits, label, data_label,
data_sentence_query, data_token_query,
query_shot_lens, id_to_token, id_to_tag)
if test:
output_file.close()
return llog, alog
def validation_by_length(model, criterion, evaluation_loader, converter, opt):
""" validation or evaluation """
n_correct = defaultdict(float)
n_norm_ED = defaultdict(float)
length_of_data = defaultdict(int)
infer_time = 0
valid_loss_avg = Averager()
for i, (image_tensors, labels) in enumerate(evaluation_loader):
batch_size = image_tensors.size(0)
# length_of_data = length_of_data + batch_size
image = image_tensors.to(device)
# For max length prediction
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)
text_for_loss, length_for_loss = converter.encode(
labels, batch_max_length=opt.batch_max_length)
start_time = time.time()
if 'CTC' in opt.Prediction:
preds = model(image, text_for_pred)
forward_time = time.time() - start_time
# Calculate evaluation loss for CTC deocder.
preds_size = torch.IntTensor([preds.size(1)] * batch_size)
# permute 'preds' to use CTCloss format
cost = criterion(
preds.log_softmax(2).permute(1, 0, 2), text_for_loss,
preds_size, length_for_loss)
# Select max probabilty (greedy decoding) then decode index to character
_, preds_index = preds.max(2)
preds_str = converter.decode(preds_index.data, preds_size.data)
else:
preds, alphas = model(image, text_for_pred, is_train=False)
forward_time = time.time() - start_time
preds = preds[:, :text_for_loss.shape[1] - 1, :]
target = text_for_loss[:, 1:] # without [GO] Symbol
cost = criterion(preds.contiguous().view(-1, preds.shape[-1]),
target.contiguous().view(-1))
# select max probabilty (greedy decoding) then decode index to character
_, preds_index = preds.max(2)
preds_str = converter.decode(preds_index, length_for_pred)
labels = converter.decode(text_for_loss[:, 1:], length_for_loss)
infer_time += forward_time
valid_loss_avg.add(cost)
# calculate accuracy & confidence score
preds_prob = F.softmax(preds, dim=2)
preds_max_prob, _ = preds_prob.max(dim=2)
confidence_score_list = []
for gt, pred, pred_max_prob in zip(labels, preds_str, preds_max_prob):
if 'Attn' in opt.Prediction:
gt = gt[:gt.find('[s]')]
pred_EOS = pred.find('[s]')
pred = pred[:
pred_EOS] # prune after "end of sentence" token ([s])
pred_max_prob = pred_max_prob[:pred_EOS]
# To evaluate 'case sensitive model' with alphanumeric and case insensitve setting.
if opt.sensitive and opt.data_filtering_off:
pred = pred.lower()
gt = gt.lower()
alphanumeric_case_insensitve = '0123456789abcdefghijklmnopqrstuvwxyz'
out_of_alphanumeric_case_insensitve = f'[^{alphanumeric_case_insensitve}]'
pred = re.sub(out_of_alphanumeric_case_insensitve, '', pred)
gt = re.sub(out_of_alphanumeric_case_insensitve, '', gt)
length_of_data[len(gt)] += 1
if pred == gt:
n_correct[len(gt)] += 1
'''
(old version) ICDAR2017 DOST Normalized Edit Distance https://rrc.cvc.uab.es/?ch=7&com=tasks
"For each word we calculate the normalized edit distance to the length of the ground truth transcription."
if len(gt) == 0:
norm_ED += 1
else:
norm_ED += edit_distance(pred, gt) / len(gt)
'''
# ICDAR2019 Normalized Edit Distance
if len(gt) == 0 or len(pred) == 0:
n_norm_ED[len(gt)] += 0
elif len(gt) > len(pred):
n_norm_ED[len(gt)] += 1 - edit_distance(pred, gt) / len(gt)
else:
n_norm_ED[len(gt)] += 1 - edit_distance(pred, gt) / len(pred)
# calculate confidence score (= multiply of pred_max_prob)
try:
confidence_score = pred_max_prob.cumprod(dim=0)[-1]
except:
confidence_score = 0 # for empty pred case, when prune after "end of sentence" token ([s])
# log = open(f'./result/{opt.exp_name}/log_all_evaluation.txt', 'a')
# log.write(f'pred: {pred}, gt: {gt}, {pred == gt}, prob: {confidence_score:0.4f}\n')
# log.close()
confidence_score_list.append(confidence_score)
# print(pred, gt, pred == gt, confidence_score.item())
accuracy = defaultdict(float)
norm_ED = defaultdict(float)
for k in n_correct.keys():
accuracy[k] = n_correct[k] / float(length_of_data[k]) * 100
norm_ED[k] = n_norm_ED[k] / float(length_of_data[k]) * 100
return valid_loss_avg.val(
), accuracy, norm_ED, preds_str, confidence_score_list, labels, infer_time, length_of_data
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 test(self, model, target_path, dataloader):
save_folder = os.path.join(self.save_path, self.args.name)
if not os.path.exists(save_folder):
raise FileNotFoundError(f'No such folders {save_folder}')
classes = model.classes
model = torch.nn.DataParallel(model).to(self.device)
model.load_state_dict(
torch.load(os.path.join(save_folder, self.args.name + '.pth'),
map_location=self.device))
loss_avg = Averager()
calc = ScoreCalc()
cer_avg = Averager()
criterion = torch.nn.CrossEntropyLoss(ignore_index=0).to(self.device)
model.eval()
pred_num = 10
pred_result = []
with tqdm(dataloader, unit="batch") as vepoch:
for batch, batch_sampler in enumerate(vepoch):
vepoch.set_description(f"Test Session / Batch {batch+1}")
with torch.no_grad():
img = batch_sampler[0].to(self.device)
text = batch_sampler[1][0].to(self.device)
length = batch_sampler[1][1].to(self.device)
preds = model(img, text[:, :-1], max(length).cpu().numpy())
target = text[:, 1:]
v_cost = criterion(
preds.contiguous().view(-1, preds.shape[-1]),
target.contiguous().view(-1))
torch.nn.utils.clip_grad_norm_(
model.parameters(),
5) # gradient clipping with 5 (Default)
loss_avg.add(v_cost)
batch_size = len(text)
pred_max = torch.argmax(
F.softmax(preds, dim=2).view(batch_size, -1, classes),
2)
calc.add(
target,
F.softmax(preds, dim=2).view(batch_size, -1, classes),
length)
word_target = dataloader.dataset.converter.decode(
target, length)
word_preds = dataloader.dataset.converter.decode(
pred_max, length)
cer_avg.add(
torch.from_numpy(
np.array(get_cer(word_preds, word_target))))
vepoch.set_postfix(loss=loss_avg.val().item(),
acc=calc.val().item(),
cer=cer_avg.val().item())
if batch % (len(vepoch) // 10) == 0:
pred = unicodedata.normalize('NFC', word_preds[0])
target = unicodedata.normalize('NFC', word_target[0])
pred_result.append(dict(target=target, pred=pred))
del batch_sampler, v_cost, pred_max, img, text, length
#save_plt(xs,os.path.join(save_folder,name),0,epoch)
log = dict()
log['loss'] = loss_avg.val().item()
log['acc'] = calc.val().item()
log['cer'] = cer_avg.val().item()
log['preds'] = pred_result
with open(os.path.join(save_folder, f'{self.args.name}_test.log'),
'w') as f:
json.dump(log, f, indent=2)
def train(self, opt):
# src, tar dataloaders
src_dataset, tar_dataset, valid_loader = self.dataloader(opt)
src_dataset_size = src_dataset.total_data_size
tar_dataset_size = tar_dataset.total_data_size
train_size = max([src_dataset_size, tar_dataset_size])
iters_per_epoch = int(train_size / opt.batch_size)
# Modify train size. Make sure both are of same size.
# Modify training loop to continue giving src loss after tar is done.
self.model.train()
self.global_discriminator.train()
self.local_discriminator.train()
start_iter = 0
if opt.continue_model != '':
self.load(opt.continue_model)
print(" [*] Load SUCCESS")
# loss averager
cls_loss_avg = Averager()
sim_loss_avg = Averager()
loss_avg = Averager()
# training loop
print('training start !')
start_time = time.time()
best_accuracy = -1
best_norm_ED = 1e+6
# i = start_iter
gamma = 0
omega = 1
epoch = 0
for step in range(start_iter, opt.num_iter + 1):
epoch = step // iters_per_epoch
if opt.decay_flag and step > (opt.num_iter // 2):
self.d_image_opt.param_groups[0]['lr'] -= (opt.lr /
(opt.num_iter // 2))
self.d_inst_opt.param_groups[0]['lr'] -= (opt.lr /
(opt.num_iter // 2))
src_image, src_labels = src_dataset.get_batch()
src_image = src_image.to(device)
src_text, src_length = self.converter.encode(
src_labels, batch_max_length=opt.batch_max_length)
tar_image, tar_labels = tar_dataset.get_batch()
tar_image = tar_image.to(device)
tar_text, tar_length = self.converter.encode(
tar_labels, batch_max_length=opt.batch_max_length)
# Set gradient to zero...
self.model.zero_grad()
# Domain classifiers
self.global_discriminator.zero_grad()
self.local_discriminator.zero_grad()
# Attention # align with Attention.forward
src_preds, src_global_feature, src_local_feature = self.model(
src_image, src_text[:, :-1])
# src_global_feature = self.model.visual_feature
# src_local_feature = self.model.Prediction.context_history
target = src_text[:, 1:] # without [GO] Symbol
src_cls_loss = self.criterion(
src_preds.view(-1, src_preds.shape[-1]),
target.contiguous().view(-1))
src_global_feature = src_global_feature.view(
src_global_feature.shape[0], -1)
src_local_feature = src_local_feature.view(
-1, src_local_feature.shape[-1])
tar_preds, tar_global_feature, tar_local_feature = self.model(
tar_image, tar_text[:, :-1], is_train=False)
# tar_global_feature = self.model.visual_feature
# tar_local_feature = self.model.Prediction.context_history
tar_global_feature = tar_global_feature.view(
tar_global_feature.shape[0], -1)
tar_local_feature = tar_local_feature.view(
-1, tar_local_feature.shape[-1])
src_local_feature, tar_local_feature = filter_local_features(
opt, src_local_feature, src_preds, tar_local_feature,
tar_preds)
# Add domain adaption elements
# setup hyperparameter
if step % 2000 == 0:
p = float(step + start_iter) / opt.num_iter
gamma = 2. / (1. + np.exp(-10 * p)) - 1
omega = 1 - 1. / (1. + np.exp(-10 * p))
self.global_discriminator.module.set_beta(gamma)
self.local_discriminator.module.set_beta(gamma)
src_d_img_score = self.global_discriminator(src_global_feature)
src_d_inst_score = self.local_discriminator(src_local_feature)
tar_d_img_score = self.global_discriminator(tar_global_feature)
tar_d_inst_score = self.local_discriminator(tar_local_feature)
src_d_img_loss = self.D_criterion(
src_d_img_score,
torch.zeros_like(src_d_img_score).to(device))
src_d_inst_loss = self.D_criterion(
src_d_inst_score,
torch.zeros_like(src_d_inst_score).to(device))
tar_d_img_loss = self.D_criterion(
tar_d_img_score,
torch.ones_like(tar_d_img_score).to(device))
tar_d_inst_loss = self.D_criterion(
tar_d_inst_score,
torch.ones_like(tar_d_inst_score).to(device))
d_img_loss = src_d_img_loss + tar_d_img_loss
d_inst_loss = src_d_inst_loss + tar_d_inst_loss
# Add domain loss
loss = src_cls_loss.mean() + omega * (d_img_loss.mean() +
d_inst_loss.mean())
loss_avg.add(loss)
cls_loss_avg.add(src_cls_loss)
sim_loss_avg.add(d_img_loss + d_inst_loss)
# frcnn backward
loss.backward()
# clip_gradient(self.model, 10.)
torch.nn.utils.clip_grad_norm_(
self.model.parameters(),
opt.grad_clip) # gradient clipping with 5 (Default)
# frcnn optimizer update
self.optimizer.step()
# domain optimizer update
self.d_inst_opt.step()
self.d_image_opt.step()
# validation part
if step % opt.valInterval == 0:
elapsed_time = time.time() - start_time
print(
f'[{step}/{opt.num_iter}] Loss: {loss_avg.val():0.5f} CLS_Loss: {cls_loss_avg.val():0.5f} SIMI_Loss: {sim_loss_avg.val():0.5f} elapsed_time: {elapsed_time:0.5f}'
)
# for log
with open(
f'./saved_models/{opt.experiment_name}/log_train.txt',
'a') as log:
log.write(
f'[{step}/{opt.num_iter}] Loss: {loss_avg.val():0.5f} elapsed_time: {elapsed_time:0.5f}\n'
)
loss_avg.reset()
cls_loss_avg.reset()
sim_loss_avg.reset()
self.model.eval()
with torch.no_grad():
valid_loss, current_accuracy, current_norm_ED, preds, labels, infer_time, length_of_data = validation(
self.model, self.criterion, valid_loader,
self.converter, opt)
self.print_prediction_result(preds, labels, log)
valid_log = f'[{step}/{opt.num_iter}] valid loss: {valid_loss:0.5f}'
valid_log += f' accuracy: {current_accuracy:0.3f}, norm_ED: {current_norm_ED:0.2f}'
print(valid_log)
log.write(valid_log + '\n')
self.model.train()
self.global_discriminator.train()
self.local_discriminator.train()
# keep best accuracy model
if current_accuracy > best_accuracy:
best_accuracy = current_accuracy
save_name = f'./saved_models/{opt.experiment_name}/best_accuracy.pth'
self.save(opt, save_name)
if current_norm_ED < best_norm_ED:
best_norm_ED = current_norm_ED
save_name = f'./saved_models/{opt.experiment_name}/best_norm_ED.pth'
self.save(opt, save_name)
best_model_log = f'best_accuracy: {best_accuracy:0.3f}, best_norm_ED: {best_norm_ED:0.2f}'
print(best_model_log)
log.write(best_model_log + '\n')
# save model per 1e+5 iter.
if (step + 1) % 1e+5 == 0:
save_name = f'./saved_models/{opt.experiment_name}/iter_{step+1}.pth'
self.save(opt, save_name)
def train(opt):
""" dataset preparation """
opt.select_data = opt.select_data.split('-')
opt.batch_ratio = opt.batch_ratio.split('-')
train_dataset = Batch_Balanced_Dataset(opt)
AlignCollate_valid = AlignCollate(imgH=opt.imgH,
imgW=opt.imgW,
keep_ratio_with_pad=opt.PAD)
valid_dataset = 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)
print('-' * 80)
""" 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 = Model(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 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
# data parallel for multi-GPU
model = torch.nn.DataParallel(model).cuda()
model.train()
if opt.continue_model != '':
print(f'loading pretrained model from {opt.continue_model}')
model.load_state_dict(torch.load(opt.continue_model))
print("Model:")
print(model)
""" setup loss """
if 'CTC' in opt.Prediction:
criterion = torch.nn.CTCLoss(zero_infinity=True).cuda()
else:
criterion = torch.nn.CrossEntropyLoss(
ignore_index=0).cuda() # ignore [GO] token = ignore index 0
# loss averager
loss_avg = 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("Optimizer:")
print(optimizer)
""" final options """
# print(opt)
with open(f'./saved_models/{opt.experiment_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.continue_model != '':
#start_iter = int(opt.continue_model.split('_')[-1].split('.')[0])
start_iter = 10000
print(f'continue to train, start_iter: {start_iter}')
start_time = time.time()
best_accuracy = -1
best_norm_ED = 1e+6
i = start_iter
while (True):
# train part
for p in model.parameters():
p.requires_grad = True
image_tensors, labels = train_dataset.get_batch()
image = image_tensors.cuda()
text, length = converter.encode(labels)
batch_size = image.size(0)
if 'CTC' in opt.Prediction:
preds = model(image, text).log_softmax(2)
preds_size = torch.IntTensor([preds.size(1)] * batch_size)
preds = preds.permute(1, 0, 2) # to use CTCLoss format
cost = criterion(preds, text, preds_size, length)
else:
preds = model(image, text)
target = text[:, 1:] # without [GO] Symbol
cost = criterion(preds.view(-1, preds.shape[-1]),
target.contiguous().view(-1))
model.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)
# validation part
if i % opt.valInterval == 0:
elapsed_time = time.time() - start_time
print(
f'[{i}/{opt.num_iter}] Loss: {loss_avg.val():0.5f} elapsed_time: {elapsed_time:0.5f}'
)
# for log
with open(f'./saved_models/{opt.experiment_name}/log_train.txt',
'a') as log:
log.write(
f'[{i}/{opt.num_iter}] Loss: {loss_avg.val():0.5f} elapsed_time: {elapsed_time:0.5f}\n'
)
loss_avg.reset()
model.eval()
valid_loss, current_accuracy, current_norm_ED, preds, labels, infer_time, length_of_data = validation(
model, criterion, valid_loader, converter, opt)
model.train()
for pred, gt in zip(preds[:5], labels[:5]):
if 'Attn' in opt.Prediction:
pred = pred[:pred.find('[s]')]
gt = gt[:gt.find('[s]')]
print(f'{pred:20s}, gt: {gt:20s}, {str(pred == gt)}')
log.write(
f'{pred:20s}, gt: {gt:20s}, {str(pred == gt)}\n')
valid_log = f'[{i}/{opt.num_iter}] valid loss: {valid_loss:0.5f}'
valid_log += f' accuracy: {current_accuracy:0.3f}, norm_ED: {current_norm_ED:0.2f}'
print(valid_log)
log.write(valid_log + '\n')
# keep best accuracy model
if current_accuracy > best_accuracy:
best_accuracy = current_accuracy
torch.save(
model.state_dict(),
f'./saved_models/{opt.experiment_name}/best_accuracy.pth'
)
if current_norm_ED < best_norm_ED:
best_norm_ED = current_norm_ED
torch.save(
model.state_dict(),
f'./saved_models/{opt.experiment_name}/best_norm_ED.pth'
)
best_model_log = f'best_accuracy: {best_accuracy:0.3f}, best_norm_ED: {best_norm_ED:0.2f}'
print(best_model_log)
log.write(best_model_log + '\n')
# save model per 1e+5 iter.
if (i + 1) % 2000 == 0:
torch.save(model.state_dict(),
f'./saved_models/{opt.experiment_name}/iter_{i+1}.pth')
if i == opt.num_iter:
print('end the training')
sys.exit()
i += 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(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 validation(model, criterion, evaluation_loader, converter, opt):
""" validation or evaluation """
n_correct = 0
norm_ED = 0
length_of_data = 0
infer_time = 0
valid_loss_avg = Averager()
for i, (image_tensors, labels) in enumerate(evaluation_loader):
batch_size = image_tensors.size(0)
length_of_data = length_of_data + batch_size
image = image_tensors.to(device)
# For max length prediction
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)
text_for_loss, length_for_loss = converter.encode(
labels, batch_max_length=opt.batch_max_length)
start_time = time.time()
if 'CTC' in opt.Prediction:
preds = model(image, text_for_pred).log_softmax(2)
forward_time = time.time() - start_time
# Calculate evaluation loss for CTC deocder.
preds_size = torch.IntTensor([preds.size(1)] * batch_size)
# permute 'preds' to use CTCloss format
cost = criterion(preds.permute(1, 0, 2), text_for_loss, preds_size,
length_for_loss)
# Select max probabilty (greedy decoding) then decode index to character
_, preds_index = preds.max(2)
preds_index = preds_index.view(-1)
preds_str = converter.decode(preds_index.data, preds_size.data)
else:
preds = model(image, text_for_pred, is_train=False)
forward_time = time.time() - start_time
preds = preds[:, :text_for_loss.shape[1] - 1, :]
target = text_for_loss[:, 1:] # without [GO] Symbol
cost = criterion(preds.contiguous().view(-1, preds.shape[-1]),
target.contiguous().view(-1))
# select max probabilty (greedy decoding) then decode index to character
_, preds_index = preds.max(2)
preds_str = converter.decode(preds_index, length_for_pred)
labels = converter.decode(text_for_loss[:, 1:], length_for_loss)
infer_time += forward_time
valid_loss_avg.add(cost)
# calculate accuracy & confidence score
preds_prob = F.softmax(preds, dim=2)
preds_max_prob, _ = preds_prob.max(dim=2)
confidence_score_list = []
for gt, pred, pred_max_prob in zip(labels, preds_str, preds_max_prob):
if 'Attn' in opt.Prediction:
gt = gt[:gt.find('[s]')]
pred_EOS = pred.find('[s]')
pred = pred[:
pred_EOS] # prune after "end of sentence" token ([s])
pred_max_prob = pred_max_prob[:pred_EOS]
if pred == gt:
n_correct += 1
if len(gt) == 0:
norm_ED += 1
else:
norm_ED += edit_distance(pred, gt) / len(gt)
# calculate confidence score (= multiply of pred_max_prob)
try:
confidence_score = pred_max_prob.cumprod(dim=0)[-1]
except:
confidence_score = 0 # for empty pred case, when prune after "end of sentence" token ([s])
confidence_score_list.append(confidence_score)
# print(pred, gt, pred==gt, confidence_score)
accuracy = n_correct / float(length_of_data) * 100
return valid_loss_avg.val(
), accuracy, norm_ED, preds_str, confidence_score_list, labels, infer_time, length_of_data
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('-')
train_dataset = Batch_Balanced_Dataset(opt)
log = open(f'./saved_models/{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:
if opt.baiduCTC:
converter = CTCLabelConverterForBaiduWarpctc(opt.character)
else:
converter = CTCLabelConverter(opt.character)
else:
converter = AttnLabelConverter(opt.character)
opt.num_class = len(converter.character)
if opt.rgb:
opt.input_channel = 3
model = Model(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 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
# data parallel for multi-GPU
model = torch.nn.DataParallel(model).to(device)
model.train()
if opt.saved_model != '':
print(f'loading pretrained model from {opt.saved_model}')
if opt.FT:
model.load_state_dict(torch.load(opt.saved_model), strict=False)
else:
model.load_state_dict(torch.load(opt.saved_model))
print("Model:")
print(model)
""" setup loss """
if 'CTC' in opt.Prediction:
if opt.baiduCTC:
# need to install warpctc. see our guideline.
from warpctc_pytorch import CTCLoss
criterion = CTCLoss()
else:
criterion = torch.nn.CTCLoss(zero_infinity=True).to(device)
else:
criterion = torch.nn.CrossEntropyLoss(ignore_index=0).to(
device) # ignore [GO] token = ignore index 0
# loss averager
loss_avg = 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("Optimizer:")
print(optimizer)
""" final options """
# print(opt)
with open(f'./saved_models/{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_model != '':
try:
start_iter = int(opt.saved_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
iteration = start_iter
while (True):
# train part
image_tensors, labels = train_dataset.get_batch()
image = image_tensors.to(device)
text, length = converter.encode(labels,
batch_max_length=opt.batch_max_length)
batch_size = image.size(0)
if 'CTC' in opt.Prediction:
preds = model(image, text)
preds_size = torch.IntTensor([preds.size(1)] * batch_size)
if opt.baiduCTC:
preds = preds.permute(1, 0, 2) # to use CTCLoss format
cost = criterion(preds, text, preds_size, length) / batch_size
else:
preds = preds.log_softmax(2).permute(1, 0, 2)
cost = criterion(preds, text, preds_size, length)
else:
preds = model(image, text[:, :-1]) # align with Attention.forward
target = text[:, 1:] # without [GO] Symbol
cost = criterion(preds.view(-1, preds.shape[-1]),
target.contiguous().view(-1))
model.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)
# validation part
if (
iteration + 1
) % opt.valInterval == 0 or iteration == 0: # To see training progress, we also conduct validation when 'iteration == 0'
elapsed_time = time.time() - start_time
# for log
with open(f'./saved_models/{opt.exp_name}/log_train.txt',
'a') as log:
model.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)
model.train()
# training loss and validation loss
loss_log = f'[{iteration+1}/{opt.num_iter}] Train loss: {loss_avg.val():0.5f}, Valid loss: {valid_loss:0.5f}, Elapsed_time: {elapsed_time:0.5f}'
loss_avg.reset()
current_model_log = f'{"Current_accuracy":17s}: {current_accuracy:0.3f}, {"Current_norm_ED":17s}: {current_norm_ED:0.2f}'
# keep best accuracy model (on valid dataset)
if current_accuracy > best_accuracy:
best_accuracy = current_accuracy
torch.save(
model.state_dict(),
f'./saved_models/{opt.exp_name}/best_accuracy.pth')
if current_norm_ED > best_norm_ED:
best_norm_ED = current_norm_ED
torch.save(
model.state_dict(),
f'./saved_models/{opt.exp_name}/best_norm_ED.pth')
best_model_log = f'{"Best_accuracy":17s}: {best_accuracy:0.3f}, {"Best_norm_ED":17s}: {best_norm_ED:0.2f}'
loss_model_log = f'{loss_log}\n{current_model_log}\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":25s} | {"Prediction":25s} | Confidence Score & T/F'
predicted_result_log = f'{dashed_line}\n{head}\n{dashed_line}\n'
for gt, pred, confidence in zip(labels[:5], preds[:5],
confidence_score[:5]):
if 'Attn' in opt.Prediction:
gt = gt[:gt.find('[s]')]
pred = pred[:pred.find('[s]')]
predicted_result_log += f'{gt:25s} | {pred:25s} | {confidence:0.4f}\t{str(pred == gt)}\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(),
f'./saved_models/{opt.exp_name}/iter_{iteration+1}.pth')
if (iteration + 1) == opt.num_iter:
print('end the training')
sys.exit()
iteration += 1
model = Convnet().cuda()
model.load_state_dict(torch.load(args.load))
model.eval()
ave_acc = Averager()
for i, batch in enumerate(loader, 1):
data, _ = [_.cuda() for _ in batch]
k = args.way * args.shot
data_shot, data_query = data[:k], data[k:]
x = model(data_shot)
x = x.reshape(args.shot, args.way, -1).mean(dim=0)
p = x
logits = euclidean_metric(model(data_query), p)
label = torch.arange(args.way).repeat(args.query)
label = label.type(torch.cuda.LongTensor)
acc = count_acc(logits, label)
ave_acc.add(acc)
print('batch {}: {:.2f}({:.2f})'.format(i,
ave_acc.item() * 100,
acc * 100))
x = None
p = None
logits = None
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
global_new, proto_new = model_reg(support_set=torch.cat(
[global_base[0], global_novel[0]]),
query_set=proto_final)
logits2 = euclidean_metric(proto_new, global_new)
loss2 = F.cross_entropy(logits2, train_gt)
similarity = F.softmax(logits2)
feature = torch.matmul(
similarity, torch.cat([global_base[0], global_novel[0]]))
logits = euclidean_metric(model_cnn(data_query), feature)
loss1 = F.cross_entropy(logits, label)
acc1 = count_acc(logits, label)
acc2 = count_acc(similarity, train_gt)
tl1.add(loss1.item())
tl2.add(loss2.item())
ta1.add(acc1)
ta2.add(acc2)
optimizer_gen.zero_grad()
optimizer_cnn.zero_grad()
optimizer_atten.zero_grad()
optimizer_global1.zero_grad()
optimizer_global2.zero_grad()
total_loss = loss1 + loss2
#loss.backward()
total_loss.backward()
if epoch > 45:
optimizer_gen.step()
for i, batch in enumerate(train_loader, 1):
data, _ = [_.cuda() for _ in batch]
label = torch.arange(args.train_way).repeat(args.query)
label = label.type(torch.cuda.LongTensor)
logits = model(data)
loss = F.cross_entropy(logits, label)
acc = count_acc(logits, label)
print('epoch {}, train {}/{}, loss={:.4f} acc={:.4f}'.format(
epoch, i, len(train_loader), loss.item(), acc))
tl.add(loss.item())
ta.add(acc)
optimizer.zero_grad()
loss.backward()
optimizer.step()
proto = None
logits = None
loss = None
tl = tl.item()
ta = ta.item()
model.eval()
def validation(model, criterion, evaluation_loader, converter, opt):
""" validation or evaluation """
n_correct = 0
norm_ED = 0
length_of_data = 0
infer_time = 0
valid_loss_avg = Averager()
for i, (image_tensors, labels) in enumerate(evaluation_loader):
batch_size = image_tensors.size(0)
length_of_data = length_of_data + batch_size
image = image_tensors.to(device)
# For max length prediction
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)
text_for_loss, length_for_loss = converter.encode(
labels, batch_max_length=opt.batch_max_length)
start_time = time.time()
if 'CTC' in opt.Prediction:
preds = model(image, text_for_pred).log_softmax(2)
forward_time = time.time() - start_time
# Calculate evaluation loss for CTC deocder.
preds_size = torch.IntTensor([preds.size(1)] * batch_size)
preds = preds.permute(1, 0, 2) # to use CTCloss format
# To avoid ctc_loss issue, disabled cudnn for the computation of the ctc_loss
# https://github.com/jpuigcerver/PyLaia/issues/16
torch.backends.cudnn.enabled = False
cost = criterion(preds, text_for_loss, preds_size, length_for_loss)
torch.backends.cudnn.enabled = True
# Select max probabilty (greedy decoding) then decode index to character
_, preds_index = preds.max(2)
preds_index = preds_index.transpose(1, 0).contiguous().view(-1)
preds_str = converter.decode(preds_index.data, preds_size.data)
else:
preds = model(image, text_for_pred, is_train=False)
forward_time = time.time() - start_time
preds = preds[:, :text_for_loss.shape[1] - 1, :]
target = text_for_loss[:, 1:] # without [GO] Symbol
cost = criterion(preds.contiguous().view(-1, preds.shape[-1]),
target.contiguous().view(-1))
# select max probabilty (greedy decoding) then decode index to character
_, preds_index = preds.max(2)
preds_str = converter.decode(preds_index, length_for_pred)
labels = converter.decode(text_for_loss[:, 1:], length_for_loss)
infer_time += forward_time
valid_loss_avg.add(cost)
# calculate accuracy.
for pred, gt in zip(preds_str, labels):
if 'Attn' in opt.Prediction:
pred = pred[:pred.find(
'[s]')] # prune after "end of sentence" token ([s])
gt = gt[:gt.find('[s]')]
if pred == gt:
n_correct += 1
if len(gt) == 0:
norm_ED += 1
else:
norm_ED += edit_distance(pred, gt) / len(gt)
accuracy = n_correct / float(length_of_data) * 100
return valid_loss_avg.val(
), accuracy, norm_ED, preds_str, labels, infer_time, length_of_data
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)
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=False, # '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)
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 = MsImageDis(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()
#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 != '':
try:
start_iter = int(opt.saved_synth_model.split('_')[-1].split('.')[0])
print(f'continue to train, start_iter: {start_iter}')
except:
pass
lexicons=[]
out_of_char = f'[^{opt.character}]'
#read lexicons file
with open(opt.lexFile,'r') as lexF:
for line in lexF:
lexWord = line[:-1]
if opt.fixedString and len(lexWord)!=opt.batch_exact_length:
continue
if len(lexWord) <= opt.batch_max_length and not(re.search(out_of_char, lexWord.lower())) and len(lexWord) >= opt.batch_min_length:
lexicons.append(lexWord)
#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
# pdb.set_trace()
for i, (image_tensors_all, labels_1_all) in enumerate(train_loader):
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_tensors, image_tensors_real, labels_gt = image_tensors_all[:disCnt], image_tensors_all[disCnt:disCnt+disCnt], labels_1_all[:disCnt]
image = image_tensors.to(device)
image_real = image_tensors_real.to(device)
batch_size = image.size(0)
labels_2 = random.sample(lexicons, batch_size)
##-----------------------------------##
#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]) # align with Attention.forward
target_1 = text_1[:, 1:] # without [GO] Symbol
preds_2 = ocrModel(images_recon_2, text_2[:, :-1]) # 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)
print('OCR shape::::::',preds_1.shape,target_1.shape,preds_2.shape,target_2.shape)
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
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))
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)
# #[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()
print('Cost:::::', cost)
# cost.backward()
recCost.backward(retain_graph=True)
disGenCost.backward(retain_graph=True)
ocrCost.backward(retain_graph=True)
styleRecCost.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 + 1) % 1e+5 == 0:
torch.save(
model.state_dict(), os.path.join(opt.exp_dir,opt.exp_name,'iter_'+str(iteration+1)+'.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
def train(opt, show_number=2, amp=False):
""" 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'
)
opt.select_data = opt.select_data.split('-')
opt.batch_ratio = opt.batch_ratio.split('-')
train_dataset = Batch_Balanced_Dataset(opt)
log = open(f'./saved_models/{opt.experiment_name}/log_dataset.txt',
'a',
encoding="utf8")
AlignCollate_valid = AlignCollate(imgH=opt.imgH,
imgW=opt.imgW,
keep_ratio_with_pad=opt.PAD,
contrast_adjust=opt.contrast_adjust)
valid_dataset, valid_dataset_log = hierarchical_dataset(
root=opt.valid_data, opt=opt)
valid_loader = torch.utils.data.DataLoader(
valid_dataset,
batch_size=min(32, opt.batch_size),
shuffle=
True, # 'True' to check training progress with validation function.
num_workers=int(opt.workers),
prefetch_factor=512,
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 = Model(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)
if opt.saved_model != '':
pretrained_dict = torch.load(opt.saved_model)
if opt.new_prediction:
model.Prediction = nn.Linear(
model.SequenceModeling_output,
len(pretrained_dict['module.Prediction.weight']))
model = torch.nn.DataParallel(model).to(device)
print(f'loading pretrained model from {opt.saved_model}')
if opt.FT:
model.load_state_dict(pretrained_dict, strict=False)
else:
model.load_state_dict(pretrained_dict)
if opt.new_prediction:
model.module.Prediction = nn.Linear(
model.module.SequenceModeling_output, opt.num_class)
for name, param in model.module.Prediction.named_parameters():
if 'bias' in name:
init.constant_(param, 0.0)
elif 'weight' in name:
init.kaiming_normal_(param)
model = model.to(device)
else:
# 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
model = torch.nn.DataParallel(model).to(device)
model.train()
print("Model:")
print(model)
count_parameters(model)
""" setup loss """
if 'CTC' in opt.Prediction:
criterion = torch.nn.CTCLoss(zero_infinity=True).to(device)
else:
criterion = torch.nn.CrossEntropyLoss(ignore_index=0).to(
device) # ignore [GO] token = ignore index 0
# loss averager
loss_avg = Averager()
# freeze some layers
try:
if opt.freeze_FeatureFxtraction:
for param in model.module.FeatureExtraction.parameters():
param.requires_grad = False
if opt.freeze_SequenceModeling:
for param in model.module.SequenceModeling.parameters():
param.requires_grad = False
except:
pass
# 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, 0.999))
optimizer = optim.Adam(filtered_parameters)
else:
optimizer = optim.Adadelta(filtered_parameters,
lr=opt.lr,
rho=opt.rho,
eps=opt.eps)
print("Optimizer:")
print(optimizer)
""" final options """
# print(opt)
with open(f'./saved_models/{opt.experiment_name}/opt.txt',
'a',
encoding="utf8") 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_model != '':
try:
start_iter = int(opt.saved_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
i = start_iter
scaler = GradScaler()
t1 = time.time()
while (True):
# train part
optimizer.zero_grad(set_to_none=True)
if amp:
with autocast():
image_tensors, labels = train_dataset.get_batch()
image = image_tensors.to(device)
text, length = converter.encode(
labels, batch_max_length=opt.batch_max_length)
batch_size = image.size(0)
if 'CTC' in opt.Prediction:
preds = model(image, text).log_softmax(2)
preds_size = torch.IntTensor([preds.size(1)] * batch_size)
preds = preds.permute(1, 0, 2)
torch.backends.cudnn.enabled = False
cost = criterion(preds, text.to(device),
preds_size.to(device), length.to(device))
torch.backends.cudnn.enabled = True
else:
preds = model(image,
text[:, :-1]) # align with Attention.forward
target = text[:, 1:] # without [GO] Symbol
cost = criterion(preds.view(-1, preds.shape[-1]),
target.contiguous().view(-1))
scaler.scale(cost).backward()
scaler.unscale_(optimizer)
torch.nn.utils.clip_grad_norm_(model.parameters(), opt.grad_clip)
scaler.step(optimizer)
scaler.update()
else:
image_tensors, labels = train_dataset.get_batch()
image = image_tensors.to(device)
text, length = converter.encode(
labels, batch_max_length=opt.batch_max_length)
batch_size = image.size(0)
if 'CTC' in opt.Prediction:
preds = model(image, text).log_softmax(2)
preds_size = torch.IntTensor([preds.size(1)] * batch_size)
preds = preds.permute(1, 0, 2)
torch.backends.cudnn.enabled = False
cost = criterion(preds, text.to(device), preds_size.to(device),
length.to(device))
torch.backends.cudnn.enabled = True
else:
preds = model(image,
text[:, :-1]) # align with Attention.forward
target = text[:, 1:] # without [GO] Symbol
cost = criterion(preds.view(-1, preds.shape[-1]),
target.contiguous().view(-1))
cost.backward()
torch.nn.utils.clip_grad_norm_(model.parameters(), opt.grad_clip)
optimizer.step()
loss_avg.add(cost)
# validation part
if (i % opt.valInterval == 0) and (i != 0):
print('training time: ', time.time() - t1)
t1 = time.time()
elapsed_time = time.time() - start_time
# for log
with open(f'./saved_models/{opt.experiment_name}/log_train.txt',
'a',
encoding="utf8") as log:
model.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, device)
model.train()
# training loss and validation loss
loss_log = f'[{i}/{opt.num_iter}] Train loss: {loss_avg.val():0.5f}, Valid loss: {valid_loss:0.5f}, Elapsed_time: {elapsed_time:0.5f}'
loss_avg.reset()
current_model_log = f'{"Current_accuracy":17s}: {current_accuracy:0.3f}, {"Current_norm_ED":17s}: {current_norm_ED:0.4f}'
# keep best accuracy model (on valid dataset)
if current_accuracy > best_accuracy:
best_accuracy = current_accuracy
torch.save(
model.state_dict(),
f'./saved_models/{opt.experiment_name}/best_accuracy.pth'
)
if current_norm_ED > best_norm_ED:
best_norm_ED = current_norm_ED
torch.save(
model.state_dict(),
f'./saved_models/{opt.experiment_name}/best_norm_ED.pth'
)
best_model_log = f'{"Best_accuracy":17s}: {best_accuracy:0.3f}, {"Best_norm_ED":17s}: {best_norm_ED:0.4f}'
loss_model_log = f'{loss_log}\n{current_model_log}\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":25s} | {"Prediction":25s} | Confidence Score & T/F'
predicted_result_log = f'{dashed_line}\n{head}\n{dashed_line}\n'
#show_number = min(show_number, len(labels))
start = random.randint(0, len(labels) - show_number)
for gt, pred, confidence in zip(
labels[start:start + show_number],
preds[start:start + show_number],
confidence_score[start:start + show_number]):
if 'Attn' in opt.Prediction:
gt = gt[:gt.find('[s]')]
pred = pred[:pred.find('[s]')]
predicted_result_log += f'{gt:25s} | {pred:25s} | {confidence:0.4f}\t{str(pred == gt)}\n'
predicted_result_log += f'{dashed_line}'
print(predicted_result_log)
log.write(predicted_result_log + '\n')
print('validation time: ', time.time() - t1)
t1 = time.time()
# save model per 1e+4 iter.
if (i + 1) % 1e+4 == 0:
torch.save(model.state_dict(),
f'./saved_models/{opt.experiment_name}/iter_{i+1}.pth')
if i == opt.num_iter:
print('end the training')
sys.exit()
i += 1
