main_ce.py

from __future__ import print_function

import argparse
import copy
import os
import math
import sys
import time

from sklearn.model_selection import train_test_split
import tensorboard_logger as tb_logger
import torch
import torch.backends.cudnn as cudnn
from torch.utils.data import DataLoader, DistributedSampler
from torchvision import transforms, datasets

import sampler
from util import AverageMeter, DoubleTransform, SubsetWithTargets, TwoCropTransform
from util import adjust_learning_rate, warmup_learning_rate, accuracy, set_optimizer, save_model
from networks.resnet_big import SupCEResNet


def parse_option():
    parser = argparse.ArgumentParser('argument for training')

    parser.add_argument('--print_freq', type=int, default=10,
                        help='print frequency')
    parser.add_argument('--save_freq', type=int, default=50,
                        help='save frequency')
    parser.add_argument('--batch_size', type=int, default=256,
                        help='batch_size')
    parser.add_argument('--num_workers', type=int, default=8,
                        help='num of workers to use')
    parser.add_argument('--epochs', type=int, default=500,
                        help='number of training epochs')

    # optimization
    parser.add_argument('--learning_rate', type=float, default=0.2,
                        help='learning rate')
    parser.add_argument('--lr_decay_epochs', type=str, default='350,400,450',
                        help='where to decay lr, can be a list')
    parser.add_argument('--lr_decay_rate', type=float, default=0.1,
                        help='decay rate for learning rate')
    parser.add_argument('--weight_decay', type=float, default=1e-4,
                        help='weight decay')
    parser.add_argument('--momentum', type=float, default=0.9,
                        help='momentum')

    # model dataset
    parser.add_argument('--model', type=str, default='resnet50')
    parser.add_argument('--dataset', type=str, default='cifar10',
                        choices=['cifar10', 'cifar100', 'imagenet100', 'imagenet'],
                        help='dataset')
    parser.add_argument('--valid_split', type=float, default=0,
                        help="proportion of train data to use for validation set")
    parser.add_argument('--size', type=int, default=32,
                        help='size of images after resizing')

    # other setting
    parser.add_argument('--cosine', action='store_true',
                        help='using cosine annealing')
    parser.add_argument('--warm', action='store_true',
                        help='warm-up for large batch training')
    parser.add_argument('--trial', type=str, default='0',
                        help='id for recording multiple runs')

    opt = parser.parse_args()

    # set the path according to the environment
    if opt.dataset == 'imagenet100':
        opt.data_folder = '/cluster/tufts/hugheslab/datasets/ImageNet100/train/'
    elif opt.dataset == 'imagenet':
        opt.data_folder = '/cluster/tufts/hugheslab/datasets/ImageNet/train/'
    else:
        opt.data_folder = './datasets/'
    opt.model_path = './save/SupCon/{}_models'.format(opt.dataset)
    opt.tb_path = './save/SupCon/{}_tensorboard'.format(opt.dataset)

    iterations = opt.lr_decay_epochs.split(',')
    opt.lr_decay_epochs = list([])
    for it in iterations:
        opt.lr_decay_epochs.append(int(it))

    opt.model_name = 'SupCE_{}_{}_lr_{}_decay_{}_bsz_{}_trial_{}'.\
        format(opt.dataset, opt.model, opt.learning_rate, opt.weight_decay,
               opt.batch_size, opt.trial)

    if opt.cosine:
        opt.model_name = '{}_cosine'.format(opt.model_name)

    # warm-up for large-batch training,
    if opt.batch_size > 256:
        opt.warm = True
    if opt.warm:
        opt.model_name = '{}_warm'.format(opt.model_name)
        opt.warmup_from = 0.01
        opt.warm_epochs = 10
        if opt.cosine:
            eta_min = opt.learning_rate * (opt.lr_decay_rate ** 3)
            opt.warmup_to = eta_min + (opt.learning_rate - eta_min) * (
                1 + math.cos(math.pi * opt.warm_epochs / opt.epochs)) / 2
        else:
            opt.warmup_to = opt.learning_rate

    opt.tb_folder = os.path.join(opt.tb_path, opt.model_name)
    if not os.path.isdir(opt.tb_folder):
        os.makedirs(opt.tb_folder)

    opt.save_folder = os.path.join(opt.model_path, opt.model_name)
    if not os.path.isdir(opt.save_folder):
        os.makedirs(opt.save_folder)

    if opt.dataset == 'cifar10':
        opt.n_cls = 10
    elif opt.dataset == 'cifar100':
        opt.n_cls = 100
    else:
        raise ValueError('dataset not supported: {}'.format(opt.dataset))

    return opt


