train.py

from __future__ import print_function
import sys
import os
import torch
import torch.nn as nn
import torch.optim as optim
import torch.backends.cudnn as cudnn
import torchvision.transforms as transforms
import torch.nn.init as init
import argparse
import numpy as np
from torch.autograd import Variable
import torch.utils.data as data
from data import *
from layers.modules import MultiBoxLoss
from layers.functions import PriorBox
import time

parser = argparse.ArgumentParser(
    description='Receptive Field Block Net Training')
parser.add_argument('-v', '--version', default='RFB_vgg',
                    help='Choose the version.')
parser.add_argument('-s', '--size', default='300',
                    help='300 or 512 input size.')
parser.add_argument('-d', '--dataset', default='VOC',
                    help='VOC, COCO or CUSTOM_VOC dataset')
parser.add_argument('--basenet', default=None, help='pretrained base model')
parser.add_argument('--jaccard_threshold', default=0.5,
                    type=float, help='Min Jaccard index for matching')
parser.add_argument('-b', '--batch_size', default=32,
                    type=int, help='Batch size for training')
parser.add_argument('--num_workers', default=8,
                    type=int, help='Number of workers used in dataloading')
parser.add_argument('--cuda', default=True,
                    type=bool, help='Use cuda to train model')
parser.add_argument('--ngpu', default=1, type=int, help='gpus')
parser.add_argument('--lr', '--learning-rate',
                    default=2e-3, type=float, help='initial learning rate')
parser.add_argument('--momentum', default=0.9, type=float, help='momentum')
parser.add_argument(
    '--resume_net', default=None, help='resume net for retraining')
parser.add_argument('--resume_epoch', default=0,
                    type=int, help='resume iter for retraining')
parser.add_argument('-max','--max_epoch', default=150,
                    type=int, help='max epoch for retraining')
parser.add_argument('--weight_decay', default=5e-4,
                    type=float, help='Weight decay for SGD')
parser.add_argument('--gamma', default=0.1,
                    type=float, help='Gamma update for SGD')
parser.add_argument('--log_iters', default=True,
                    type=bool, help='Print the loss at each iteration')
parser.add_argument('--save_folder', default='./weights/',
                    help='Location to save checkpoint models')
parser.add_argument('--summary_writer', default=None,
                    help='summary writing folder location')
args = parser.parse_args()

if not os.path.exists(args.save_folder):
    os.mkdir(args.save_folder)

if not args.summary_writer == None:
    from tensorboardX import SummaryWriter
    summary = SummaryWriter(args.summary_writer)
else:
    from tensorboardX import SummaryWriter
    summary = SummaryWriter()

if args.dataset == 'VOC':
    train_sets = [('2007', 'trainval'), ('2012', 'trainval')]
    cfg = (VOC_300, VOC_512)[args.size == '512']
    num_classes = 21
elif args.dataset == 'COCO':
    train_sets = [('2014', 'train'),('2014', 'valminusminival')]
    cfg = (COCO_300, COCO_512)[args.size == '512']
    num_classes = 81
else:
    train_sets = 'Customset'
    cfg = (COCO_300, COCO_512)[args.size == '512']
    num_classes = 2

# Version checking
if args.version == 'RFB_vgg':
    from models.RFB_Net_vgg import build_net
elif args.version == 'RFB_E_vgg':
    from models.RFB_Net_E_vgg import build_net
elif args.version == 'RFB_mobile':
    from models.RFB_Net_mobile import build_net
    cfg = mobile_300
elif args.version == 'DRFB_mobile':
    from models.DRFB_Net_mobile import build_net
    cfg = mobile_300
elif args.version == 'SSD_vgg':
    from models.SSD_vgg import build_net
    cfg = (VOC_SSDVGG_300, COCO_SSDVGG_300)[args.dataset == 'COCO']
elif args.version == 'SSD_mobile':
    from models.SSD_lite_mobilenet_v1 import build_net
    cfg = mobile_300
#=============TEST==============
elif args.version == 'Mix_base':
    from models.Mix_Det_mobile_basic import build_net
    cfg = mobile_300
elif args.version == 'Mix_module':
    from models.Mix_Det_mobile_module import build_net
    cfg = mobile_300
