train.py

from __future__ import print_function
import argparse
import sys
import time
import torch
import torch.nn as nn
import torch.optim as optim
import torch.backends.cudnn as cudnn
from torch.autograd import Variable
import torch.utils.data as data
import torchvision
import torchvision.transforms as transforms
from data_loader import SYSUData, RegDBData, TestData
from data_manager import *
from eval_metrics import eval_sysu, eval_regdb
from model import embed_net
from utils import *
from tripletloss import TripletLoss
from RandomErasing import RandomErasing
from model import Normalize
from closs import CovLoss

parser = argparse.ArgumentParser(description='PyTorch Cross-Modality Training')
parser.add_argument('--dataset', default='sysu',
                    help='dataset name: regdb or sysu]')
parser.add_argument('--lr', default=3.5e-5, type=float,
                    help='learning rate')  # 3.5e-5
parser.add_argument('--optim', default='adam', type=str, help='optimizer')
parser.add_argument('--arch', default='resnet50', type=str,
                    help='network baseline:resnet18 or resnet50')
parser.add_argument('--resume', '-r', default='', type=str,
                    help='resume from checkpoint')
parser.add_argument('--test-only', action='store_true', help='test only')
parser.add_argument('--model_path', default='../Log/save_model/', type=str,
                    help='model save path')
parser.add_argument('--save_epoch', default=20, type=int,
                    metavar='s', help='save model every 10 epochs')
parser.add_argument('--log_path', default='log/', type=str,
                    help='log save path')
parser.add_argument('--workers', default=4, type=int, metavar='N',
                    help='number of data loading workers (default: 4)')
parser.add_argument('--low-dim', default=512, type=int,
                    metavar='D', help='feature dimension')
parser.add_argument('--img_w', default=144, type=int,
                    metavar='imgw', help='img width')
parser.add_argument('--img_h', default=288, type=int,
                    metavar='imgh', help='img height')
parser.add_argument('--batch-size', default=32, type=int,
                    metavar='B', help='training batch size')
parser.add_argument('--test-batch', default=64, type=int,
                    metavar='tb', help='testing batch size')
parser.add_argument('--method', default='id', type=str,
                    metavar='m', help='method type')
parser.add_argument('--drop', default=0.0, type=float,
                    metavar='drop', help='dropout ratio')
parser.add_argument('--trial', default=1, type=int,
                    metavar='t', help='trial (only for RegDB dataset)')
parser.add_argument('--gpu', default='0', type=str,
                    help='gpu device ids for CUDA_VISIBLE_DEVICES')
parser.add_argument('--mode', default='all', type=str, help='all or indoor')

args = parser.parse_args()
os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu
np.random.seed(0)

final_dim = 2048
dataset = args.dataset
if dataset == 'sysu':
    data_path = './SYSUMM01/'
    log_path = args.log_path + 'sysu_log/'
    test_mode = [1, 2]  # thermal to visible
elif dataset == 'regdb':
    data_path = 'RegDB/'
    log_path = args.log_path + 'regdb_log/'
    test_mode = [2, 1]  # visible to thermal

checkpoint_path = args.model_path

if not os.path.isdir(log_path):
    os.makedirs(log_path)
if not os.path.isdir(checkpoint_path):
    os.makedirs(checkpoint_path)

if args.method == 'id':
    suffix = dataset + '_id_bn_relu'
suffix = suffix + '_drop_{}'.format(args.drop)
suffix = suffix + '_lr_{:1.1e}'.format(args.lr)
suffix = suffix + '_dim_{}'.format(args.low_dim)
if not args.optim == 'sgd':
    suffix = suffix + '_' + args.optim
suffix = suffix + '_' + args.arch
if dataset == 'regdb':
    suffix = suffix + '_trial_{}'.format(args.trial)

test_log_file = open(log_path + suffix + '.txt', "w")
sys.stdout = Logger(log_path + suffix + '_os.txt')

device = 'cuda' if torch.cuda.is_available() else 'cpu'
best_acc = 0  # best test accuracy
start_epoch = 0
feature_dim = args.low_dim