def set_loader(opt, contrast_trans=False, for_test=False):
    # dataset specific normalization
    if opt.dataset == 'cifar10':
        mean = (0.4914, 0.4822, 0.4465)
        std = (0.2023, 0.1994, 0.2010)
    elif opt.dataset == 'cifar100':
        mean = (0.5071, 0.4867, 0.4408)
        std = (0.2675, 0.2565, 0.2761)
    elif opt.dataset == 'cifar2':
        mean = (0.4977, 0.4605, 0.4160)
        std = (0.2537, 0.2481, 0.2535)
    elif opt.dataset == 'aircraft':
        mean = (0.4804, 0.5115, 0.5348)
        std = (0.1561, 0.1555, 0.1810)
    elif opt.dataset == 'cars':
        mean = (0.4707, 0.4601, 0.4549)
        std = (0.2319, 0.2318, 0.2373)
    elif opt.dataset == 'food101':
        mean = (0.5456, 0.4430, 0.3423)
        std = (0.2096, 0.2186, 0.2165)
    elif opt.dataset == 'pet':
        mean = (0.4786, 0.4462, 0.3962)
        std = (0.2073, 0.2042, 0.2053)
    elif opt.dataset == 'dtd':
        mean = (0.5301, 0.4734, 0.4243)
        std = (0.1250, 0.1246, 0.1211)
    elif opt.dataset == 'flowers':
        mean = (0.4312, 0.3786, 0.2944)
        std = (0.2385, 0.1858, 0.1986)
    elif opt.dataset == 'imagenet100' or opt.dataset == 'imagenet':
        mean = (0.485, 0.456, 0.406)
        std = (0.229, 0.224, 0.225)
    elif opt.dataset == 'path':
        mean = eval(opt.mean)
        std = eval(opt.std)
    else:
        raise ValueError('dataset not supported: {}'.format(opt.dataset))
    normalize = transforms.Normalize(mean=mean, std=std)

    # contrastive data transforms
    if contrast_trans:
        # first image is non-augmented, second is lightly augmented
        test_transform = DoubleTransform(
            transforms.Compose([
                transforms.Resize([opt.size, opt.size]),
                transforms.ToTensor(),
                normalize,
            ]),
            transforms.Compose([
                transforms.RandomResizedCrop(size=opt.size, scale=(0.2, 1.)),
                transforms.RandomHorizontalFlip(),
                transforms.ToTensor(),
                normalize,
            ]))
        if for_test:
            train_transform = test_transform
        else:
            # both images heavily augmented
            train_transform = TwoCropTransform(transforms.Compose([
                transforms.RandomResizedCrop(size=opt.size, scale=(0.2, 1.)),
                transforms.RandomHorizontalFlip(),
                transforms.RandomApply([
                    transforms.ColorJitter(0.4, 0.4, 0.4, 0.1)
                ], p=0.8),
                transforms.RandomGrayscale(p=0.2),
                transforms.ToTensor(),
                normalize,
            ]))
    # non-contrastive data transforms
    else:
        train_transform = transforms.Compose([
            transforms.RandomResizedCrop(size=opt.size, scale=(0.2, 1.)),
            transforms.RandomHorizontalFlip(),
            transforms.ToTensor(),
            normalize,
        ])
        test_transform = transforms.Compose([
            transforms.Resize([opt.size, opt.size]),
            transforms.ToTensor(),
            normalize,
        ])

    # construct dataset
    if opt.dataset == 'cifar10' or opt.dataset == 'cifar2':
        train_dataset = datasets.CIFAR10(root=opt.data_folder,
                                         transform=train_transform,
                                         download=True)
        test_dataset = datasets.CIFAR10(root=opt.data_folder,
                                        train=False,
                                        transform=test_transform)
        # produce CIFAR-2 (cats vs dogs) from CIFAR-10
        if opt.dataset == 'cifar2':
            classes = torch.Tensor([train_dataset.class_to_idx[name] for name in ["cat", "dog"]])
            cls_idx_map = {int(classes[0]): 0, int(classes[1]): 1}
            train_ind = torch.where(torch.isin(torch.Tensor(train_dataset.targets), classes))[0]
            train_dataset.target_transform = lambda x: cls_idx_map[x]
            train_dataset = SubsetWithTargets(train_dataset, train_ind)
            test_ind = torch.where(torch.isin(torch.Tensor(test_dataset.targets), classes))[0]
            test_dataset.target_transform = lambda x: cls_idx_map[x]
            test_dataset = SubsetWithTargets(test_dataset, test_ind)
    elif opt.dataset == 'cifar100':
        train_dataset = datasets.CIFAR100(root=opt.data_folder,
                                          transform=train_transform,
                                          download=True)
        test_dataset = datasets.CIFAR100(root=opt.data_folder,
                                         train=False,
                                         transform=test_transform)
    elif opt.dataset == 'aircraft':
        train_dataset = datasets.FGVCAircraft(root=opt.data_folder,
                                              split="trainval",
                                              transform=train_transform,
                                              download=True)
        train_dataset.targets = train_dataset._labels
        test_dataset = datasets.FGVCAircraft(root=opt.data_folder,
                                             split="test",
                                             transform=test_transform)
        test_dataset.targets = test_dataset._labels
    elif opt.dataset == 'cars':
        train_dataset = datasets.StanfordCars(root=opt.data_folder,
                                              transform=train_transform)
        train_dataset.targets = [label for _, label in train_dataset._samples]
        test_dataset = datasets.StanfordCars(root=opt.data_folder,
                                             split="test",
                                             transform=test_transform)
        test_dataset.targets = [label for _, label in test_dataset._samples]
    elif opt.dataset == 'food101':
        train_dataset = datasets.Food101(root=opt.data_folder,
                                         transform=train_transform,
                                         download=True)
        train_dataset.targets = train_dataset._labels
        test_dataset = datasets.Food101(root=opt.data_folder,
                                        split="test",
                                        transform=test_transform,
                                        download=True)
        test_dataset.targets = test_dataset._labels
    elif opt.dataset == 'pet':
        train_dataset = datasets.OxfordIIITPet(root=opt.data_folder,
                                               transform=train_transform,
                                               download=True)
        train_dataset.targets = train_dataset._labels
        test_dataset = datasets.OxfordIIITPet(root=opt.data_folder,
                                              split="test",
                                              transform=test_transform,
                                              download=True)
        test_dataset.targets = test_dataset._labels
    elif opt.dataset == 'dtd':
        train_dataset = datasets.DTD(root=opt.data_folder,
                                     transform=train_transform,
                                     download=True)
        train_dataset.targets = train_dataset._labels
        test_dataset = datasets.DTD(root=opt.data_folder,
                                    split="test",
                                    transform=test_transform,
                                    download=True)
        test_dataset.targets = test_dataset._labels
    elif opt.dataset == 'flowers':
        train_dataset = datasets.Flowers102(root=opt.data_folder,
                                            transform=train_transform,
                                            download=True)
        train_dataset.targets = train_dataset._labels
        test_dataset = datasets.Flowers102(root=opt.data_folder,
                                           split="test",
                                           transform=test_transform,
                                           download=True)
        test_dataset.targets = test_dataset._labels
    elif opt.dataset == 'imagenet100' or opt.dataset == 'imagenet' or opt.dataset == 'path':
        train_dataset = datasets.ImageFolder(root=opt.data_folder,
                                             transform=train_transform)
        test_dataset = datasets.ImageFolder(root=opt.data_folder + "../val/",
                                            transform=test_transform)
    else:
        raise ValueError(opt.dataset)

    if opt.valid_split == 0:
        valid_loader = None
    # split validation off of train
    else:
        old_train = train_dataset
        old_val = copy.deepcopy(train_dataset)
        old_val.transform = test_transform
        train_ind, valid_ind = train_test_split(
            range(len(old_train)), stratify=old_train.targets, test_size=opt.valid_split,
            random_state=12345)
        train_dataset = SubsetWithTargets(old_train, train_ind)
        valid_dataset = SubsetWithTargets(old_val, valid_ind)
        # construct validation data loader
        if "device" in opt:
            valid_loader = DataLoader(
                valid_dataset, num_workers=opt.num_workers, pin_memory=True,
                batch_size=opt.batch_size,
                sampler=DistributedSampler(valid_dataset) if "device" in opt else None)
        else:
            valid_loader = DataLoader(
                valid_dataset, num_workers=opt.num_workers, pin_memory=True,
                batch_sampler=sampler.my_sampler(valid_dataset, opt.batch_size))

    # construct train and test data loaders
    if "device" in opt:
        train_loader = DataLoader(
            train_dataset, num_workers=opt.num_workers, pin_memory=True,
            batch_size=opt.batch_size,
            sampler=DistributedSampler(train_dataset) if "device" in opt else None)
        test_loader = DataLoader(
            test_dataset, num_workers=opt.num_workers, pin_memory=True,
            batch_size=opt.batch_size,
            sampler=DistributedSampler(test_dataset) if "device" in opt else None)
    elif not for_test:
        train_loader = DataLoader(
            train_dataset, num_workers=opt.num_workers, pin_memory=True,
            batch_sampler=sampler.my_sampler(train_dataset, opt.batch_size))
        test_loader = DataLoader(
            test_dataset, num_workers=opt.num_workers, pin_memory=True,
            batch_sampler=sampler.my_sampler(test_dataset, opt.batch_size))
    else:
        train_loader = DataLoader(
            train_dataset, num_workers=opt.num_workers, pin_memory=True,
            batch_size=opt.batch_size)
        test_loader = DataLoader(
            test_dataset, num_workers=opt.num_workers, pin_memory=True,
            batch_size=opt.batch_size)

    # compute mean and STD used above (use test transform without normalization)
    # code from https://discuss.pytorch.org/t/about-normalization-using-pre-trained
    # -vgg16-networks/23560/5?u=kuzand
    # mean = 0.
    # std = 0.
    # nb_samples = 0.
    # for data, _ in train_loader:
    #     batch_samples = data.size(0)
    #     data = data.view(batch_samples, data.size(1), -1)
    #     mean += data.mean(2).sum(0)
    #     std += data.std(2).sum(0)
    #     nb_samples += batch_samples
    # mean /= nb_samples
    # std /= nb_samples
    # print(mean)
    # print(std)

    return train_loader, valid_loader, test_loader