#===============================
else:
    print('ERROR::UNKNOWN VERSION')
    sys.exit()

img_dim = (300,512)[args.size=='512']
rgb_means = ((103.94,116.78,123.68), (104, 117, 123))[args.version == 'RFB_vgg' or args.version == 'RFB_E_vgg']
p = (0.2, 0.6)[args.version == 'RFB_vgg' or args.version == 'RFB_E_vgg']
#num_classes = (21, 81)[args.dataset == 'COCO']
batch_size = args.batch_size
weight_decay = 0.0005
gamma = 0.1
momentum = 0.9

net = build_net('train', img_dim, num_classes)
#print(net)
if args.resume_net == None:
    if args.basenet == None:
        print('!!Model training from scratch!!')
    else:
        print('!!Loading backbone network weight files!!')
        base_weights = torch.load(args.basenet)
        net.base.load_state_dict(base_weights)

    def xavier(param):
        init.xavier_uniform(param)

    def weights_init(m):
        for key in m.state_dict():
            if key.split('.')[-1] == 'weight':
                # for only weight files
                if 'conv' in key:
                    #print(key.split('.'))
                    if(m.state_dict()[key].dim()==1):
                        m.state_dict()[key][...] = 1
                    else:
                        init.kaiming_normal_(m.state_dict()[key], mode='fan_out')
                if 'bn' in key:
                    #print(key)
                    m.state_dict()[key][...] = 1
            elif key.split('.')[-1] == 'bias':
                #print(key.split('.'))
                m.state_dict()[key][...] = 0
            


    print('Initializing weights...')
# initialize newly added layers' weights with kaiming_normal method
    net.extras.apply(weights_init)
    net.loc.apply(weights_init)
    net.conf.apply(weights_init)
    net.Norm.apply(weights_init)
    if args.version == 'RFB_E_vgg':
        net.reduce.apply(weights_init)
        net.up_reduce.apply(weights_init)

else:
# load resume network
    print('Loading resume network...')
    state_dict = torch.load(args.resume_net)
    # create new OrderedDict that does not contain `module.`
    from collections import OrderedDict
    new_state_dict = OrderedDict()
    for k, v in state_dict.items():
        head = k[:7]
        if head == 'module.':
            name = k[7:] # remove `module.`
        else:
            name = k
        new_state_dict[name] = v
    net.load_state_dict(new_state_dict)

if args.ngpu > 1:
    net = torch.nn.DataParallel(net, device_ids=list(range(args.ngpu)))

if args.cuda:
    net.cuda()
    cudnn.benchmark = True


optimizer = optim.SGD(net.parameters(), lr=args.lr,
                      momentum=args.momentum, weight_decay=args.weight_decay)
#optimizer = optim.RMSprop(net.parameters(), lr=args.lr,alpha = 0.9, eps=1e-08,
#                      momentum=args.momentum, weight_decay=args.weight_decay)

criterion = MultiBoxLoss(num_classes, 0.5, True, 0, True, 3, 0.5, False)
priorbox = PriorBox(cfg)
with torch.no_grad():
    priors = priorbox.forward()
    if args.cuda:
        priors = priors.cuda()

