def cross_domain(model_path, device='cpu'):
model = ConceptTagger.load(model_path, device)
config = model.config
log = Logger(os.path.join(config.save_dir, '_tgt.txt'), level='info')
dataDict = getNERdata(dataSetName=config.dataset,
dataDir=config.data_dir,
desc_path=config.description_path,
cross_domain=config.cross_domain,
target_domain=config.target_domain)
model.to(device)
test_metric_pre, test_metric_rec, test_metric_f1 = evaluate(
model, dataDict['target']['test'], config.batch_size, log)
log.logger.info("test_pre : %.4f, test_rec : %.4f, test_f1 : %.4f" %
(test_metric_pre, test_metric_rec, test_metric_f1))
def cross_domain(model_path, device='cuda:1'):
model = Bilstm_LabelEmbedding.load(model_path, device)
config = model.config
log = Logger(os.path.join(config.save_dir, '_tgt.txt'), level='info')
dataDict = getNERdata(dataSetName=config.dataset,
dataDir=config.data_dir,
desc_path=config.description_path,
cross_domain=config.cross_domain,
exemplar_num=config.exemplar_num,
target_domain=config.target_domain)
tgt_label2Idx = ExtractLabelsFromTokens(dataDict['target']['test'])
model.label2Idx = tgt_label2Idx
Emb = Bilstm_LabelEmbedding.BuildLabelEmbedding(
model.embedding, model.word2Idx, tgt_label2Idx, model.description,
dataDict['exemplar_test'], config.embedding_method,
config.encoder_method, device)
description_emb_test = BuildEmb.buildembedding(
model.embedding, model.word2Idx, tgt_label2Idx,
dataDict['description'], None, 'description', 'wordembedding',
config.device)
model.LabelEmbedding = torch.cat((Emb, description_emb_test), 1)
if config.crf:
model.crf.labelembedding = model.crf.buildCRFLabelEmbedding(
model.LabelEmbedding)
model.crf.num_tags = model.LabelEmbedding.size(0)
model.to(device)
(test_metric_pre, test_metric_rec,
test_metric_f1), d = evaluate(model, dataDict['target']['test'],
config.batch_size, log)
f = open(os.path.join(config.save_dir, 'result.txt'), 'a+')
js = json.dumps(d)
f.write(js + '\n')
f.close()
# print("test_pre : %.4f, test_rec : %.4f, test_f1 : %.4f" % (test_metric_pre, test_metric_rec, test_metric_f1),
# file=sys.stderr)
log.logger.info("test_pre : %.4f, test_rec : %.4f, test_f1 : %.4f" %
(test_metric_pre, test_metric_rec, test_metric_f1))
def train(config):
dataDict = getNERdata(dataSetName=config.dataset,
dataDir=config.data_dir,
desc_path=config.description_path,
cross_domain=config.cross_domain,
exemplar_num=config.exemplar_num,
target_domain=config.target_domain)
emb, word2Idx = readTokenEmbeddings(config.embed_file)
char2Idx = getCharIdx()
label2Idx = ExtractLabelsFromTokens(dataDict['source']['train'])
label2IdxForDev = ExtractLabelsFromTokens(dataDict['target']['dev'])
label2IdxForTest = ExtractLabelsFromTokens(dataDict['target']['test'])
print(label2IdxForDev)
print(dataDict['exemplar_dev'])
DevLabelEmbedding = Bilstm_LabelEmbedding.BuildLabelEmbedding(
emb, word2Idx, label2IdxForDev, dataDict['description'],
dataDict['exemplar_dev'], config.embedding_method,
config.encoder_method, config.device)
TestLabelEmbedding = Bilstm_LabelEmbedding.BuildLabelEmbedding(
emb, word2Idx, label2IdxForTest, dataDict['description'],
dataDict['exemplar_test'], config.embedding_method,
config.encoder_method, config.device)
max_batch_size = math.ceil(
len(dataDict['source']['train']) / config.batch_size)
model = Bilstm_LabelEmbedding(config, emb, word2Idx, label2Idx, char2Idx,
dataDict['description'],
dataDict['exemplar_train'])
model.train()
model = model.to(config.device)
hist_valid_scores = []
patience = num_trial = 0
train_iter = 0
optimizer = torch.optim.Adam(model.parameters(), lr=config.learning_rate)
train_time = time.time()
config.save_dir = config.save_dir + config.target_domain + '/'
if not os.path.exists(config.save_dir):
os.mkdir(config.save_dir)
if os.path.exists(os.path.join(config.save_dir, 'params')):
os.remove(os.path.join(config.save_dir, 'params'))
log = Logger(os.path.join(config.save_dir, '_src.txt'), level='info')
for epoch in range(config.epoch):
for da in data_generator(dataDict['source']['train'],
config.batch_size):
train_iter += 1
x = da[0]
y = da[1]
loss = model(x, y, 'train')
optimizer.zero_grad()
loss.backward()
optimizer.step()
if train_iter % config.log_every == 0:
# print(
# 'epoch %d, iter %d, loss %.2f, time elapsed %.2f sec' %
# (epoch, train_iter, loss, time.time() - train_time),
# file=sys.stderr)
log.logger.info(
'epoch %d, iter %d, loss %.2f, time elapsed %.2f sec' %
(epoch, train_iter, loss, time.time() - train_time))
train_time = time.time()
if train_iter % config.log_valid == 0:
trainLabelEmbedding = model.LabelEmbedding
trainLabel2Idx = model.label2Idx
model.label2Idx = label2IdxForDev
model.LabelEmbedding = DevLabelEmbedding
if config.crf:
model.crf.labelembedding = model.crf.buildCRFLabelEmbedding(
model.LabelEmbedding)
model.crf.num_tags = model.LabelEmbedding.size(0)
(valid_metric_pre, valid_metric_rec,
valid_metric_f1), d = evaluate(model,
dataDict['target']['dev'],
config.batch_size, log)
model.label2Idx = label2IdxForTest
model.LabelEmbedding = TestLabelEmbedding
if config.crf:
model.crf.labelembedding = model.crf.buildCRFLabelEmbedding(
model.LabelEmbedding)
model.crf.num_tags = model.LabelEmbedding.size(0)
(test_metric_pre, test_metric_rec,
test_metric_f1), d = evaluate(model,
dataDict['target']['test'],
config.batch_size, log)
model.label2Idx = label2Idx
model.LabelEmbedding = trainLabelEmbedding
if config.crf:
model.crf.labelembedding = model.crf.buildCRFLabelEmbedding(
model.LabelEmbedding)
model.crf.num_tags = model.LabelEmbedding.size(0)
# print("val_pre : %.4f, val_rec : %.4f, val_f1 : %.4f" % (valid_metric_pre, valid_metric_rec, valid_metric_f1), file=sys.stderr)
# print("test_pre : %.4f, test_rec : %.4f, test_f1 : %.4f" % (test_metric_pre, test_metric_rec, test_metric_f1), file=sys.stderr)
log.logger.info(
"val_pre : %.4f, val_rec : %.4f, val_f1 : %.4f" %
(valid_metric_pre, valid_metric_rec, valid_metric_f1))
log.logger.info(
"test_pre : %.4f, test_rec : %.4f, test_f1 : %.4f" %
(test_metric_pre, test_metric_rec, test_metric_f1))
is_better = len(
hist_valid_scores
) == 0 or valid_metric_f1 > max(hist_valid_scores)
hist_valid_scores.append(valid_metric_f1)
if is_better:
patience = 0
# print('save currently the best model to [%s]' % (config.save_dir + 'model'), file=sys.stderr)
log.logger.info('save currently the best model to [%s]' %
(config.save_dir + 'model'))
model.save(config.save_dir + 'model')
# also save the optimizers' state
torch.save(optimizer.state_dict(),
config.save_dir + 'optim')
elif patience < config.patience:
patience += 1
log.logger.info('hit patience %d' % patience)
# print('hit patience %d' % patience, file=sys.stderr)
if patience == int(config.patience):
num_trial += 1
log.logger.info('hit #%d trial' % num_trial)
# print('hit #%d trial' % num_trial, file=sys.stderr)
if num_trial == config.max_num_trial:
log.logger.info('early stop!')
# print('early stop!', file=sys.stderr)
exit(0)
lr = optimizer.param_groups[0]['lr'] * config.lr_decay
log.logger.info(
'load previously best model and decay learning rate to %f'
% lr)
# print('load previously best model and decay learning rate to %f' % lr, file=sys.stderr)
# load model
params = torch.load(
config.save_dir + 'model',
map_location=lambda storage, loc: storage)
model.load_state_dict(params['state_dict'])
model = model.to(config.device)
log.logger.info('restore parameters of the optimizers')
# print('restore parameters of the optimizers', file=sys.stderr)
optimizer.load_state_dict(
torch.load(config.save_dir + 'optim'))
# set new lr
for param_group in optimizer.param_groups:
param_group['lr'] = lr
# reset patience
patience = 0
def main(args):
log_write = Logger('./log.txt', 'LogFile')
log_write.set_names([
'Total loss', 'Classified loss', 'Y location loss', 'X Refine loss',
'Learning Rate'
])
data_loader = IC15Loader(args.size_list)
gt_files = data_loader.gt_paths
train_loader = torch.utils.data.DataLoader(data_loader,
batch_size=args.batch_size,
shuffle=True,
num_workers=args.num_worker,
drop_last=True,
pin_memory=True)
model = CTPN_Model(base_model=args.base_model,
pretrained=args.pretrain).cuda()
critetion = CTPNLoss().cuda()
if (args.restore != ''):
model.load_state_dict(torch.load(args.restore))
if (args.optimizer == 'SGD'):
optimizer = optim.SGD(model.parameters(),
lr=args.lr,
momentum=0.99,
weight_decay=5e-4)
else:
optimizer = optim.Adam(model.parameters(),
lr=args.lr,
weight_decay=5e-4)
scheduler = optim.lr_scheduler.StepLR(optimizer,
step_size=args.step_size,
gamma=args.gamma)
model.train()
for epoch in range(args.train_epochs):
loss_total_list = []
loss_cls_list = []
loss_ver_list = []
loss_refine_list = []
for batch_idx, (imgs, img_scales, im_shapes, gt_path_indexs,
im_infos) in enumerate(train_loader):
data_loader.get_random_train_size()
optimizer.zero_grad()
image = Variable(imgs.cuda())
score_pre, vertical_pred, side_refinement = model(image)
score_pre = score_pre.permute(0, 2, 3, 1)
vertical_pred = vertical_pred.permute(0, 2, 3, 1)
side_refinement = side_refinement.permute(0, 2, 3, 1)
batch_res_polys = get_bboxes(imgs, gt_files, gt_path_indexs,
img_scales, im_shapes)
batch_loss_tatal = []
batch_loss_cls = []
batch_loss_ver = []
batch_loss_refine = []
for i in range(image.shape[0]):
image_ori = (imgs[i].numpy() * 255).transpose((1, 2, 0)).copy()
gt_boxes = np.array(batch_res_polys[i])
rpn_labels, rpn_bbox_targets, rpn_bbox_inside_weights, rpn_bbox_outside_weights = anchor_target_layer(
image_ori, score_pre[i].cpu().unsqueeze(0), gt_boxes,
im_infos[i].numpy())
rpn_labels = toTensor(rpn_labels)
rpn_bbox_targets = toTensor(rpn_bbox_targets)
rpn_bbox_inside_weights = toTensor(rpn_bbox_inside_weights)
rpn_bbox_outside_weights = toTensor(rpn_bbox_outside_weights)
loss_tatal, loss_cls, loss_ver, loss_refine = critetion(
score_pre[i].unsqueeze(0), vertical_pred[i].unsqueeze(0),
rpn_labels, rpn_bbox_targets)
batch_loss_tatal.append(loss_tatal)
batch_loss_cls.append(loss_cls)
batch_loss_ver.append(loss_ver)
batch_loss_refine.append(loss_refine)
del (loss_tatal)
del (loss_cls)
del (loss_ver)
del (loss_refine)
loss_tatal = sum(batch_loss_tatal) / len(batch_loss_tatal)
loss_cls = sum(batch_loss_cls) / len(batch_loss_cls)
loss_ver = sum(batch_loss_ver) / len(batch_loss_ver)
loss_refine = sum(batch_loss_refine) / len(batch_loss_refine)
loss_tatal.backward()
optimizer.step()
loss_total_list.append(loss_tatal.item())
loss_cls_list.append(loss_cls.item())
loss_ver_list.append(loss_ver.item())
loss_refine_list.append(loss_refine.item())
if (batch_idx % args.show_step == 0):
log = '({epoch}/{epochs}/{batch_i}/{all_batch}) | loss_tatal: {loss1:.4f} | loss_cls: {loss2:.4f} | loss_ver: {loss3:.4f} | loss_refine: {loss4:.4f} | Lr: {lr}'.format(
epoch=epoch,
epochs=args.train_epochs,
batch_i=batch_idx,
all_batch=len(train_loader),
loss1=loss_tatal.item(),
loss2=loss_cls.item(),
loss3=loss_ver.item(),
loss4=loss_refine.item(),
lr=scheduler.get_lr()[0])
print(log)
log_write.append([
loss_tatal.item(),
loss_cls.item(),
loss_ver.item(),
loss_refine.item(),
scheduler.get_lr()[0]
])
print(
'--------------------------------------------------------------------------------------------------------'
)
log_write.set_split(
['---------', '----------', '--------', '----------', '--------'])
print(
"epoch_loss_total:{loss1:.4f} | epoch_loss_cls:{loss2:.4f} | epoch_loss_ver:{loss3:.4f} | epoch_loss_ver:{loss4:.4f} | Lr:{lr}"
.format(loss1=np.mean(loss_total_list),
loss2=np.mean(loss_cls_list),
loss3=np.mean(loss_ver_list),
loss4=np.mean(loss_refine_list),
lr=scheduler.get_lr()[0]))
log_write.append([
np.mean(loss_total_list),
np.mean(loss_cls_list),
np.mean(loss_ver_list),
np.mean(loss_refine_list),
scheduler.get_lr()[0]
])
print(
'-------------------------------------------------------------------------------------------------------'
)
log_write.set_split(
['---------', '----------', '--------', '----------', '--------'])
if (epoch % args.epoch_save == 0 and epoch != 0):
torch.save(
model.state_dict(),
os.path.join(args.checkpoint, 'ctpn_' + str(epoch) + '.pth'))
scheduler.step()
log_write.close()