DSDG/DUM/train.py

from __future__ import print_function, division
import torch
import matplotlib.pyplot as plt
import argparse, os
import numpy as np
from torch.utils.data import DataLoader
from torchvision import transforms

from models.CDCNs_u import Conv2d_cd, CDCN_u

from Load_OULUNPUcrop_train import Spoofing_train_g, SeparateBatchSampler, Normaliztion, ToTensor, \
    RandomHorizontalFlip, Cutout, RandomErasing
from Load_OULUNPUcrop_valtest import Spoofing_valtest, Normaliztion_valtest, ToTensor_valtest

import torch.nn.functional as F
import torch.nn as nn
import torch.optim as optim

from utils import AvgrageMeter, performances

# Dataset root
train_image_dir = '/export2/home/wht/oulu_img_crop/train_file_flod/'
val_image_dir = '/export2/home/wht/oulu_img_crop/dev_file_flod/'
test_image_dir = '/export2/home/wht/oulu_img_crop/test_file_flod/'

train_map_dir = '/export2/home/wht/oulu_img_crop/train_depth_flod/'
val_map_dir = '/export2/home/wht/oulu_img_crop/dev_depth_flod/'
test_map_dir = '/export2/home/wht/oulu_img_crop/test_depth_flod/'

train_list = '/export2/home/wht/oulu_img_crop/protocols/Protocol_1/Train_g.txt'
val_list = '/export2/home/wht/oulu_img_crop/protocols/Protocol_1/Dev.txt'
test_list = '/export2/home/wht/oulu_img_crop/protocols/Protocol_1/Test.txt'


def contrast_depth_conv(input):
    ''' compute contrast depth in both of (out, label) '''
    '''
        input  32x32
        output 8x32x32
    '''

    kernel_filter_list = [
        [[1, 0, 0], [0, -1, 0], [0, 0, 0]], [[0, 1, 0], [0, -1, 0], [0, 0, 0]], [[0, 0, 1], [0, -1, 0], [0, 0, 0]],
        [[0, 0, 0], [1, -1, 0], [0, 0, 0]], [[0, 0, 0], [0, -1, 1], [0, 0, 0]],
        [[0, 0, 0], [0, -1, 0], [1, 0, 0]], [[0, 0, 0], [0, -1, 0], [0, 1, 0]], [[0, 0, 0], [0, -1, 0], [0, 0, 1]]
    ]

    kernel_filter = np.array(kernel_filter_list, np.float32)

    kernel_filter = torch.from_numpy(kernel_filter.astype(np.float)).float().cuda()
    # weights (in_channel, out_channel, kernel, kernel)
    kernel_filter = kernel_filter.unsqueeze(dim=1)

    input = input.unsqueeze(dim=1).expand(input.shape[0], 8, input.shape[1], input.shape[2])

    contrast_depth = F.conv2d(input, weight=kernel_filter, groups=8)

    return contrast_depth


class Contrast_depth_loss(nn.Module):
    def __init__(self):
        super(Contrast_depth_loss, self).__init__()
        return

    def forward(self, out, label):
        contrast_out = contrast_depth_conv(out)
        contrast_label = contrast_depth_conv(label)

        criterion_MSE = nn.MSELoss().cuda()

        loss = criterion_MSE(contrast_out, contrast_label)

        return loss


