extract.py

import argparse
import os
import time
import h5py
import numpy

import torch
import torch.nn as nn
import torch.nn.parallel
import torch.backends.cudnn as cudnn
from torch.autograd import Variable

import torchvision.transforms as transforms
import torchvision.datasets as datasets

import models.convnets as convnets
import vqa.datasets as datasets
from vqa.lib.dataloader import DataLoader
from vqa.lib.logger import AvgMeter

parser = argparse.ArgumentParser(description='Extract')
parser.add_argument('--dataset', default='coco',
                    choices=['coco', 'vgenome', 'clevr'],
                    help='dataset type: coco (default) | vgenome | clevr')
parser.add_argument('--dir_data', default='data/coco',
                    help='dir dataset to download or/and load images')
parser.add_argument('--data_split', default='train', type=str,
                    help='Options: (default) train | val | test')
parser.add_argument('--arch', '-a', default='resnet152',
                    choices=convnets.model_names,
                    help='model architecture: ' +
                        ' | '.join(convnets.model_names) +
                        ' (default: fbresnet152)')
parser.add_argument('--conv', action='store_true', help='Whether to output the conv feature map')
parser.add_argument('--workers', default=4, type=int, 
                    help='number of data loading workers (default: 4)')
parser.add_argument('--batch_size', '-b', default=80, type=int, 
                    help='mini-batch size (default: 80)')
parser.add_argument('--mode', default='both', type=str,
                    help='Options: att | noatt | (default) both')
parser.add_argument('--size', default=448, type=int,
                    help='Image size (448 for noatt := avg pooling to get 224) (default:448)')
parser.add_argument('--dilation', type=int, default=1, help='[dilation] for pertrained conv layers')


def main():
    global args
    args = parser.parse_args()

    print("=> using pre-trained model '{}'".format(args.arch))
    model = convnets.factory({'arch':args.arch, 'conv':args.conv, 'dilation': args.dilation}, cuda=True, data_parallel=True)

    extract_name = 'arch,{}_size,{}'.format(args.arch, args.size)
    if args.conv:
        extract_name += '_conv'

    print('The output will be dumped to: {}'.format(extract_name))
    normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406],
                                     std=[0.229, 0.224, 0.225])

    if args.dataset == 'coco':
        if 'coco' not in args.dir_data:
            raise ValueError('"coco" string not in dir_data')
        dataset = datasets.COCOImages(args.data_split, dict(dir=args.dir_data), 
            transform=transforms.Compose([
                transforms.Scale(args.size),
                transforms.CenterCrop(args.size),
                transforms.ToTensor(),
                normalize,
            ]))
    elif args.dataset == 'vgenome':
        if args.data_split != 'train':
            raise ValueError('train split is required for vgenome')
        if 'vgenome' not in args.dir_data:
            raise ValueError('"vgenome" string not in dir_data')
        dataset = datasets.VisualGenomeImages(args.data_split, dict(dir=args.dir_data), 
            transform=transforms.Compose([
                transforms.Scale(args.size),
                transforms.CenterCrop(args.size),
                transforms.ToTensor(),
                normalize,
            ]))
    elif args.dataset == 'clevr':
        if 'clevr' not in args.dir_data:
            raise ValueError('"clevr" string not in dir_data')
        dataset = datasets.CLEVRImages(args.data_split, dict(dir=args.dir_data), 
            transform=transforms.Compose([
                transforms.Scale(args.size),
                transforms.CenterCrop(args.size),
                transforms.ToTensor(),
                normalize,
            ]))

    data_loader = DataLoader(dataset,
        batch_size=args.batch_size, shuffle=False,
        num_workers=args.workers, pin_memory=True)

    dir_extract = os.path.join(args.dir_data, 'extract', extract_name)
    path_file = os.path.join(dir_extract, args.data_split + 'set')
    os.system('mkdir -p ' + dir_extract)

    extract(data_loader, model, path_file, args.mode)


def extract(data_loader, model, path_file, mode):
    path_hdf5 = path_file + '.hdf5'
    path_txt = path_file + '.txt'
    hdf5_file = h5py.File(path_hdf5, 'w')

    # estimate output shapes
    output = model(Variable(torch.ones(1, 3, args.size, args.size),
                            volatile=True))

    nb_images = len(data_loader.dataset)
    if mode == 'both' or mode == 'att':
        shape_att = (nb_images, output.size(1), output.size(2), output.size(3))
        print('Warning: shape_att={}'.format(shape_att))
        hdf5_att = hdf5_file.create_dataset('att', shape_att,
                                            dtype='f')#, compression='gzip')
    if mode == 'both' or mode == 'noatt':
        shape_noatt = (nb_images, output.size(1))
        print('Warning: shape_noatt={}'.format(shape_noatt))
        hdf5_noatt = hdf5_file.create_dataset('noatt', shape_noatt,
                                              dtype='f')#, compression='gzip')

    model.eval()

    batch_time = AvgMeter()
    data_time  = AvgMeter()
    begin = time.time()
    end = time.time()

    idx = 0
    for i, input in enumerate(data_loader):
        input_var = Variable(input['visual'], volatile=True)
        output_att = model(input_var)
        batch_size = output_att.size(0)
        nb_regions = output_att.size(2) * output_att.size(3)
        output_noatt = output_att.sum(3).sum(2).div(nb_regions).view(batch_size, -1)
        
        if mode == 'both' or mode == 'att':
            hdf5_att[idx:idx+batch_size]   = output_att.data.cpu().numpy()
        if mode == 'both' or mode == 'noatt':
            hdf5_noatt[idx:idx+batch_size] = output_noatt.data.cpu().numpy()
        idx += batch_size

        #torch.cuda.synchronize()
        batch_time.update(time.time() - end)
        end = time.time()

        if i % 1 == 0:
            print('Extract: [{0}/{1}]\t'
                  'Time {batch_time.val:.3f} ({batch_time.avg:.3f})\t'
                  'Data {data_time.val:.3f} ({data_time.avg:.3f})\t'.format(
                   i, len(data_loader),
                   batch_time=batch_time,
                   data_time=data_time,))
            
    hdf5_file.close()

    # Saving image names in the same order than extraction
    with open(path_txt, 'w') as handle:
        for name in data_loader.dataset.dataset.imgs:
            handle.write(name + '\n')

    end = time.time() - begin
    print('Finished in {}m and {}s'.format(int(end/60), int(end%60)))


if __name__ == '__main__':
    main()