def set_model(opt):
    model = SupCEResNet(name=opt.model, num_classes=opt.n_cls)
    criterion = torch.nn.CrossEntropyLoss()

    if torch.cuda.is_available():
        if torch.cuda.device_count() > 1:
            model = torch.nn.DataParallel(model)
        model = model.cuda()
        criterion = criterion.cuda()
        cudnn.benchmark = True

    return model, criterion


def train(train_loader, model, criterion, optimizer, epoch, opt):
    """one epoch training"""
    model.train()

    batch_time = AverageMeter()
    data_time = AverageMeter()
    losses = AverageMeter()
    top1 = AverageMeter()

    end = time.time()
    for idx, (images, labels) in enumerate(train_loader):
        data_time.update(time.time() - end)

        images = images.cuda(non_blocking=True)
        labels = labels.cuda(non_blocking=True)
        bsz = labels.shape[0]

        # warm-up learning rate
        warmup_learning_rate(opt, epoch, idx, len(train_loader), optimizer)

        # compute loss
        output = model(images)
        loss = criterion(output, labels)

        # update metric
        losses.update(loss.item(), bsz)
        acc1, acc5 = accuracy(output, labels, topk=(1, 5))
        top1.update(acc1[0], bsz)

        # SGD
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        # measure elapsed time
        batch_time.update(time.time() - end)
        end = time.time()

        # print info
        if (idx + 1) % opt.print_freq == 0:
            print('Train: [{0}][{1}/{2}]\t'
                  'BT {batch_time.val:.3f} ({batch_time.avg:.3f})\t'
                  'DT {data_time.val:.3f} ({data_time.avg:.3f})\t'
                  'loss {loss.val:.3f} ({loss.avg:.3f})\t'
                  'Acc@1 {top1.val:.3f} ({top1.avg:.3f})'.format(
                      epoch, idx + 1, len(train_loader), batch_time=batch_time,
                      data_time=data_time, loss=losses, top1=top1))
            sys.stdout.flush()

    return losses.avg, top1.avg


