train_RFN.py

# Copyright (c) SenseTime. All Rights Reserved.

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
from __future__ import unicode_literals

import argparse
import logging
import os
import time
import math
import json
import random
import numpy as np

os.environ["CUDA_VISIBLE_DEVICES"] = "1,2"

import torch
import torch.nn as nn
from torch.utils.data import DataLoader
from tensorboardX import SummaryWriter
from torch.nn.utils import clip_grad_norm_
from torch.utils.data.distributed import DistributedSampler

from DFAT.utils.lr_scheduler import build_lr_scheduler
from DFAT.utils.log_helper import init_log, print_speed, add_file_handler
from DFAT.utils.distributed import dist_init, DistModule, reduce_gradients,\
        average_reduce, get_rank, get_world_size
from DFAT.utils.model_load import load_pretrain, restore_from
from DFAT.utils.average_meter import AverageMeter
from DFAT.utils.misc import describe, commit
from DFAT.models.model_builder import ModelBuilder
from DFAT.datasets.dataset_RFN import TrkDataset
from DFAT.core.config import cfg
import pdb


logger = logging.getLogger('global')
parser = argparse.ArgumentParser(description='siamrpn tracking')
#/data/Disk_B/zhangyong/DFAT   --> ..
parser.add_argument('--cfg', type=str, default='./experiments/siam_base/config.yaml',
                    help='configuration of tracking')
parser.add_argument('--seed', type=int, default=123456,
                    help='random seed')
parser.add_argument('--local_rank', type=int, default=0,
                    help='compulsory for pytorch launcer')
args = parser.parse_args()


