O2020/train_model.py

import os
import numpy as np
import pandas as pd
import random
from utils.dataset_config import dataset_config
import json
import sys
import yaml
import time
import gc
from tqdm.auto import tqdm

#Importing functions from utils
from utils.misc import AverageMeter, Logger, DETRModel, sample_metrics, seed_everything, get_data
import utils.misc as util

#Torch
import torch
from torch.utils.data import DataLoader
from torch.utils.tensorboard import SummaryWriter

#sklearn
from sklearn.model_selection import train_test_split

#CV
import cv2

#Importing dataset class
from utils.dataset import CustomDataset


'''
Funtion that is used to do one epoch of training of the model
'''
def train_one_epoch(data_loader, model, criterion, optimizer, device, max_norm, null_train_boxes, scheduler=None):

    start_epoch = time.time()

    model.train()
    criterion.train()

    tk0 = tqdm(data_loader, total=len(data_loader), leave=False)
    log = None

    for step, (sample, targets) in enumerate(tk0):

        start_time = time.time()

        batch_size = len(sample)
        sample = list(image.to(device) for image in sample)
        targets = [{k: v.to(device) for k, v in t.items()} for t in targets]

        optimizer.zero_grad()

        output = model(sample)
        loss_dict = criterion(output, targets)

        if log is None:
            log = {k:AverageMeter() for k in loss_dict}
            log['total_loss'] = AverageMeter()
            log['avg_iter_time'] = AverageMeter()
            for i in range(len(null_train_boxes)):
                log[f'avg_prec_{i}'] = AverageMeter(id=i)
                log[f'avg_rec_{i}'] = AverageMeter(id=i)
                log[f'avg_f1_{i}'] = AverageMeter(id=i)

        weight_dict = criterion.weight_dict
        total_loss = sum(loss_dict[k] * weight_dict[k] for k in loss_dict.keys() if k in weight_dict)
        total_loss.backward()

        if max_norm > 0:
            torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm)

        optimizer.step()

        if scheduler is not None:
            scheduler.step()

        for k,v in loss_dict.items():
            log[k].update(v.item(),batch_size)

        end_time = time.time()

        log['total_loss'].update(total_loss.item(),batch_size)
        log['avg_iter_time'].update(end_time-start_time)

        metrics = sample_metrics(output, targets)
        for i in range(len(null_train_boxes)):
            class_metric = metrics[:,i]
            log[f'avg_prec_{i}'].update(class_metric[0])
            log[f'avg_rec_{i}'].update(class_metric[1])
            log[f'avg_f1_{i}'].update(class_metric[2])

        tk0.set_postfix({k:v.avg(null_train_boxes) for k,v in log.items() if v.id is None})

    end_epoch = time.time()
    train_log = {'T/'+k:v.avg(null_train_boxes) for k,v in log.items()}
    train_log.update({'T/epoch_time': end_epoch-start_epoch})
    return train_log

'''
Funtion that is used to do one epoch of validation of the model
'''
def evaluate(data_loader, model, criterion, device, null_val_boxes):
    model.eval()
    criterion.eval()
    log = None

    with torch.no_grad():

        tk0 = tqdm(data_loader, total=len(data_loader),leave=False)
        weight_dict = criterion.weight_dict

        for step, (sample, targets) in enumerate(tk0):

            batch_size = len(sample)
            sample = list(image.to(device) for image in sample)
            targets = [{k: v.to(device) for k, v in t.items()} for t in targets]

            output = model(sample)
            loss_dict = criterion(output, targets)

            if log is None:
                log = {k:AverageMeter() for k in loss_dict}
                log['total_loss'] = AverageMeter()
                for i in range(len(null_val_boxes)):
                    log[f'avg_prec_{i}'] = AverageMeter(id=i)
                    log[f'avg_rec_{i}'] = AverageMeter(id=i)
                    log[f'avg_f1_{i}'] = AverageMeter(id=i)

            for k,v in loss_dict.items():
                log[k].update(v.item(),batch_size)

            total_loss = sum(loss_dict[k] * weight_dict[k] for k in loss_dict.keys() if k in weight_dict)

            log['total_loss'].update(total_loss.item(), batch_size)

            metrics = sample_metrics(output, targets)
            for i in range(len(null_val_boxes)):
                class_metric = metrics[:,i]
                log[f'avg_prec_{i}'].update(class_metric[0])
                log[f'avg_rec_{i}'].update(class_metric[1])
                log[f'avg_f1_{i}'].update(class_metric[2])

            tk0.set_postfix({k:v.avg(null_val_boxes) for k,v in log.items() if not v.id})

    valid_log = {'V/'+k:v.avg(null_val_boxes) for k,v in log.items()}
    return valid_log