def validate(val_loader, model, criterion, opt):
    """validation"""
    model.eval()

    batch_time = AverageMeter()
    losses = AverageMeter()
    top1 = AverageMeter()

    with torch.no_grad():
        end = time.time()
        for idx, (images, labels) in enumerate(val_loader):
            images = images.float().cuda()
            labels = labels.cuda()
            bsz = labels.shape[0]

            # forward
            output = model(images)
            loss = criterion(output, labels)

            # update metric
            losses.update(loss.item(), bsz)
            acc1, acc5 = accuracy(output, labels, topk=(1, 5))
            top1.update(acc1[0], bsz)

            # measure elapsed time
            batch_time.update(time.time() - end)
            end = time.time()

            if idx % opt.print_freq == 0:
                print('Test: [{0}/{1}]\t'
                      'Time {batch_time.val:.3f} ({batch_time.avg:.3f})\t'
                      'Loss {loss.val:.4f} ({loss.avg:.4f})\t'
                      'Acc@1 {top1.val:.3f} ({top1.avg:.3f})'.format(
                          idx, len(val_loader), batch_time=batch_time,
                          loss=losses, top1=top1))

    print(' * Acc@1 {top1.avg:.3f}'.format(top1=top1))
    return losses.avg, top1.avg


def main():
    best_acc = 0
    opt = parse_option()

    # build data loader
    train_loader, val_loader = set_loader(opt)

    # build model and criterion
    model, criterion = set_model(opt)

    # build optimizer
    optimizer = set_optimizer(opt, model)

    # tensorboard
    logger = tb_logger.Logger(logdir=opt.tb_folder, flush_secs=2)

    # training routine
    for epoch in range(1, opt.epochs + 1):
        adjust_learning_rate(opt, optimizer, epoch)

        # train for one epoch
        time1 = time.time()
        loss, train_acc = train(train_loader, model,
                                criterion, optimizer, epoch, opt)
        time2 = time.time()
        print('epoch {}, total time {:.2f}'.format(epoch, time2 - time1))

        # tensorboard logger
        logger.log_value('train_loss', loss, epoch)
        logger.log_value('train_acc', train_acc, epoch)
        logger.log_value(
            'learning_rate', optimizer.param_groups[0]['lr'], epoch)

        # evaluation
        loss, val_acc = validate(val_loader, model, criterion, opt)
        logger.log_value('val_loss', loss, epoch)
        logger.log_value('val_acc', val_acc, epoch)

        if val_acc > best_acc:
            best_acc = val_acc

        if epoch % opt.save_freq == 0:
            save_file = os.path.join(
                opt.save_folder, 'ckpt_epoch_{epoch}.pth'.format(epoch=epoch))
            save_model(model, optimizer, opt, epoch, save_file)

    # save the last model
    save_file = os.path.join(
        opt.save_folder, 'last.pth')
    save_model(model, optimizer, opt, opt.epochs, save_file)

    print('best accuracy: {:.2f}'.format(best_acc))


if __name__ == '__main__':
    main()