print('==> Loading data..')
# Data loading code
randomerase = RandomErasing(probability=0.5)
normalize = transforms.Normalize(
    mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
transform_train = transforms.Compose([
    transforms.ToPILImage(),
    transforms.Pad(10),
    transforms.RandomCrop((args.img_h, args.img_w)),
    transforms.RandomHorizontalFlip(),
    transforms.ToTensor(),
    normalize,
    randomerase,
])
transform_test = transforms.Compose([
    transforms.ToPILImage(),
    transforms.Resize((args.img_h, args.img_w)),
    transforms.ToTensor(),
    normalize,
])

end = time.clock()
if dataset == 'sysu':
    # training set
    trainset = SYSUData(data_path,  transform=transform_train)
    # generate the idx of each person identity
    color_pos, thermal_pos = GenIdx(
        trainset.train_color_label, trainset.train_thermal_label)

    # testing set
    query_img, query_label, query_cam = process_query_sysu(
        data_path, mode=args.mode)
    gall_img, gall_label, gall_cam = process_gallery_sysu(
        data_path, mode=args.mode, trial=0)

elif dataset == 'regdb':
    # training set
    trainset = RegDBData(data_path, args.trial, transform=transform_train)
    # generate the idx of each person identity
    color_pos, thermal_pos = GenIdx(
        trainset.train_color_label, trainset.train_thermal_label)

    # testing set
    query_img, query_label = process_test_regdb(
        data_path, trial=args.trial, modal='visible')
    gall_img, gall_label = process_test_regdb(
        data_path, trial=args.trial, modal='thermal')

gallset = TestData(gall_img, gall_label, transform=transform_test,
                   img_size=(args.img_w, args.img_h))
queryset = TestData(query_img, query_label,
                    transform=transform_test, img_size=(args.img_w, args.img_h))

# testing data loader
gall_loader = data.DataLoader(
    gallset, batch_size=args.test_batch, shuffle=False, num_workers=args.workers)
query_loader = data.DataLoader(
    queryset, batch_size=args.test_batch, shuffle=False, num_workers=args.workers)

n_class = len(np.unique(trainset.train_color_label))
nquery = len(query_label)
ngall = len(gall_label)

print('Dataset {} statistics:'.format(dataset))
print('  ------------------------------')
print('  subset   | # ids | # images')
print('  ------------------------------')
print('  visible  | {:5d} | {:8d}'.format(
    n_class, len(trainset.train_color_label)))
print('  thermal  | {:5d} | {:8d}'.format(
    n_class, len(trainset.train_thermal_label)))
print('  ------------------------------')
print('  query    | {:5d} | {:8d}'.format(len(np.unique(query_label)), nquery))
print('  gallery  | {:5d} | {:8d}'.format(len(np.unique(gall_label)), ngall))
print('  ------------------------------')
print('Data Loading Time:\t {:.3f}'.format(time.clock()-end))


print('==> Building model..')
net = embed_net(final_dim, n_class, drop=args.drop,
                arch=args.arch)
net.to(device)
cudnn.benchmark = True

if len(args.resume) > 0:
    model_path = checkpoint_path + args.resume
    if os.path.isfile(model_path):
        print('==> loading checkpoint {}'.format(args.resume))
        checkpoint = torch.load(model_path)
        start_epoch = checkpoint['epoch']
        net.load_state_dict(checkpoint['net'])
        print('==> loaded checkpoint {} (epoch {})'
              .format(args.resume, checkpoint['epoch']))
    else:
        print('==> no checkpoint found at {}'.format(args.resume))

if args.method == 'id':
    criterion = nn.CrossEntropyLoss()
    criterion.to(device)


#**********************#

if args.optim == 'sgd':
    optimizer = optim.SGD([
        {'params': net.parameters(), 'lr': args.lr}, ],
        weight_decay=5e-4, momentum=0.9, nesterov=True)
elif args.optim == 'adam':
    optimizer = optim.Adam([
        {'params': net.parameters(), 'lr': args.lr}, ],
        weight_decay=5e-4)


def adjust_learning_rate(optimizer, epoch):
    """Sets the learning rate to the initial LR decayed by 10 every 30 epochs"""
    if epoch < 10 and epoch >= 0:
        lr = args.lr * (epoch + 1)
    elif epoch >= 10 and epoch < 40:
        lr = args.lr * 10
    elif epoch >= 40 and epoch < 70:
        lr = args.lr
    else:
        lr = args.lr * 0.1

    optimizer.param_groups[0]['lr'] = lr
    return lr


def train(epoch):
    current_lr = adjust_learning_rate(optimizer, epoch)
    train_loss = AverageMeter()
    data_time = AverageMeter()
    batch_time = AverageMeter()
    correct = 0
    total = 0
    # New
    tripletloss_global = TripletLoss(args.batch_size, 4)
    Covloss = CovLoss(batchsize=args.batch_size, num_instance=4)
    softmax = nn.Softmax(dim=1)
    l2norm = Normalize()

    # switch to train mode
    net.train()
    end = time.clock()
    for batch_idx, (input1, input2, label1, label2) in enumerate(trainloader):
        input1 = Variable(input1.cuda())
        input2 = Variable(input2.cuda())
        label1 = Variable(label1.cuda())
        label2 = Variable(label2.cuda())
        labels = torch.cat((label1, label2), 0)
        inputs = torch.cat((input1, input2), 0)

        data_time.update(time.clock() - end)
        outputs, feat, _, outputs3, feat3, _ = net(inputs)
        if args.method == 'id':
            loss = criterion(outputs, labels)
            loss = loss + criterion(outputs3, labels)
            score = softmax(outputs)

            score3 = softmax(outputs3)

            _, predicted = torch.max(score.data, 1)
            correct += predicted.eq(labels).sum().item()
            _, predicted3 = torch.max(score3.data, 1)
            correct3 = predicted3.eq(labels).sum().item()
            correct = correct + correct3
        feat = l2norm(feat)
        feat3 = l2norm(feat3)
        closs = Covloss(feat)
        closs3 = Covloss(feat3)
        triloss = tripletloss_global(
            feat, labels)
        triloss3 = tripletloss_global(
            feat3, labels)
        loss = triloss + loss + triloss3 + closs + closs3

        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        train_loss.update(loss.item(), 2*input1.size(0))

        total += labels.size(0)

        # measure elapsed time
        batch_time.update(time.clock() - end)
        end = time.clock()
        if batch_idx % 10 == 0:
            print('Epoch: [{}][{}/{}] '
                  'Time: {batch_time.val:.3f} ({batch_time.avg:.3f}) '
                  'Data: {data_time.val:.3f} ({data_time.avg:.3f}) '
                  'lr:{} '
                  'Loss: {train_loss.val:.4f} ({train_loss.avg:.4f}) '
                  'Accu: {:.2f}' .format(
                      epoch, batch_idx, len(trainloader), current_lr,
                      50.*correct/total, batch_time=batch_time,
                      data_time=data_time, train_loss=train_loss))


def test(epoch):
    # switch to evaluation mode
    l2norm = Normalize()
    net.eval()
    print('Extracting Gallery Feature...')
    start = time.clock()
    ptr = 0
    gall_feat = np.zeros((ngall, final_dim+1024))
    with torch.no_grad():
        for batch_idx, (input, label) in enumerate(gall_loader):
            batch_num = input.size(0)
            input = Variable(input.cuda())
            _, _, feat, _, _, feat3 = net(input)
            feat = l2norm(feat)
            feat3 = l2norm(feat3)
            feat = torch.cat((feat, feat3), dim=1)
            gall_feat[ptr:ptr+batch_num, :] = feat.detach().cpu().numpy()
            ptr = ptr + batch_num
    print('Extracting Time:\t {:.3f}'.format(time.clock()-start))

    # switch to evaluation mode
    net.eval()
    print('Extracting Query Feature...')
    start = time.clock()
    ptr = 0
    query_feat = np.zeros((nquery, final_dim+1024))
    with torch.no_grad():
        for batch_idx, (input, label) in enumerate(query_loader):
            batch_num = input.size(0)
            input = Variable(input.cuda())
            _, _, feat, _, _, feat3 = net(input)
            feat = l2norm(feat)
            feat3 = l2norm(feat3)
            feat = torch.cat((feat, feat3), dim=1)
            query_feat[ptr:ptr+batch_num, :] = feat.detach().cpu().numpy()
            ptr = ptr + batch_num
    print('Extracting Time:\t {:.3f}'.format(time.clock()-start))

    start = time.clock()
    distmat = np.matmul(query_feat, np.transpose(gall_feat))

    # evaluation
    if dataset == 'regdb':
        cmc, mAP = eval_regdb(-distmat, query_label, gall_label)
    elif dataset == 'sysu':
        cmc, mAP = eval_sysu(-distmat, query_label,
                             gall_label, query_cam, gall_cam)
    print('Evaluation Time:\t {:.3f}'.format(time.clock()-start))
    return cmc, mAP


# training
print('==> Start Training...')

for epoch in range(start_epoch, 121):

    print('==> Preparing Data Loader...')
    # identity sampler
    sampler = IdentitySampler(trainset.train_color_label,
                              trainset.train_thermal_label, color_pos, thermal_pos, args.batch_size)
    trainset.cIndex = sampler.index1  # color index
    trainset.tIndex = sampler.index2  # thermal index
    trainloader = data.DataLoader(trainset, batch_size=args.batch_size,
                                  sampler=sampler, num_workers=args.workers, drop_last=True)

    # training
    train(epoch)

    if epoch > 0 and epoch % 2 == 0:
        print('Test Epoch: {}'.format(epoch))
        print('Test Epoch: {}'.format(epoch), file=test_log_file)
        # testing
        cmc, mAP = test(epoch)

        print('FC:   Rank-1: {:.2%} | Rank-5: {:.2%} | Rank-10: {:.2%}| mAP: {:.2%}'.format(
            cmc[0], cmc[4], cmc[9], mAP))
        print('FC:   Rank-1: {:.2%} | Rank-5: {:.2%} | Rank-10: {:.2%}| mAP: {:.2%}'.format(
            cmc[0], cmc[4], cmc[9], mAP), file=test_log_file)
        test_log_file.flush()

        # save model
        if cmc[0] > best_acc:  # not the real best for sysu-mm01
            best_acc = cmc[0]
            state = {
                'net': net.state_dict(),
                'cmc': cmc,
                'mAP': mAP,
                'epoch': epoch,
            }
            torch.save(state, checkpoint_path +
                       suffix + 'origin_origin_best.t')

        # save model every 20 epochs
        if epoch > 10 and epoch % 10 == 0:
            state = {
                'net': net.state_dict(),
                'cmc': cmc,
                'mAP': mAP,
                'epoch': epoch,
            }
            torch.save(state, checkpoint_path + suffix +
                       'origin_origin_epoch_{}.t'.format(epoch))