def seed_torch(seed=0):
    random.seed(seed)
    os.environ['PYTHONHASHSEED'] = str(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed(seed)
    torch.backends.cudnn.benchmark = False    #spend some time to find the most properly inplementation of conv for the network and be faster
    torch.backends.cudnn.deterministic = True #every time the conv algorithm returned is settled


def build_data_loader():
    logger.info("build train dataset")
    # train_dataset
    train_dataset = TrkDataset()
    logger.info("build dataset done")

    train_sampler = None
    # if get_world_size() > 1:
    #     train_sampler = DistributedSampler(train_dataset)
    train_loader = DataLoader(train_dataset,
                              batch_size=cfg.TRAIN.BATCH_SIZE,
                              num_workers=cfg.TRAIN.NUM_WORKERS,
                              pin_memory=True,
                              sampler=train_sampler)
    return train_loader


def build_opt_lr(model, current_epoch=0):
    for param in model.backbone.parameters():
        param.requires_grad = False
    for m in model.backbone.modules():
        if isinstance(m, nn.BatchNorm2d):
            m.eval()
    for param in model.neck.parameters():
        param.requires_grad = False
    for m in model.neck.modules():
        if isinstance(m, nn.BatchNorm2d):
            m.eval()
        #############################################
    for param in model.RFN.parameters():
        param.requires_grad = True
    for m in model.RFN.modules():
        if isinstance(m, nn.BatchNorm2d):
            m.train()
        #############################################
    for param in model.rpn_head.parameters():
        param.requires_grad = False
    for m in model.rpn_head.modules():
        if isinstance(m, nn.BatchNorm2d):
            m.eval()

        # set RFN trainable
    trainable_params = [{'params': filter(lambda x: x.requires_grad, model.RFN.parameters()),
                         'lr': cfg.TRAIN.BASE_LR}]
    trainable_params += [{'params': filter(lambda x: x.requires_grad, model.rpn_head.parameters()),
                         'lr': cfg.TRAIN.BASE_LR}]
    # + \
    # [{'params': filter(lambda x: x.requires_grad, model.rpn_head.parameters()),
    # 'lr': cfg.TRAIN.BASE_LR}]

    if current_epoch >= cfg.BACKBONE.TRAIN_EPOCH:
        for layer in cfg.BACKBONE.TRAIN_LAYERS:
            for param in getattr(model.backbone, layer).parameters():
                param.requires_grad = True
            for m in getattr(model.backbone, layer).modules():
                if isinstance(m, nn.BatchNorm2d):
                    m.train()
        for param in model.neck.parameters():
            param.requires_grad = True
        for m in model.neck.modules():
            if isinstance(m, nn.BatchNorm2d):
                m.train()

        trainable_params += [{'params': list(filter(lambda x: x.requires_grad, model.neck.parameters())),
                              'lr': cfg.TRAIN.BASE_LR}]

    optimizer = torch.optim.SGD(trainable_params, cfg.TRAIN.BASE_LR,
                                momentum=cfg.TRAIN.MOMENTUM,
                                weight_decay=cfg.TRAIN.WEIGHT_DECAY)

    lr_scheduler = build_lr_scheduler(optimizer, epochs=cfg.TRAIN.EPOCH)
    lr_scheduler.step(cfg.TRAIN.START_EPOCH)
    return optimizer, lr_scheduler


def log_grads(model, tb_writer, tb_index):
    def weights_grads(model):
        grad = {}
        weights = {}
        for name, param in model.named_parameters():
            if param.grad is not None:
                grad[name] = param.grad
                weights[name] = param.data
        return grad, weights

    grad, weights = weights_grads(model)
    feature_norm, rpn_norm = 0, 0
    for k, g in grad.items():
        _norm = g.data.norm(2)
        weight = weights[k]
        w_norm = weight.norm(2)
        if 'feature' in k:
            feature_norm += _norm ** 2
        else:
            rpn_norm += _norm ** 2

        tb_writer.add_scalar('grad_all/'+k.replace('.', '/'),
                             _norm, tb_index)
        tb_writer.add_scalar('weight_all/'+k.replace('.', '/'),
                             w_norm, tb_index)
        tb_writer.add_scalar('w-g/'+k.replace('.', '/'),
                             w_norm/(1e-20 + _norm), tb_index)
    tot_norm = feature_norm + rpn_norm
    tot_norm = tot_norm ** 0.5
    feature_norm = feature_norm ** 0.5
    rpn_norm = rpn_norm ** 0.5

    tb_writer.add_scalar('grad/tot', tot_norm, tb_index)
    tb_writer.add_scalar('grad/feature', feature_norm, tb_index)
    tb_writer.add_scalar('grad/rpn', rpn_norm, tb_index)

def BNtoFixed(m):
    class_name = m.__class__.__name__
    if class_name.find('BatchNorm') != -1:
        m.eval()

def train(train_loader, model, optimizer, lr_scheduler, tb_writer):
    cur_lr = lr_scheduler.get_cur_lr()
    # rank = get_rank()

    average_meter = AverageMeter()
    model.train()
    model.module.backbone.eval()
    model.module.neck.eval()
    ################################
    model.module.RFN.train()
    ################################
    # model.module.rpn_head.eval()
    model.module.backbone.apply(BNtoFixed)
    # model.module.rpn_head.apply(BNtoFixed)
    model = model.cuda()

    def is_valid_number(x):
        return not(math.isnan(x) or math.isinf(x) or x > 1e4)

    # world_size = get_world_size()
    num_per_epoch = len(train_loader.dataset) // \
        cfg.TRAIN.EPOCH // cfg.TRAIN.BATCH_SIZE
    start_epoch = cfg.TRAIN.START_EPOCH
    epoch = start_epoch

    if not os.path.exists(cfg.TRAIN.SNAPSHOT_DIR):
        os.makedirs(cfg.TRAIN.SNAPSHOT_DIR)

    print('******')
    # print the para and its name
    for name, param in model.named_parameters():
        if param.requires_grad:
            print(name)
    print('******')

    # logger.info("model\n{}".format(describe(model.module)))
    end = time.time()


    for idx, data in enumerate(train_loader):
        # data['anchor_iou'] = anchor_iou
        if epoch != idx // num_per_epoch + start_epoch:
            epoch = idx // num_per_epoch + start_epoch

            torch.save(
                    {'epoch': epoch,
                    'state_dict': model.module.state_dict(),
                    'optimizer': optimizer.state_dict()},
                    cfg.TRAIN.SNAPSHOT_DIR+'/checkpoint_e%d.pth' % (epoch))

            if epoch == cfg.TRAIN.EPOCH:
                return

            if cfg.BACKBONE.TRAIN_EPOCH == epoch:
                logger.info('start training backbone.')
                optimizer, lr_scheduler = build_opt_lr(model.module, epoch)

                # print('******')
                # # print the para and its name
                # for name, param in model.named_parameters():
                #     if param.requires_grad:
                #         print(name)
                # print('******')

                # logger.info("model\n{}".format(describe(model.module)))

            lr_scheduler.step(epoch)
            cur_lr = lr_scheduler.get_cur_lr()
            logger.info('epoch: {}'.format(epoch+1))

        # logger.info('start training backbone.')
        # optimizer, lr_scheduler = build_opt_lr(model.module, epoch)
        # # logger.info("model\n{}".format(describe(model.module)))
        # lr_scheduler.step(epoch)
        # cur_lr = lr_scheduler.get_cur_lr()

        tb_idx = idx
        if idx % num_per_epoch == 0 and idx != 0:
            for idx, pg in enumerate(optimizer.param_groups):
                logger.info('epoch {} lr {}'.format(epoch+1, pg['lr']))
                tb_writer.add_scalar('lr/group{}'.format(idx+1),
                                    pg['lr'], tb_idx)

        data_time = time.time() - end
        tb_writer.add_scalar('time/data', data_time, tb_idx)

        outputs = model(data)
        # loss = outputs['total_loss']

        loss = torch.mean(outputs['total_loss'])

        if is_valid_number(loss.data.item()):
            optimizer.zero_grad()
            loss.backward()
            # reduce_gradients(model)

            if cfg.TRAIN.LOG_GRADS:
                log_grads(model.module, tb_writer, tb_idx)

            # clip gradient
            clip_grad_norm_(model.parameters(), cfg.TRAIN.GRAD_CLIP)
            optimizer.step()

        batch_time = time.time() - end
        batch_info = {}
        batch_info['batch_time'] = batch_time
        batch_info['data_time'] = data_time
        for k, v in sorted(outputs.items()):
            batch_info[k] = v.mean().data.item()


        average_meter.update(**batch_info)

        for k, v in batch_info.items():
            tb_writer.add_scalar(k, v, tb_idx)

        if (idx+1) % cfg.TRAIN.PRINT_FREQ == 0:
            #check the weight for three RPN blocks
            # weight_rpn = str(outputs['cls_w']) + '+' + str(outputs['loc_w']) + '\n'
            # logger.info(weight_rpn)
            # if cfg.FUSION_pred_cur.TYPE == "weights":
            #     weight_old_new = str(outputs['balance_cls']) + '+' + str(outputs['balance_loc']) + '\n'
            #     logger.info(weight_old_new)


            info = "Epoch: [{}][{}/{}] lr: {:.6f}\n".format(
                        epoch+1, (idx+1) % num_per_epoch,
                        num_per_epoch, cur_lr)
            for cc, (k, v) in enumerate(batch_info.items()):
                if cc % 2 == 0:
                    info += ("\t{:s}\t").format(
                            getattr(average_meter, k))
                else:
                    info += ("{:s}\n").format(
                            getattr(average_meter, k))
            logger.info(info)
            print_speed(idx+1+start_epoch*num_per_epoch,
                        average_meter.batch_time.avg,
                        cfg.TRAIN.EPOCH * num_per_epoch)
        end = time.time()


def main():
    # rank, world_size = dist_init()
    logger.info("init done")

    # load cfg
    cfg.merge_from_file(args.cfg)

    if not os.path.exists(cfg.TRAIN.LOG_DIR):
        os.makedirs(cfg.TRAIN.LOG_DIR)
    init_log('global', logging.INFO)
    if cfg.TRAIN.LOG_DIR:
        add_file_handler('global',
            os.path.join(cfg.TRAIN.LOG_DIR, cfg.TRAIN.LOGFILE),
            logging.INFO)

    logger.info("Version Information: \n{}\n".format(commit()))
    logger.info("config \n{}".format(json.dumps(cfg, indent=4)))

    # build dataset loader
    train_loader = build_data_loader()

    # anchor_iou = ioum(torch.from_numpy(train_loader.dataset.anchor_target.anchors.all_anchors[0]).cuda())
    # anchor_iou_index = high_iou_index(anchor_iou)
    # create model
    # model = ModelBuilder(anchor_iou_index).cuda()
    model = ModelBuilder().cuda()
    # load pretrained backbone weights
    if cfg.BACKBONE.PRETRAINED:
        cur_path = os.path.dirname(os.path.realpath(__file__))#__file__ is a built-in para
        backbone_path = os.path.join(cur_path, cfg.BACKBONE.PRETRAINED)
        load_pretrain(model.backbone, backbone_path)
        # load_pretrain(model.backbone_tir, backbone_path)

    # dist_model = nn.DataParallel(model, list(range(torch.cuda.device_count()))).cuda()   #choose gpu from available setted before
    used = [0, 1]
    dist_model = nn.DataParallel(model, used).cuda()

    # create tensorboard writer
    tb_writer = SummaryWriter(cfg.TRAIN.LOG_DIR)


    # build optimizer and lr_scheduler
    optimizer, lr_scheduler = build_opt_lr(model,
                                           cfg.TRAIN.START_EPOCH)

    # resume training
    if cfg.TRAIN.RESUME:
        logger.info("resume from {}".format(cfg.TRAIN.RESUME))
        assert os.path.isfile(cfg.TRAIN.RESUME), \
            '{} is not a valid file.'.format(cfg.TRAIN.RESUME)
        model, optimizer, cfg.TRAIN.START_EPOCH = \
            restore_from(model, optimizer, cfg.TRAIN.RESUME)

        optimizer, lr_scheduler = build_opt_lr(model,
                                               cfg.TRAIN.START_EPOCH)

    # load pretrain
    elif cfg.TRAIN.PRETRAINED:
        load_pretrain(model, cfg.TRAIN.PRETRAINED)
    # dist_model = DistModule(model)

    logger.info(lr_scheduler)
    logger.info("model prepare done")


    # start training
    train(train_loader, dist_model, optimizer, lr_scheduler, tb_writer)


if __name__ == '__main__':
    seed_torch(args.seed)
    main()