train_knn_pointnetvlad.py

#!/usr/bin/env python3
# This file is covered by the LICENSE file in the root of this project.



'''
Version:
 - train pointNetVlad model on a retrieval-based task  using kitti
 
'''
import argparse
import subprocess
import datetime
import yaml
from shutil import copyfile
import os
import shutil
import tqdm
from sklearn.metrics import accuracy_score, precision_score, recall_score, f1_score
import signal, sys
from sklearn.neighbors import NearestNeighbors
from torch import optim
import torch 
import numpy as np

from networks.pointnetvlad import PointNetVlad

from inference import eval_net
from dataset_utils.kitti import pcl_parser_test as test_pars
from dataset_utils.kitti import pcl_parser_train as train_pars

from utils.session_plot_utils import pose_plots,metric_plots,loss_plots,distribution_plots

from datetime import datetime

import random
from torch.utils.data import DataLoader, random_split
from utils.utils import dump_info


if __name__ == '__main__':
  parser = argparse.ArgumentParser("./infer.py")
  parser.add_argument(
      '--dataset', '-d',
      type=str,
      default = "kitti",
      required=False,
      help='Dataset to train with. No Default',
  )
  parser.add_argument(
      '--corr', '-c',
      type=str,
      default= ''
  )

  parser.add_argument(
      '--model', '-m',
      type=str,
      required=False,
      default='pointnetvlad',
      help='Directory to get the trained model.'
  )

  parser.add_argument(
      '--sess_cfg', '-f',
      type=str,
      required=False,
      #default='cross_val_00',
      default='cosine_small_session',
      help='Directory to get the trained model.'
  )

  parser.add_argument(
      '--pretrained', '-p',
      type=str,
      required=False,
      default="darknet53-512",
      help='Directory to get the trained model.'
  )

  parser.add_argument(
      '--debug', '-b',
      type=int,
      required=False,
      default=False,
      help='Directory to get the trained model.'
  )

  parser.add_argument(
      '--plot',
      type=int,
      required=False,
      default=1,
      help='Directory to get the trained model.'
  )

  parser.add_argument(
      '--results',
      type=str,
      required=False,
      default='session_results.txt',
      help='Directory to get the trained model.'
  )

  FLAGS, unparsed = parser.parse_known_args()

  # print summary of what we will do
  print("----------")
  print("INTERFACE:")
  print("Dataset:", FLAGS.dataset)
  print("Sequence: ", FLAGS.corr)
  print("Model: ", FLAGS.model)
  print("Debug flag: ", FLAGS.debug)
  print("Pretrained flag: ", FLAGS.pretrained)
  print("----------\n")

  # open arch config file
  cfg_file = os.path.join('dataset_utils',FLAGS.dataset,'data_cfg_pointnetvlad.yaml')
  try:
    print("Opening data config file: %s" % cfg_file)
    DATA = yaml.safe_load(open(cfg_file , 'r'))
  except Exception as e:
    print(e)
    print("Error opening data yaml file.")
    quit()

  session_cfg_file = os.path.join('sessions', FLAGS.sess_cfg + '.yaml')
  try:
    print("Opening session config file: %s" % session_cfg_file)
    SESSION = yaml.safe_load(open(session_cfg_file, 'r'))
  except Exception as e:
    print(e)
    print("Error opening arch yaml file.")
    quit()
  
  model_cfg_file = os.path.join('model_cfg', FLAGS.model + '.yaml')
  try:
    print("Opening model config file: %s" % model_cfg_file)
    ARCH = yaml.safe_load(open(model_cfg_file, 'r'))
  except Exception as e:
    print(e)
    print("Error opening arch yaml file.")
    quit()

  
  ###################################################################### 
  
  debug_flag   = FLAGS.debug
  dataset_name = FLAGS.dataset
  
  
  data_root_path = DATA['dataset']['path']['root']
  
  training_setup = SESSION['train']
  
  print("[WARN] Loading training parm from yam file")

  # Loading dataset 
  train_dataset = train_pars.Parser(
                                  dataset = DATA["dataset"],
                                  session = SESSION['train'])
  test_dataset = test_pars.Parser(
                                  dataset = DATA["dataset"],
                                  session = SESSION['test'])

  train_loader = train_dataset.get_set()
  val_loader   = test_dataset.get_set()
  triplet_idx  = test_dataset.get_triplets() 
  
  # Loading the network and pretrained weights
  model = PointNetVlad(global_feat=True, feature_transform=True, max_pool=False,
                                    output_dim=256, num_points=ARCH['max_points'])
  model_name = "PointNetVlad"
  
  # Load Pretrained weights  
  pretrained_root = SESSION['pretrained_root']
  pretrained = os.path.join(pretrained_root,FLAGS.pretrained + '.pth')
  if os.path.isfile(pretrained):
    pretrained_to_load = pretrained

  else: 
    # If no pretrained model is given, try to load a pretrained model with the session name
    session_name  = FLAGS.sess_cfg
    pretrained_to_load = os.path.join(pretrained_root,model_name + '_' + session_name + '.pth')
  try:
    # Verify if the selected pretrained exists 
    if os.path.isfile(pretrained_to_load):
      model.load_state_dict(torch.load(pretrained_to_load))
      print("[INF] Pretrained model was loaded: " + pretrained_to_load)
    else: 
      print("[INF] No pretrained model loaded!")
  except:
      print("[WRN] Something went wrong while loading the pretrained model!")


  # Device configuration
  device = 'cpu'
  if torch.cuda.is_available():
    device = 'cuda:1'
    torch.cuda.empty_cache()

  model.to(device)
  model.train()
  
  # Loss function
  loss_cf = SESSION['loss_function']
  print("\n----------")
  print("Loss Function Param:")
  print("Margin: ", loss_cf['margin'])
  print("Reduction: ", loss_cf['reduction'])

  criterion = torch.nn.CosineEmbeddingLoss(margin  = loss_cf['margin'],reduction= loss_cf['reduction'])

  print("\n----------")
  print("Data Param:")
  print("Root:              ", data_root_path)
  print("Batch Size:        ", training_setup['batch_size'])
  #print("Dataset Fraction:  ", fraction)
  
  # Optimizer parameterization
  lr            = ARCH['optimizer']['lr']
  w_decay       = ARCH['optimizer']['w_decay']
  amsgrad       = ARCH['optimizer']['amsgrad']
  epsilon       = ARCH['optimizer']['epsilon_w']
  betas         = tuple(ARCH['optimizer']['betas'])

  print("\n----------")
  print("Optimization Param:")
  print("Learning rate: ", lr)
  print("Weight decay:  ",w_decay)
  print("amsgrad:       ",amsgrad)
  print("epsilon:       ", epsilon)
  print("Betas:         ", betas)
  print("----------\n")

  optimizer = optim.Adam(
              model.parameters(), 
              lr = lr,
              weight_decay=w_decay,
              eps = epsilon,
              amsgrad = amsgrad,
              betas = betas
              )


  if FLAGS.plot == 1:
    # Create plot instances 
    plot_loss = loss_plots('Loss')

    plot_train_distro = distribution_plots([0,1],'train')
    plot_val_distro   = distribution_plots([0,1],'val')
    plot_val_metrics  = metric_plots('val metrics')

    plot_pose = pose_plots('map')
    plot_pose.update(ref = triplet_idx['poses'])

  # Create a global dictionary with all important variables 
  global_val_score = {'f1':-1,
                      'r':-1,
                      'p':-1,
                      'a':-1,
                      'loss':2,
                      'fps':-1,
                      'epoch':-1 }
  mean_fps_array= []
  
  # Root path to store models weights
 
  if not os.path.isdir(pretrained_root):
    # If it does not exist then create the folder 
    os.makedirs(pretrained_root)

  # Training parameters
  epochs          = training_setup['max_epochs']
  VAL_EPOCH       = training_setup['report_val']
  eval_criterion  = torch.nn.CosineSimilarity(dim=1)

  print("[INF] Loaded Model: " + model_name)
  print("[INF] Device: " + device)
  print("[INF] Result File: " + FLAGS.results)
  print("[INF] Session: " + SESSION['name'])
  print("[INF] Retrieval top candidates: %d"%(SESSION['retrieval']['top_cand']))
  print("[INF] Retrieval range threshold: %d"%(SESSION['retrieval']['range_thres']))

  try:
    for epoch in range(epochs):
          
      running_loss  = 0
      gt_true = np.array([])
      sim_bag = np.array([])
      loss_bag = np.array([])
      
      sub_epoch = epoch
      xx = np.array([])
      itr = 1/len(train_loader)
      
      for projA,projB, gt_label in tqdm.tqdm(train_loader):
        
        projA = projA.to(device)
        projB = projB.to(device)

        gt_label = gt_label.to(device).view(-1,1)
        
        optimizer.zero_grad()

        # compute output
        fa = model(projA)
        fb = model(projB)

        loss = criterion(fa,fb,gt_label)
          
        loss.backward()
        optimizer.step()
        
        scores = eval_criterion(fa,fb).detach().cpu().numpy()
        
        labels = gt_label.detach().cpu()
        loss = loss.detach().item()
        running_loss += loss
        
        xx = np.append(xx,round(sub_epoch,3))
        sub_epoch += itr
        gt_true = np.append(gt_true,labels)
        sim_bag = np.append(sim_bag,scores)
        #loss_bag = np.append(loss_bag,loss)

      train_loss = running_loss/len(train_loader)
      gt_true[gt_true==-1] = 0
      
      if FLAGS.plot == 1:
        plot_loss.update(data='mean',scores=train_loss,x=epoch+1)
      
      print("train epoch : {}/{}, loss: {:.6f}".format(epoch, epochs, train_loss))
      
      # Plotting stuff

      weigths_file = '%s_%s.pth' % (model_name,FLAGS.sess_cfg)
      trained_weights = os.path.join(pretrained_root,weigths_file)
      torch.save(model.state_dict(), trained_weights)

      if epoch % VAL_EPOCH == 0:
        

        metric,val_poses,text = eval_net( model,
                                          test_dataset, 
                                          device,
                                          top_candid = SESSION['retrieval']['top_cand'],
                                          range_thres = SESSION['retrieval']['range_thres']
                                          ) 
        
        mean_fps_array.append(metric['fps'])
        
        if FLAGS.plot == 1:
          plot_val_metrics.update(epoch = epoch,f1 = metric['f1'],acc = metric['a'])
          #plot_val_distro.update( loops['labels'],loops['scores'])
          plot_pose.update(query = val_poses['query'],
                    tn =  val_poses['tn'],
                    tp = val_poses['tp'],
                    #fp  = val_poses['fp'],
                    fn =  val_poses['fn'])
      
        if metric['f1'] > global_val_score['f1']: 
          
          weigths_file = '%s_%s_best.pth' % (model_name,SESSION['name'])
          trained_weights = os.path.join(pretrained_root,weigths_file)
          torch.save(model.state_dict(), trained_weights)
          # Overwite
          global_val_score  = metric
          global_val_score['loss']  = train_loss
          global_val_score['epoch'] = epoch

          print("[INF] weights stored at: " + weigths_file)
  
  except KeyboardInterrupt:
    print("[INF] CTR + C")
  
  except ValueError:
    print("[INF] Error: {} ".format(ValueError))

  root = 'results/' + model_name + '_' + SESSION['name']
  
  if FLAGS.plot == 1:
    plot_loss.save_data_file(root) 
    plot_val_metrics.save_data_file(root)
  
  text_to_store = {}
  text_to_store['model'] = model_name
  text_to_store['session'] =   SESSION['name']
  text_to_store['F1'] =  round(global_val_score['f1'],3)
  text_to_store['R'] =   round(global_val_score['r'],3)
  text_to_store['P'] =   round(global_val_score['p'],3)
  text_to_store['A'] =   round(global_val_score['a'],3)
  text_to_store['epoch'] =  "%d/%d"%(global_val_score['epoch'],epochs)
  # text_to_store['param'] = model.get_parm_size()
  text_to_store['FPS'] =   round(np.mean(mean_fps_array),1)

  output_txt = dump_info( FLAGS.results, text_to_store,'a')
  print("[INF] " + output_txt)