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'])