def train():
    net.train()

    save_dir = os.path.join(args.save_folder, args.version + '_' + args.dataset)
    if not os.path.exists(save_dir):
        os.mkdir(save_dir)

    # loss counters
    loc_loss = 0  # epoch
    conf_loss = 0
    epoch = 0 + args.resume_epoch
    print('Loading Dataset...')

    if args.dataset == 'VOC':
        dataset = VOCDetection(VOCroot, train_sets, preproc(
            img_dim, rgb_means, p), AnnotationTransform())
    elif args.dataset == 'COCO':
        dataset = COCODetection(COCOroot, train_sets, preproc(
            img_dim, rgb_means, p))
    elif args.dataset == 'custom':
        dataset = CustomDetection(Customroot, train_sets, preproc(
            img_dim, rgb_means, p), AnnoTransform_custom())
    else:
        print('Choose dataset among VOC, COCO and custom')
        return
    epoch_size = len(dataset) // args.batch_size
    max_iter = args.max_epoch * epoch_size

    if args.dataset == 'VOC':
        stepvalues = (150 * epoch_size, 200 * epoch_size, 250 * epoch_size)
    elif args.dataset == 'COCO':
        stepvalues = (90 * epoch_size, 120 * epoch_size, 140 * epoch_size)
        #stepvalues = (100 * epoch_size, 130 * epoch_size, 150 * epoch_size)
    else: # for custom dataset
        stepvalues = (90 * epoch_size, 120 * epoch_size, 140 * epoch_size)
    #stepvalues = (stepvalues_VOC,stepvalues_COCO)[args.dataset=='COCO']
    print('Training',args.version, 'on', dataset.name)
    step_index = 0

    if args.resume_epoch > 0:
        start_iter = args.resume_epoch * epoch_size
    else:
        start_iter = 0

    if args.summary_writer != None:
        print('TENSORBOARD: $ tensorboard --logdir=' + args.summary_writer)
    else:
        print('TENSORBOARD: $ tensorboard --logdir=./runs/folder_name')
    lr = args.lr
    for iteration in range(start_iter, max_iter):
        if iteration % epoch_size == 0:
            # create batch iterator
            batch_iterator = iter(data.DataLoader(dataset, batch_size,
                                                  shuffle=True, num_workers=args.num_workers, collate_fn=detection_collate))
            loc_loss = 0
            conf_loss = 0
            if (epoch % 10 == 0 and epoch > 0) or (epoch % 5 == 0 and epoch > 150):
                filedir = os.path.join(save_dir, 'epoches_'+ repr(epoch) + '.pth')
                torch.save(net.state_dict(), filedir)
            epoch += 1

        load_t0 = time.time()
        if iteration in stepvalues:
            step_index += 1
        lr = adjust_learning_rate(optimizer, args.gamma, epoch, step_index, iteration, epoch_size)

        # load train data
        images, targets = next(batch_iterator)
        
        #print(np.sum([torch.sum(anno[:,-1] == 2) for anno in targets]))

        if args.cuda:
            images = Variable(images.cuda())
            targets = [Variable(anno.cuda()) for anno in targets]
        else:
            images = Variable(images)
            targets = [Variable(anno) for anno in targets]
        # forward
        t0 = time.time()
        out = net(images)

        # backprop
        optimizer.zero_grad()
        loss_l, loss_c = criterion(out, priors, targets)
        loss = loss_l + loss_c
                
        loss.backward()
        optimizer.step()
        t1 = time.time()
        loc_loss += loss_l.item()
        conf_loss += loss_c.item()
        load_t1 = time.time()
        if iteration % 10 == 0:
            status = 'Epoch:' + repr(epoch) + ' || epochiter: ' + repr(iteration % epoch_size) \
            + '/' + repr(epoch_size) + '|| Total iter ' + repr(iteration) \
            + ' || L: %.4f C: %.4f||' % (loss_l.item(),loss_c.item()) \
            + 'Batch time: %.4f sec. ||' % (load_t1 - load_t0) + 'LR: %.8f' % (lr) + '\r'
            sys.stdout.write(status)
            sys.stdout.flush()
            summary.add_scalar('loss/loss_l', loss_l.item(), iteration)
            summary.add_scalar('loss/loss_c', loss_c.item(), iteration)
            summary.add_scalar('learning_rate', lr, iteration)
            summary.add_scalars('loss/loss', {"loss_l": loss_l.item(),
                                            "loss_c": loss_c.item(),
                                            "loss": loss.item()}, iteration)
    filedir = os.path.join(save_dir, 'Final_epoch.pth')
    torch.save(net.state_dict(), filedir)

def adjust_learning_rate(optimizer, gamma, epoch, step_index, iteration, epoch_size):
    """Sets the learning rate 
    # Adapted from PyTorch Imagenet example:
    # https://github.com/pytorch/examples/blob/master/imagenet/main.py
    """
    if epoch < 6:
        lr = 1e-6 + (args.lr-1e-6) * iteration / (epoch_size * 5) 
    else:
        lr = args.lr * (gamma ** (step_index))
    for param_group in optimizer.param_groups:
        param_group['lr'] = lr
    return lr


if __name__ == '__main__':
    train()