def training(self):
self.DNet.eval()
self.RNet.train()
loss_t = {'sum': 0}
# train start
print('train start =====> RNet')
logger('RNet train start\n', self.logfile)
for i, batch in enumerate(self.dataloader):
# load data
img = batch['img'].cuda()
# generate line pairs per image
pairwise = generate_line_pair_RNet(
training_line=batch['training_data'],
line_batch_size=self.cfg.batch_size['train_pair'],
mode='training')
# extract ref & tar line features from detector
f_ref = self.forward_model.run_feature_extractor(
img, pairwise['ref'], self.DNet)
f_tar = self.forward_model.run_feature_extractor(
img, pairwise['tar'], self.DNet)
# model
out = self.RNet(f_ref, f_tar)
# loss
loss = self.loss_fn(out, pairwise['label'])
# optimize
self.optimizer.zero_grad()
loss.backward()
self.optimizer.step()
loss_t['sum'] += loss.item()
# display
if i % 100 == 0:
print('iter %d' % i)
self.visualize.display_for_train_RNet(batch, pairwise, out, i)
logger("Loss : %5f\n" % loss.item(), self.logfile)
# logger
logger("Average Loss : %5f\n" % (loss_t['sum'] / len(self.dataloader)),
self.logfile)
print("Average Loss : %5f\n" % (loss_t['sum'] / len(self.dataloader)))
# save model
self.ckpt = {
'epoch': self.epoch,
'model': self.RNet,
'optimizer': self.optimizer,
'val_result': self.val_result
}
save_model(checkpoint=self.ckpt,
param='checkpoint_RNet_final',
path=self.cfg.weight_dir + '')
def training(self):
self.DNet.train()
loss_t = {'sum': 0, 'cls': 0, 'reg': 0}
# train start
print('train start =====> DNet')
logger('DNet train start\n', self.logfile)
for i, batch in enumerate(self.dataloader):
# shuffle idx with pos:neg = 4:6
idx = torch.randperm(self.cfg.batch_size['train_line'])
batch['train_data'] = batch['train_data'].cuda()
batch['train_data'][0, :] = batch['train_data'][0, idx]
# load data
img = batch['img'].cuda()
candidates = batch['train_data'][:, :, :4]
gt_cls = batch['train_data'][:, :, 4:6]
gt_reg = batch['train_data'][:, :, 6:]
# model
out = self.DNet(img=img, line_pts=candidates)
# loss
loss, loss_cls, loss_reg = self.loss_fn(output=out,
gt_cls=gt_cls,
gt_reg=gt_reg)
# optimize
self.optimizer.zero_grad()
loss.backward()
self.optimizer.step()
loss_t['sum'] += loss.item()
loss_t['cls'] += loss_cls.item()
loss_t['reg'] += loss_reg.item()
# display
if i % 100 == 0:
print('iter %d' % i)
self.visualize.display_for_train_detector(batch, out, i)
logger(
"Loss : %5f, "
"Loss_cls : %5f, "
"Loss_reg : %5f\n" %
(loss.item(), loss_cls.item(), loss_reg.item()),
self.logfile)
# logger
logger(
"Average Loss : %5f %5f %5f\n" %
(loss_t['sum'] / len(self.dataloader), loss_t['cls'] /
len(self.dataloader), loss_t['reg'] / len(self.dataloader)),
self.logfile)
print("Average Loss : %5f %5f %5f\n" %
(loss_t['sum'] / len(self.dataloader), loss_t['cls'] /
len(self.dataloader), loss_t['reg'] / len(self.dataloader)))
# save model
self.ckpt = {
'epoch': self.epoch,
'model': self.DNet,
'optimizer': self.optimizer,
'val_result': self.val_result
}
save_model(checkpoint=self.ckpt,
param='checkpoint_DNet_final',
path=self.cfg.weight_dir)
def training(self):
self.SLNet.train()
loss_t = {'sum': 0, 'cls': 0, 'reg': 0}
correct = 0
error = 0
# train start
print('train start =====> SLNet')
logger('SLNet train start\n', self.logfile)
for i, batch in enumerate(self.dataloader):
# shuffle idx with pos:neg = 4:6
idx = torch.randperm(self.cfg.batch_size['train_line'])
batch['train_data'] = batch['train_data'].cuda()
batch['train_data'][0, :] = batch['train_data'][0, idx]
# load data
img = batch['img'].cuda()
candidates = batch['train_data'][:, :, :4].float()
gt_cls = batch['train_data'][:, :, 4:8]
gt_reg = batch['train_data'][:, :, 8:]
# model
out = self.SLNet(img=img, line_pts=candidates)
# loss
loss, loss_cls, loss_reg = self.loss_fn(output=out,
gt_cls=gt_cls,
gt_reg=gt_reg)
# accuracy
for j in range(batch['train_data'].shape[0]):
outans = torch.argmax(out['cls'][j], dim=1)
gtans = torch.argmax(gt_cls[j], dim=1)
for k in range(outans.shape[0]):
if outans[k] == gtans[k]:
correct += 1
else:
error += 1
# optimize
self.optimizer.zero_grad()
loss.backward()
self.optimizer.step()
loss_t['sum'] += loss.item()
loss_t['cls'] += loss_cls.item()
loss_t['reg'] += loss_reg.item()
# display
if i % 100 == 0:
print('iter %d' % i)
self.visualize.display_for_train_detector(batch, out, i)
logger(
"Loss : %5f, "
"Loss_cls : %5f, "
"Loss_reg : %5f\n" %
(loss.item(), loss_cls.item(), loss_reg.item()),
self.logfile)
# logger
logger(
"Average Loss : %5f %5f %5f\n" %
(loss_t['sum'] / len(self.dataloader), loss_t['cls'] /
len(self.dataloader), loss_t['reg'] / len(self.dataloader)),
self.logfile)
print("Average Loss : %5f %5f %5f\n" %
(loss_t['sum'] / len(self.dataloader), loss_t['cls'] /
len(self.dataloader), loss_t['reg'] / len(self.dataloader)))
wandb.log({
"accuracy": correct / (correct + error),
"loss_sum": loss_t['sum'] / len(self.dataloader),
"loss_cls": loss_t['cls'] / len(self.dataloader),
"loss_reg": loss_t['reg'] / len(self.dataloader),
})
# save model
self.ckpt = {
'epoch': self.epoch,
'model': self.SLNet,
'optimizer': self.optimizer,
'val_result': self.val_result
}
save_model(checkpoint=self.ckpt,
param='checkpoint_SLNet_final',
path=self.cfg.weight_dir)
torch.save(self.SLNet.state_dict(),
os.path.join(wandb.run.dir, 'model.pt'))
def training(self):
loss_t = {'sum': 0, 'score': 0, 'node_loss': 0, 'edge_loss': 0}
# train start
self.s_net.eval()
self.h_net.train()
print('train start')
logger('train start\n', self.logfile)
for case in range(self.cfg.case_num):
rmdir(path=self.cfg.dir['out'] + 'train/display/' + str(case))
for i, batch in enumerate(self.dataloader):
# load data
img = batch['img'].cuda()
a_idx = batch['a_idx'].cuda()
d_idx = batch['d_idx'].cuda()
img_name = batch['img_name'][0]
# H_Net training data generator
self.generator.update_batch(
a_idx,
d_idx,
img_name=img_name,
dataset_name='{}_train'.format(self.dataset_name),
pickle_dir=self.cfg.dir['out'] +
'preprocess/{}_train/pickle/'.format(self.dataset_name))
train_data1, train_data2 = self.generator.run_for_training_data(
a_idx, d_idx)
self.count['node'] += train_data1['node_count']
self.count['edge'] += train_data2['edge_count']
self.count['tot'] = (self.count['node'] + self.count['edge'])
# model
self.h_net.forward_encoder(img, train_data2, is_training=True)
out1 = self.h_net.forward_node_score(train_data1, is_training=True)
out2 = self.h_net.forward_edge_score(train_data2, is_training=True)
out_f = dict(out1, **out2)
# loss
loss, node_loss, edge_loss = self.loss_fn(out=out_f,
train_data1=train_data1,
train_data2=train_data2)
# optimize
self.optimizer.zero_grad()
loss.backward()
self.optimizer.step()
loss_t['sum'] += loss.item()
loss_t['node_loss'] += node_loss.item()
loss_t['edge_loss'] += edge_loss.item()
# display
if i % self.cfg.disp_step == 0:
print('img iter %d' % i)
self.visualize.display_for_train_node_score(
batch, out_f, train_data1, i)
self.visualize.display_for_train_edge_score(
batch, out_f, train_data2, i)
if i % self.cfg.disp_step == 0:
logger(
"Loss : %5f, node_Loss : %5f, edge_Loss : %5f\n" %
(loss.item(), node_loss.item(), edge_loss.item()),
self.logfile)
# logger
logger(
"Average Loss : %5f %5f %5f\n" %
(loss_t['sum'] / i, loss_t['node_loss'] / i,
loss_t['edge_loss'] / i), self.logfile)
print("Average Loss : %5f %5f %5f\n" %
(loss_t['sum'] / i, loss_t['node_loss'] / i,
loss_t['edge_loss'] / i))
label_list = ['tot', 'node', 'edge']
for name in label_list:
logger(
"%s label distribution\n%s\n" %
(name,
np.round(
(self.count[name] / np.sum(self.count[name])), 3) * 100),
self.logfile)
print(
"%s label distribution\n%s\n" %
(name,
np.round(
(self.count[name] / np.sum(self.count[name])), 3) * 100))
# save model
self.ckpt = {
'epoch': self.epoch,
'model': self.h_net,
'optimizer': self.optimizer,
'val_result': self.val_result
}
save_model(checkpoint=self.ckpt,
param='checkpoint_final',
path=self.cfg.dir['weight'])
def training(self):
loss_t = {'sum': 0, 'prob': 0, 'off_a': 0, 'off_d': 0}
# train start
self.s_net.train()
print('train start')
logger('train start\n', self.logfile)
for i, batch in enumerate(self.dataloader):
# load data
img = batch['img'].cuda()
prob = batch['prob'].cuda()
offset = batch['offset'].cuda()
img_name = batch['img_name'][0]
# model
out_cls = self.s_net.forward_for_cls(img, is_training=True)
out_reg = self.s_net.forward_for_reg(
l_feat=out_cls['l_feat'],
idx=torch.arange(out_cls['l_feat'].shape[2]).cuda())
out = dict(out_cls, **out_reg)
# loss
loss = self.loss_fn(out=out, gt_prob=prob, gt_offset=offset)
# optimize
self.optimizer.zero_grad()
loss['sum'].backward()
self.optimizer.step()
loss_t['sum'] += loss['sum'].item()
loss_t['prob'] += loss['prob'].item()
loss_t['off_a'] += loss['off_a'].item()
loss_t['off_d'] += loss['off_d'].item()
# display
if i % self.cfg.disp_step == 0:
print('img iter %d ==> %s' % (i, img_name))
self.visualize.display_for_train_reg(batch, out, i)
if i % self.cfg.disp_step == 0:
logger(
"%d %s ==> Loss : %5f, Loss_prob : %5f, Loss_off_a : %5f, Loss_off_d : %5f\n"
%
(i, img_name, loss['sum'].item(), loss['prob'].item(),
loss['off_a'].item(), loss['off_d'].item()),
self.logfile)
# logger
logger(
"Average Loss : %5f %5f %5f %5f\n" %
(loss_t['sum'] / i, loss_t['prob'] / i, loss_t['off_a'] / i,
loss_t['off_d'] / i), self.logfile)
print("Average Loss : %5f %5f %5f %5f\n" %
(loss_t['sum'] / i, loss_t['prob'] / i, loss_t['off_a'] / i,
loss_t['off_d'] / i))
# save model
self.ckpt = {
'epoch': self.epoch,
'model': self.s_net,
'optimizer': self.optimizer,
'val_result': self.val_result
}
save_model(checkpoint=self.ckpt,
param='checkpoint_final',
path=self.cfg.dir['weight'])