'''
Main function that does the training of a DETR model with the parameters loaded
from the configs.yaml file. Both the parameters in that file and the directories
of the files needed, which are set in the dir.yaml file, can be changed.
'''
def run(parameters, directories, ids_dict):
    #Retrieving directories information
    csv_path = directories['csv_path']
    train_test_path = directories['train_test_path']
    train_dir = os.path.join(train_test_path, "train")
    val_dir = os.path.join(train_test_path, "val")
    model_file = directories['model_file']
    output_path = directories['output_path']

    #Retrieving parameters
    seed = parameters['seed']
    null_class_coef = parameters['null_class_coef']
    num_classes = parameters['num_classes']
    num_queries = parameters['num_queries']
    batch_size = parameters['batch_size']
    LR = parameters['LR']
    lr_dict = parameters['lr_dict']
    epochs = parameters['epochs']
    max_norm = parameters['max_norm']
    model_name = parameters['model_name']
    dev = parameters['dev']
    weight_dict = parameters['weight_dict']
    pre_trained = parameters['pre_trained']
    from_scratch = parameters['from_scratch']
    lr_drop = parameters['lr_drop']

    #Setting random processes seed to assure reproducibility
    seed_everything(seed)

    #Gathering information about training data in a pandas dataframe
    train_ids, val_ids, test_ids = ids_dict.values()

    train_csv = os.path.join(csv_path, 'train.csv')
    train_marking = pd.read_csv(train_csv)
    train_image_data = train_marking.groupby('image_id')

    val_csv = os.path.join(csv_path, 'val.csv')
    val_marking = pd.read_csv(val_csv)
    val_image_data = val_marking.groupby('image_id')

    #Retrieval of necessary information for the creation of the dataset and
    #split of previously called train set into train and validation sets
    train_info = [get_data(img_id, train_image_data) for img_id in train_ids]
    validation_info =[get_data(img_id, val_image_data) for img_id in val_ids]
    null_train_images = [x['image_id'] for x in train_info if len(x['boxes'])==0]
    null_val_images = [x['image_id'] for x in validation_info if len(x['boxes'])==0]

    print(f'Total number of images in train set: {len(train_info)}')
    print('Images with no boxes in train set: ', len(null_train_images))
    null_train_boxes = np.zeros(num_classes)
    for img_id in train_info:
        for i in range(len(img_id['boxes'])):
            null_train_boxes[img_id['labels'][i]] += 1
    print(f'Total number of images in validation set: {len(validation_info)}')
    print('Images with no boxes in validation set: ', len(null_val_images))
    null_val_boxes = np.zeros(num_classes)
    for img_id in validation_info:
        for i in range(len(img_id['boxes'])):
            null_val_boxes[img_id['labels'][i]] += 1

    #Initialization of both the train and validation sets
    train_ds = CustomDataset(train_info, train_dir, image_set = 'train')
    valid_ds = CustomDataset(validation_info, val_dir, image_set = 'val')

    #Initialization of tensorboard's SummaryWriter to keep track of the
    #evolution of the training phase of the model
    tb = SummaryWriter()

    gc.collect()

    #Creation DataLoader for the train and validations sets, which are used
    #to retrieve the batches during the training phase
    train_data_loader = DataLoader(
        train_ds,
        batch_size=batch_size,
        shuffle=True,
        num_workers=2,
        collate_fn=util.collate_fn
        )
    valid_data_loader = DataLoader(
        valid_ds,
        batch_size=batch_size,
        shuffle=True,
        num_workers=2,
        collate_fn=util.collate_fn
        )

    #Device where tensors and model will be stored
    device = torch.device(dev)

    #Initialization of DETR
    if pre_trained or from_scratch:

        #Adding path to DETR's repository
        detr_path = directories['detr_path']
        sys.path.extend([detr_path])

        #Importing classes that are required from DETR's repository
        from models.detr import DETR,SetCriterion
        from models.backbone import Backbone,Joiner
        from models.position_encoding import PositionEmbeddingSine
        from models.transformer import Transformer

        hidden_dim = 256
        backbone = Backbone(model_name, train_backbone=True, return_interm_layers=False, dilation=True)
        pos_enc = PositionEmbeddingSine(hidden_dim // 2, normalize=True)
        backbone_with_pos_enc = Joiner(backbone, pos_enc)
        backbone_with_pos_enc.num_channels = backbone.num_channels
        transformer = Transformer(d_model=hidden_dim, return_intermediate_dec=True)
        model = DETR(backbone_with_pos_enc, transformer, num_classes=num_classes, num_queries=num_queries, aux_loss=False)

        if pre_trained:
            try:
                model.load_state_dict(torch.load(model_file))
                print(f'Model weights loaded from {model_file}')
            except:
                model_dir_info = model_file.split('.')
                model_dir = model_dir_info[0]
                filename = model_dir_info[1]
                sys.exit(f'File {model_dir} does not contain file {filename}.')
        else:
            print('Model initiated without pre-trained weights.')
    else:
        model = DETRModel(model_name=model_name, num_classes=num_classes)
        print('Model weights loaded from https://github.com/facebookresearch/detr.')

    model.to(device)

    '''
    Tried to add graph of model in tensorboard but did not work so far

    (images, targets, image_ids) = next(iter(train_data_loader))
    tb.add_graph(model, images[0])
    '''

    #Creation of the loss function
    cost_bbox = weight_dict['loss_bbox']
    cost_giou = weight_dict['loss_giou']
    cost_class = weight_dict['loss_ce']

    from models.matcher import HungarianMatcher
    from models.detr import SetCriterion

    matcher = HungarianMatcher(cost_class=cost_class, cost_bbox=cost_bbox, cost_giou=cost_giou)
    losses = ['labels', 'boxes', 'cardinality']

    criterion = SetCriterion(num_classes,
                             matcher, weight_dict,
                             eos_coef=null_class_coef,
                             losses=losses)
    criterion = criterion.to(device)

    #Selecting different learning rates for each part of DETR and initialization
    #of the optimizer
    param_dic = { 'backbone': [p for n,p in model.named_parameters()
                              if ('backbone' in n) and p.requires_grad],
                 'transformer': [p for n,p in model.named_parameters()
                                 if (('transformer' in n) or ('input_proj' in n)) and p.requires_grad],
                 'embed': [p for n,p in model.named_parameters()
                                 if (('class_embed' in n) or ('bbox_embed' in n) or ('query_embed' in n))
                           and p.requires_grad]}

    #Setting optimizer, which is the method for doing the backpropagation os loss values
    optimizer = torch.optim.AdamW([{'params': v, 'lr': lr_dict.get(k,1)*LR} for k,v in param_dic.items()], lr=LR, weight_decay=1e-4)
    #Setting scheduler, which changes the learning rate of the optimizer after lr_drop iterations
    if lr_drop is not "None":
        scheduler = torch.optim.lr_scheduler.StepLR(optimizer, lr_drop)
    else:
        scheduler = None

    #Setting best total loss to infinity to later save the model that minimizes
    #this value
    logger = util.Logger('train_metrics')
    best_total_loss = float('inf')
    header_printed = True

    #Beginning of the training process
    for epoch in range(epochs):

        train_log = train_one_epoch(train_data_loader, model, criterion, optimizer, device, max_norm, null_train_boxes)
        valid_log = evaluate(valid_data_loader, model, criterion, device, null_val_boxes)

        complete_log = {}
        complete_log.update(train_log)
        complete_log.update(valid_log)
        logger.save(complete_log, epoch+1)
        keys = sorted(complete_log.keys())

        if header_printed:
            print(' '.join(map(lambda k: f'{k[:15]:8}', keys)))
            header_printed = False
        print(' '.join(map(lambda k: f'{complete_log[k]:8.3f}'[:11],keys)))

        #Logging information to tensorboard
        for k,v in complete_log.items():
            tb.add_scalar(k,v,epoch+1)

        if complete_log['V/total_loss'] < best_total_loss:
            best_total_loss = complete_log['V/total_loss']
            print('Best model found at epoch {}'.format(epoch+1))
            torch.save(model.state_dict(), os.path.join(output_path, 'best_model.pth'))
        torch.save(model.state_dict(), 'checkpoint_model.pth')

        if epoch == epochs-1:
            torch.save(model.state_dict(), os.path.join(output_path, 'final_model.pth'))

if __name__ == "__main__":
    #Attempting to read dir.yaml file that contains directories information
    try:
        with open(os.path.join("configs","dir.yaml"), "r") as fp:
            directories = yaml.load(fp, Loader=yaml.FullLoader)
    except:
        sys.exit('File dir.yaml does not exist in the current folder')

    #Attempting to read parameters.yaml file that contains parameters information
    try:
        with open(os.path.join("configs","parameters.yaml"), "r") as fp:
            parameters = yaml.load(fp, Loader=yaml.FullLoader)
    except:
        sys.exit('File parameters.yaml does not exist in the current folder')

    #Attempting to read id_dict.json file that contains sets' images ids
    try:
        with open("ids_dict.json", "r") as fp:
            ids_dict = json.load(fp)
    except:
        sys.exit('File ids_dict.json does not exist in the current folder')

    run(parameters, directories, ids_dict)