train.py

import os
import torch.utils.data
from torch import nn
from torch.nn import DataParallel
from datetime import datetime
from config import BATCH_SIZE, SAVE_FREQ, RESUME, SAVE_DIR, TEST_FREQ, TOTAL_EPOCH, MODEL_PRE
from config import CASIA_DATA_DIR, LFW_DATA_DIR
from core import model, head
from core.utils import init_log
from dataloader.CASIA_Face_loader import CASIA_Face
from dataloader.LFW_loader import LFW
from dataloader.augmenter import Augmenter
from torch.optim import lr_scheduler
from torchsummary import summary
import torch.optim as optim
import time
from lfw_eval import parseList, evaluation_10_fold
import numpy as np
import scipy.io
import sys
from data_utils import extract_deep_feature

if __name__ == '__main__':
    # other init
    start_epoch = 1
    save_dir = os.path.join(SAVE_DIR, MODEL_PRE + datetime.now().strftime('%Y%m%d_%H%M%S'))
    if os.path.exists(save_dir):
        raise NameError('model dir exists!')
    os.makedirs(save_dir)
    logging = init_log(save_dir)
    _print = logging.info

    # define trainloader and testloader
    print('defining casia dataloader...')
    trainset = CASIA_Face(root=CASIA_DATA_DIR, augmenter=Augmenter(0.2, 0.2, 0.2))
    trainloader = torch.utils.data.DataLoader(trainset, batch_size=BATCH_SIZE,
                                              shuffle=True, num_workers=5, drop_last=False)

    # nl: left_image_path
    # nr: right_image_path
    print('defining lfw dataloader...')
    nl, nr, folds, flags = parseList(root=LFW_DATA_DIR, name="lfw-112x112")
    testdataset = LFW(nl, nr)
    testloader = torch.utils.data.DataLoader(testdataset, batch_size=32,
                                             shuffle=False, num_workers=5, drop_last=False)

    # define model
    print('defining vargfacenet model...')
    net = model.VarGFaceNet()

    if RESUME:
        print("Resume Training...")
        ckpt = torch.load(RESUME)
        net.load_state_dict(ckpt['net_state_dict'])
        start_epoch = ckpt['epoch'] + 1

    net = net.cuda()
    summary(net, input_size=(3, 112, 112))

    # NLLLoss
    nllloss = nn.CrossEntropyLoss().cuda()
    # CenterLoss
    lmcl_loss = head.build_head(head_type='adaface',
                                embedding_size=512,
                                class_num=trainset.class_nums,
                                m=0.4,
                                h=0.333,
                                s=64.,
                                t_alpha=0.01).cuda()


    if torch.cuda.device_count() > 1:
        print("Train with MultiGPUs")
        net = DataParallel(net)
        lmcl_loss = DataParallel(lmcl_loss)

    criterion = [nllloss, lmcl_loss]

    # optimzer4nn
    # optimizer4nn = optim.Adam(net.parameters(), lr=0.001, weight_decay=0.0005)
    # sheduler_4nn = lr_scheduler.StepLR(optimizer4nn, 20, gamma=0.5)

    # optimzer4center
    # optimizer4center = optim.Adam(lmcl_loss.parameters(), lr=0.01)
    # sheduler_4center = lr_scheduler.StepLR(optimizer4center, 20, gamma=0.5)
    
    # All Optimizer
    optimizer = optim.SGD([
        {'params': net.parameters(), 'weight_decay': 5e-4},
        {'params': lmcl_loss.parameters(), 'weight_decay': 5e-4}
    ], lr=0.01, momentum=0.9, nesterov=True)

    scheduler = lr_scheduler.MultiStepLR(optimizer, milestones=[10, 18, 22], gamma=0.1)

    best_acc = 0.0
    best_epoch = 0
    for epoch in range(start_epoch, TOTAL_EPOCH+1):
        # exp_lr_scheduler.step()
        # optimizer4nn.step()
        # optimizer4center.step()
        # optimizer.step()
        scheduler.step()
             
        # train model
        _print('Train Epoch: {}/{} ...'.format(epoch, TOTAL_EPOCH))
        net.train()

        train_total_loss = 0.0
        total = 0
        since = time.time()
        total_step = len(trainloader)
        for i, data in enumerate(trainloader):
            sys.stdout.write("\r Step: {0}/{1}".format(i, total_step))
            sys.stdout.flush()
            img, label = data[0].cuda(), data[1].cuda()
            batch_size = img.size(0)

            raw_logits, norms = net(img)

            mlogits = criterion[1](raw_logits, norms, label)
            total_loss = criterion[0](mlogits, label)

            # optimizer4nn.zero_grad() 
            # optimizer4center.zero_grad()
            optimizer.zero_grad()

            total_loss.backward()
            
            # optimizer4nn.step()
            # optimizer4center.step()
            optimizer.step() 

            train_total_loss += total_loss.item() * batch_size
            total += batch_size

        train_total_loss = train_total_loss / total
        time_elapsed = time.time() - since
        loss_msg = '    total_loss: {:.4f} time: {:.0f}m {:.0f}s'\
            .format(train_total_loss, time_elapsed // 60, time_elapsed % 60)
        _print(loss_msg)

        # test model on lfw
        if epoch % TEST_FREQ == 0:
            net.eval()
            featureLs = None
            featureRs = None
            _print('Test Epoch: {} ...'.format(epoch))
            total_step = len(testloader)
            for i, data in enumerate(testloader):
                sys.stdout.write("\r Step: {0}/{1}".format(i, total_step))
                sys.stdout.flush()
                
                for i in range(len(data)):
                    data[i] = data[i].cuda()
                    
                res = []
                for d in data:
                    out, _ = net(d)
                    fliped_image = torch.flip(d, dims=[3])
                    flipped_embedding, flipped_ = net(fliped_image)
                    embedding = extract_deep_feature(out, _, flipped_embedding, flipped_)
                    res.append(embedding)
                    
                featureL = np.concatenate((res[0], res[1]), 1)
                featureR = np.concatenate((res[2], res[3]), 1)
                
                if featureLs is None:
                    featureLs = featureL
                else:
                    featureLs = np.concatenate((featureLs, featureL), 0)
                    
                if featureRs is None:
                    featureRs = featureR
                else:
                    featureRs = np.concatenate((featureRs, featureR), 0)

            result = {'fl': featureLs, 'fr': featureRs, 'fold': folds, 'flag': flags}
            # save tmp_result
            if not os.path.exists('./result'):
                os.makedirs('./result')
            scipy.io.savemat('./result/tmp_result.mat', result)
            accs = evaluation_10_fold('./result/tmp_result.mat')
            _print('    ave: {:.4f}'.format(np.mean(accs) * 100))

        # save model
        if epoch % SAVE_FREQ == 0:
            msg = 'Saving checkpoint: {}'.format(epoch)
            _print(msg)
            if torch.cuda.device_count() > 1:
                net_state_dict = net.module.state_dict()
            else:
                net_state_dict = net.state_dict()
            if not os.path.exists(save_dir):
                os.mkdir(save_dir)
            torch.save({
                'epoch': epoch,
                'net_state_dict': net_state_dict},
                os.path.join(save_dir, '%03d.ckpt' % epoch))
    print('finishing training')