def train_test():
    isExists = os.path.exists(args.log)
    if not isExists:
        os.makedirs(args.log)
    log_file = open(args.log + '/' + args.log + '_log_P1.txt', 'a')

    log_file.write('Oulu-NPU, P1:\n ')
    log_file.flush()

    print('train from scratch!\n')
    log_file.write('train from scratch!\n')
    log_file.write('lr:%.6f, lamda_kl:%.6f , batchsize:%d\n' % (args.lr, args.kl_lambda, args.batchsize))
    log_file.flush()

    model = CDCN_u(basic_conv=Conv2d_cd, theta=0.7)
    # model = ResNet18_u()

    model = model.cuda()
    model = torch.nn.DataParallel(model)

    lr = args.lr
    optimizer = optim.Adam(model.parameters(), lr=lr, weight_decay=0.00005)
    scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=args.step_size, gamma=args.gamma)

    print(model)

    criterion_absolute_loss = nn.MSELoss().cuda()
    criterion_contrastive_loss = Contrast_depth_loss().cuda()

    for epoch in range(args.epochs):
        if (epoch + 1) % args.step_size == 0:
            lr *= args.gamma

        loss_absolute_real = AvgrageMeter()
        loss_absolute_fake = AvgrageMeter()
        loss_contra_real = AvgrageMeter()
        loss_contra_fake = AvgrageMeter()
        loss_kl_real = AvgrageMeter()
        loss_kl_fake = AvgrageMeter()

        ###########################################
        '''                train             '''
        ###########################################
        model.train()

        # load random 16-frame clip data every epoch
        train_data = Spoofing_train_g(train_list, train_image_dir, train_map_dir,
                                      transform=transforms.Compose(
                                          [RandomErasing(), RandomHorizontalFlip(), ToTensor(), Cutout(),
                                           Normaliztion()]))
        train_real_idx, train_fake_idx = train_data.get_idx()
        batch_sampler = SeparateBatchSampler(train_real_idx, train_fake_idx, batch_size=args.batchsize, ratio=args.ratio)
        dataloader_train = DataLoader(train_data, num_workers=8, batch_sampler=batch_sampler)

        for i, sample_batched in enumerate(dataloader_train):
            # get the inputs
            inputs, map_label, spoof_label = sample_batched['image_x'].cuda(), sample_batched['map_x'].cuda(), \
                                             sample_batched['spoofing_label'].cuda()

            optimizer.zero_grad()

            # forward + backward + optimize
            mu, logvar, map_x, x_concat, x_Block1, x_Block2, x_Block3, x_input = model(inputs)

            mu_real = mu[:int(args.batchsize * args.ratio), :, :]
            logvar_real = logvar[:int(args.batchsize * args.ratio), :, :]
            map_x_real = map_x[:int(args.batchsize * args.ratio), :, :]
            map_label_real = map_label[:int(args.batchsize * args.ratio), :, :]

            absolute_loss_real = criterion_absolute_loss(map_x_real, map_label_real)
            contrastive_loss_real = criterion_contrastive_loss(map_x_real, map_label_real)
            kl_loss_real = -(1 + logvar_real - (mu_real - map_label_real).pow(2) - logvar_real.exp()) / 2
            kl_loss_real = kl_loss_real.sum(dim=1).sum(dim=1).mean()
            kl_loss_real = args.kl_lambda * kl_loss_real

            mu_fake = mu[int(args.batchsize * args.ratio):, :, :]
            logvar_fake = logvar[int(args.batchsize * args.ratio):, :, :]
            map_x_fake = map_x[int(args.batchsize * args.ratio):, :, :]
            map_label_fake = map_label[int(args.batchsize * args.ratio):, :, :]

            absolute_loss_fake = 0.1 * criterion_absolute_loss(map_x_fake, map_label_fake)
            contrastive_loss_fake = 0.1 * criterion_contrastive_loss(map_x_fake, map_label_fake)
            kl_loss_fake = -(1 + logvar_fake - (mu_fake - map_label_fake).pow(2) - logvar_fake.exp()) / 2
            kl_loss_fake = kl_loss_fake.sum(dim=1).sum(dim=1).mean()
            kl_loss_fake = 0.1 * args.kl_lambda * kl_loss_fake

            absolute_loss = absolute_loss_real + absolute_loss_fake
            contrastive_loss = contrastive_loss_real + contrastive_loss_fake
            kl_loss = kl_loss_real + kl_loss_fake

            loss = absolute_loss + contrastive_loss + kl_loss

            loss.backward()

            optimizer.step()

            n = inputs.size(0)
            loss_absolute_real.update(absolute_loss_real.data, n)
            loss_absolute_fake.update(absolute_loss_fake.data, n)
            loss_contra_real.update(contrastive_loss_real.data, n)
            loss_contra_fake.update(contrastive_loss_fake.data, n)
            loss_kl_real.update(kl_loss_real.data, n)
            loss_kl_fake.update(kl_loss_fake.data, n)

        scheduler.step()
        # whole epoch average
        print(
            'epoch:%d, Train:  Absolute_loss: real=%.4f,fake=%.4f, '
            'Contrastive_loss: real=%.4f,fake=%.4f, kl_loss: real=%.4f,fake=%.4f' % (
                epoch + 1, loss_absolute_real.avg, loss_absolute_fake.avg, loss_contra_real.avg, loss_contra_fake.avg,
                loss_kl_real.avg, loss_kl_fake.avg))

        # validation/test
        if epoch < 200:
            epoch_test = 200
        else:
            epoch_test = 50
        # epoch_test = 1
        if epoch % epoch_test == epoch_test - 1:
            model.eval()

            with torch.no_grad():
                ###########################################
                '''                val             '''
                ###########################################
                # val for threshold
                val_data = Spoofing_valtest(val_list, val_image_dir, val_map_dir,
                                            transform=transforms.Compose([Normaliztion_valtest(), ToTensor_valtest()]))
                dataloader_val = DataLoader(val_data, batch_size=1, shuffle=False, num_workers=4)

                map_score_list = []

                for i, sample_batched in enumerate(dataloader_val):
                    # get the inputs
                    inputs, spoof_label = sample_batched['image_x'].cuda(), sample_batched['spoofing_label'].cuda()
                    val_maps = sample_batched['val_map_x'].cuda()  # binary map from PRNet

                    optimizer.zero_grad()

                    mu, logvar, map_x, x_concat, x_Block1, x_Block2, x_Block3, x_input = model(inputs.squeeze(0))
                    score_norm = mu.sum(dim=1).sum(dim=1) / val_maps.squeeze(0).sum(dim=1).sum(dim=1)
                    map_score = score_norm.mean()
                    map_score_list.append('{} {}\n'.format(map_score, spoof_label[0][0]))

                map_score_val_filename = args.log + '/' + args.log + '_map_score_val.txt'
                with open(map_score_val_filename, 'w') as file:
                    file.writelines(map_score_list)

                ###########################################
                '''                test             '''
                ##########################################
                # test for ACC
                test_data = Spoofing_valtest(test_list, test_image_dir, test_map_dir,
                                             transform=transforms.Compose([Normaliztion_valtest(), ToTensor_valtest()]))
                dataloader_test = DataLoader(test_data, batch_size=1, shuffle=False, num_workers=4)

                map_score_list = []

                for i, sample_batched in enumerate(dataloader_test):
                    # get the inputs
                    inputs, spoof_label = sample_batched['image_x'].cuda(), sample_batched['spoofing_label'].cuda()
                    test_maps = sample_batched['val_map_x'].cuda()

                    optimizer.zero_grad()
                    mu, logvar, map_x, x_concat, x_Block1, x_Block2, x_Block3, x_input = model(inputs.squeeze(0))
                    score_norm = mu.sum(dim=1).sum(dim=1) / test_maps.squeeze(0).sum(dim=1).sum(dim=1)
                    map_score = score_norm.mean()
                    map_score_list.append('{} {}\n'.format(map_score, spoof_label[0][0]))

                map_score_test_filename = args.log + '/' + args.log + '_map_score_test.txt'
                with open(map_score_test_filename, 'w') as file:
                    file.writelines(map_score_list)

                #############################################################
                #       performance measurement both val and test
                #############################################################
                val_threshold, test_threshold, val_ACC, val_ACER, test_ACC, test_APCER, test_BPCER, test_ACER, test_ACER_test_threshold = performances(
                    map_score_val_filename, map_score_test_filename)

                print('epoch:%d, Val:  val_threshold= %.4f, val_ACC= %.4f, val_ACER= %.4f' % (
                    epoch + 1, val_threshold, val_ACC, val_ACER))
                log_file.write('\n epoch:%d, Val:  val_threshold= %.4f, val_ACC= %.4f, val_ACER= %.4f \n' % (
                    epoch + 1, val_threshold, val_ACC, val_ACER))

                print('epoch:%d, Test:  ACC= %.4f, APCER= %.4f, BPCER= %.4f, ACER= %.4f' % (
                    epoch + 1, test_ACC, test_APCER, test_BPCER, test_ACER))
                log_file.write('epoch:%d, Test:  ACC= %.4f, APCER= %.4f, BPCER= %.4f, ACER= %.4f \n' % (
                    epoch + 1, test_ACC, test_APCER, test_BPCER, test_ACER))
                log_file.flush()

        if epoch % epoch_test == epoch_test - 1:
            # save the model until the next improvement
            torch.save(model.state_dict(), args.log + '/' + args.log + '_%d.pkl' % (epoch + 1))

    print('Finished Training')
    log_file.close()

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="save quality using landmarkpose model")
    parser.add_argument('--gpus', type=str, default='0, 1, 2, 3', help='the gpu id used for predict')
    parser.add_argument('--lr', type=float, default=0.0001, help='initial learning rate')
    parser.add_argument('--batchsize', type=int, default=64, help='initial batchsize')
    parser.add_argument('--step_size', type=int, default=500, help='how many epochs lr decays once')  # 500
    parser.add_argument('--gamma', type=float, default=0.5, help='gamma of optim.lr_scheduler.StepLR, decay of lr')
    parser.add_argument('--kl_lambda', type=float, default=0.001, help='')
    parser.add_argument('--ratio', type=float, default=0.75, help='real and fake in batchsize ')
    parser.add_argument('--echo_batches', type=int, default=50, help='how many batches display once')  # 50
    parser.add_argument('--epochs', type=int, default=1600, help='total training epochs')
    parser.add_argument('--log', type=str, default="CDCN_U_P1", help='log and save model name')
    parser.add_argument('--finetune', action='store_true', default=False, help='whether finetune other models')

    args = parser.parse_args()
    train_test()