def main(): """ train dataloader """ args = parser.parse_args() data_loader = TrainDataLoader(args.train_path, check=True) if not os.path.exists(args.weight_dir): os.makedirs(args.weight_dir) """ compute max_batches """ for root, dirs, files in os.walk(args.train_path): for dirname in dirs: dir_path = os.path.join(root, dirname) args.max_batches += len(os.listdir(dir_path)) inter = args.max_batches // 10 print('Max batches:{} in one epoch '.format(args.max_batches)) """ Model on gpu """ model = SiameseRPN() model = model.cuda() cudnn.benchmark = True """ loss and optimizer """ criterion = MultiBoxLoss() optimizer = torch.optim.SGD(model.parameters(), lr=args.lr, momentum=args.momentum, weight_decay=args.weight_decay) """ load weights """ init_weights(model) if not args.checkpoint_path == None: assert os.path.isfile( args.checkpoint_path), '{} is not valid checkpoint_path'.format( args.checkpoint_path) try: checkpoint = torch.load(args.checkpoint_path) start = checkpoint['epoch'] model.load_state_dict(checkpoint['state_dict']) optimizer.load_state_dict(checkpoint['optimizer']) except: start = 0 init_weights(model) else: start = 0 """ train phase """ closses, rlosses, tlosses = AverageMeter(), AverageMeter(), AverageMeter() for epoch in range(start, args.max_epoches): cur_lr = adjust_learning_rate(args.lr, optimizer, epoch, gamma=0.1) index_list = range(data_loader.__len__()) #for example in range(args.max_batches): for example in range(900): ret = data_loader.__get__(random.choice(index_list)) template = ret['template_tensor'].cuda() detection = ret['detection_tensor'].cuda() pos_neg_diff = ret['pos_neg_diff_tensor'].cuda( ) if ret['pos_neg_diff_tensor'] is not None else None cout, rout = model(template, detection) predictions = (cout, rout) targets = pos_neg_diff area = ret['area_target_in_resized_detection'] num_pos = len(np.where(pos_neg_diff == 1)[0]) if area == 0 or num_pos == 0 or pos_neg_diff is None: continue closs, rloss, loss, reg_pred, reg_target, pos_index, neg_index = criterion( predictions, targets) # debug for class cout = cout.squeeze().permute(1, 2, 0).reshape(-1, 2) cout = cout.cpu().detach().numpy() print(cout.shape) score = 1 / (1 + np.exp(cout[:, 0] - cout[:, 1])) print(score[pos_index]) print(score[neg_index]) #time.sleep(1) # debug for reg tmp_dir = '/home/song/srpn/tmp/visualization/7_train_debug_pos_anchors' if not os.path.exists(tmp_dir): os.makedirs(tmp_dir) detection = ret['detection_cropped_resized'].copy() draw = ImageDraw.Draw(detection) pos_anchors = ret['pos_anchors'].copy() # pos anchor的回归情况 x = pos_anchors[:, 0] + pos_anchors[:, 2] * reg_pred[ pos_index, 0].cpu().detach().numpy() y = pos_anchors[:, 1] + pos_anchors[:, 3] * reg_pred[ pos_index, 1].cpu().detach().numpy() w = pos_anchors[:, 2] * np.exp(reg_pred[pos_index, 2].cpu().detach().numpy()) h = pos_anchors[:, 3] * np.exp(reg_pred[pos_index, 3].cpu().detach().numpy()) x1s, y1s, x2s, y2s = x - w // 2, y - h // 2, x + w // 2, y + h // 2 for i in range(2): x1, y1, x2, y2 = x1s[i], y1s[i], x2s[i], y2s[i] draw.line([(x1, y1), (x2, y1), (x2, y2), (x1, y2), (x1, y1)], width=1, fill='red') #predict # 应当的gt x = pos_anchors[:, 0] + pos_anchors[:, 2] * reg_target[ pos_index, 0].cpu().detach().numpy() y = pos_anchors[:, 1] + pos_anchors[:, 3] * reg_target[ pos_index, 1].cpu().detach().numpy() w = pos_anchors[:, 2] * np.exp( reg_target[pos_index, 2].cpu().detach().numpy()) h = pos_anchors[:, 3] * np.exp( reg_target[pos_index, 3].cpu().detach().numpy()) x1s, y1s, x2s, y2s = x - w // 2, y - h // 2, x + w // 2, y + h // 2 for i in range(2): x1, y1, x2, y2 = x1s[i], y1s[i], x2s[i], y2s[i] draw.line([(x1, y1), (x2, y1), (x2, y2), (x1, y2), (x1, y1)], width=1, fill='green') #gt # 找分数zui da de, m_indexs = np.argsort(score)[::-1][:5] for m_index in m_indexs: diff = reg_pred[m_index].cpu().detach().numpy() anc = ret['anchors'][m_index] x = anc[0] + anc[0] * diff[0] y = anc[1] + anc[1] * diff[1] w = anc[2] * np.exp(diff[2]) h = anc[3] * np.exp(diff[3]) x1, y1, x2, y2 = x - w // 2, y - h // 2, x + w // 2, y + h // 2 draw.line([(x1, y1), (x2, y1), (x2, y2), (x1, y2), (x1, y1)], width=2, fill='black') save_path = osp.join( tmp_dir, 'epoch_{:04d}_{:04d}_{:02d}.jpg'.format(epoch, example, i)) detection.save(save_path) closs_ = closs.cpu().item() if np.isnan(closs_): sys.exit(0) #loss = closs + rloss closses.update(closs.cpu().item()) rlosses.update(rloss.cpu().item()) tlosses.update(loss.cpu().item()) optimizer.zero_grad() loss.backward() optimizer.step() #time.sleep(1) print( "Epoch:{:04d}\texample:{:08d}/{:08d}({:.2f})\tlr:{:.7f}\tcloss:{:.4f}\trloss:{:.4f}\ttloss:{:.4f}" .format(epoch, example + 1, args.max_batches, 100 * (example + 1) / args.max_batches, cur_lr, closses.avg, rlosses.avg, tlosses.avg)) if epoch % 5 == 0: file_path = os.path.join( args.weight_dir, 'epoch_{:04d}_weights.pth.tar'.format(epoch)) state = { 'epoch': epoch + 1, 'state_dict': model.state_dict(), 'optimizer': optimizer.state_dict(), } torch.save(state, file_path)
def main(): """ train dataloader """ args = parser.parse_args() data_loader = TrainDataLoader(args.train_path, check = args.debug) if not os.path.exists(args.weight_dir): os.makedirs(args.weight_dir) """ compute max_batches """ for root, dirs, files in os.walk(args.train_path): for dirname in dirs: dir_path = os.path.join(root, dirname) args.max_batches += len(os.listdir(dir_path)) print('max_batches: {}'.format(args.max_batches)) """ Model on gpu """ model = SiameseRPN() model = model.cuda() cudnn.benchmark = True """ loss and optimizer """ criterion = MultiBoxLoss() optimizer = torch.optim.SGD(model.parameters(), lr=args.lr, momentum=args.momentum, weight_decay = args.weight_decay) """ load weights """ init_weights(model) if not args.checkpoint_path == None: assert os.path.isfile(args.checkpoint_path), '{} is not valid checkpoint_path'.format(args.checkpoint_path) try: checkpoint = torch.load(args.checkpoint_path) start = checkpoint['epoch'] model.load_state_dict(checkpoint['state_dict']) optimizer.load_state_dict(checkpoint['optimizer']) except: start = 0 init_weights(model) else: start = 0 """ train phase """ closses, rlosses, tlosses = AverageMeter(), AverageMeter(), AverageMeter() steps = 0 #print('data_loader length: {}'.format(len(data_loader))) for epoch in range(start, args.max_epoches): cur_lr = adjust_learning_rate(args.lr, optimizer, epoch, gamma=0.1) index_list = range(data_loader.__len__()) example_index = 0 for example in range(args.max_batches): ret = data_loader.__get__(random.choice(index_list)) template = ret['template_tensor'].cuda() detection= ret['detection_tensor'].cuda() pos_neg_diff = ret['pos_neg_diff_tensor'].cuda() cout, rout = model(template, detection) predictions, targets = (cout, rout), pos_neg_diff closs, rloss, loss, reg_pred, reg_target, pos_index, neg_index = criterion(predictions, targets) closs_ = closs.cpu().item() if np.isnan(closs_): sys.exit(0) closses.update(closs.cpu().item()) rlosses.update(rloss.cpu().item()) tlosses.update(loss.cpu().item()) optimizer.zero_grad() loss.backward() optimizer.step() steps += 1 cout = cout.cpu().detach().numpy() score = 1/(1 + np.exp(cout[:,0]-cout[:,1])) # ++++++++++++ post process below just for debug ++++++++++++++++++++++++ # ++++++++++++++++++++ v1.0 add penalty +++++++++++++++++++++++++++++++++ if ret['pos_anchors'] is not None: penalty_k = 0.055 tx, ty, tw, th = ret['template_target_xywh'].copy() tw *= ret['template_cropprd_resized_ratio'] th *= ret['template_cropprd_resized_ratio'] anchors = ret['anchors'].copy() w = anchors[:,2] * np.exp(reg_pred[:, 2].cpu().detach().numpy()) h = anchors[:,3] * np.exp(reg_pred[:, 3].cpu().detach().numpy()) eps = 1e-2 change_w = np.maximum(w/(tw+eps), tw/(w+eps)) change_h = np.maximum(h/(th+eps), th/(h+eps)) penalty = np.exp(-(change_w + change_h - 1) * penalty_k) pscore = score * penalty else: pscore = score # +++++++++++++++++++ v1.0 add window default cosine ++++++++++++++++++++++ score_size = 17 window_influence = 0.42 window = (np.outer(np.hanning(score_size), np.hanning(score_size)).reshape(17,17,1) + np.zeros((1, 1, 5))).reshape(-1) pscore = pscore * (1 - window_influence) + window * window_influence score_old = score score = pscore #from 0.2 - 0.7 # ++++++++++++++++++++ debug for class ++++++++++++++++++++++++++++++++++++ if example_index%1000 == 0: print(score[pos_index]) # this should tend to be 1 print(score[neg_index]) # this should tend to be 0 # ++++++++++++++++++++ debug for reg ++++++++++++++++++++++++++++++++++++++ tmp_dir = '/home/ly/chz/Siamese-RPN-pytorch/code_v1.0/tmp/visualization/7_check_train_phase_debug_pos_anchors' if not os.path.exists(tmp_dir): os.makedirs(tmp_dir) detection = ret['detection_cropped_resized'].copy() draw = ImageDraw.Draw(detection) pos_anchors = ret['pos_anchors'].copy() if ret['pos_anchors'] is not None else None if pos_anchors is not None: # draw pos anchors x = pos_anchors[:, 0] y = pos_anchors[:, 1] w = pos_anchors[:, 2] h = pos_anchors[:, 3] x1s, y1s, x2s, y2s = x - w//2, y - h//2, x + w//2, y + h//2 for i in range(16): x1, y1, x2, y2 = x1s[i], y1s[i], x2s[i], y2s[i] draw.line([(x1, y1), (x2, y1), (x2, y2), (x1, y2), (x1, y1)], width=1, fill='white') # pos anchor # pos anchor transform to red box after prediction x = pos_anchors[:,0] + pos_anchors[:, 2] * reg_pred[pos_index, 0].cpu().detach().numpy() y = pos_anchors[:,1] + pos_anchors[:, 3] * reg_pred[pos_index, 1].cpu().detach().numpy() w = pos_anchors[:,2] * np.exp(reg_pred[pos_index, 2].cpu().detach().numpy()) h = pos_anchors[:,3] * np.exp(reg_pred[pos_index, 3].cpu().detach().numpy()) x1s, y1s, x2s, y2s = x - w//2, y - h//2, x + w//2, y + h//2 for i in range(16): x1, y1, x2, y2 = x1s[i], y1s[i], x2s[i], y2s[i] draw.line([(x1, y1), (x2, y1), (x2, y2), (x1, y2), (x1, y1)], width=1, fill='red') # predict(white -> red) # pos anchor should be transformed to green gt box, if red and green is same, it is overfitting x = pos_anchors[:,0] + pos_anchors[:, 2] * reg_target[pos_index, 0].cpu().detach().numpy() y = pos_anchors[:,1] + pos_anchors[:, 3] * reg_target[pos_index, 1].cpu().detach().numpy() w = pos_anchors[:,2] * np.exp(reg_target[pos_index, 2].cpu().detach().numpy()) h = pos_anchors[:,3] * np.exp(reg_target[pos_index, 3].cpu().detach().numpy()) x1s, y1s, x2s, y2s = x - w//2, y-h//2, x + w//2, y + h//2 for i in range(16): x1, y1, x2, y2 = x1s[i], y1s[i], x2s[i], y2s[i] draw.line([(x1, y1), (x2, y1), (x2, y2), (x1, y2), (x1, y1)], width=1, fill='green') # gt (white -> green) x1, y1, x3, y3 = x1s[0], y1s[0], x2s[0], y2s[0] else: x1, y1, x3, y3 = 0, 0, 0, 0 # top1 proposal after nms (white) if example_index%1000 == 0: save_path = osp.join(tmp_dir, 'epoch_{:010d}_{:010d}_anchor_pred.jpg'.format(epoch, example)) detection.save(save_path) example_index = example_index+1 # +++++++++++++++++++ v1.0 restore ++++++++++++++++++++++++++++++++++++++++ ratio = ret['detection_cropped_resized_ratio'] detection_cropped = ret['detection_cropped'].copy() detection_cropped_resized = ret['detection_cropped_resized'].copy() original = Image.open(ret['detection_img_path']) x_, y_ = ret['detection_tlcords_of_original_image'] draw = ImageDraw.Draw(original) w, h = original.size """ un resized """ x1, y1, x3, y3 = x1/ratio, y1/ratio, y3/ratio, y3/ratio """ un cropped """ x1 = np.clip(x_ + x1, 0, w-1).astype(np.int32) # uncropped #target_of_original_img y1 = np.clip(y_ + y1, 0, h-1).astype(np.int32) x3 = np.clip(x_ + x3, 0, w-1).astype(np.int32) y3 = np.clip(y_ + y3, 0, h-1).astype(np.int32) draw.line([(x1, y1), (x3, y1), (x3, y3), (x1, y3), (x1, y1)], width=3, fill='yellow') #save_path = osp.join(tmp_dir, 'epoch_{:010d}_{:010d}_restore.jpg'.format(epoch, example)) #original.save(save_path) print("Epoch:{:04d}\texample:{:06d}/{:06d}({:.2f})%\tsteps:{:010d}\tlr:{:.7f}\tcloss:{:.4f}\trloss:{:.4f}\ttloss:{:.4f}".format(epoch, example+1, args.max_batches, 100*(example+1)/args.max_batches, steps, cur_lr, closses.avg, rlosses.avg, tlosses.avg )) if steps % 1 == 0: file_path = os.path.join(args.weight_dir, 'weights-{:07d}.pth.tar'.format(steps)) state = { 'epoch' :epoch+1, 'state_dict' :model.state_dict(), 'optimizer' : optimizer.state_dict(), } torch.save(state, file_path)