main_supcon.py

from __future__ import print_function

import os
import sys
import argparse
import time
import math
import timm
import tensorboard_logger as tb_logger
import torch
import torch.backends.cudnn as cudnn
from torchvision import transforms, datasets
import datetime

from question_loader import (Question1Dataset, Question2Dataset, Question3Dataset, Question4Dataset,
                             Group2Dataset, Group3Dataset)
from util import TwoCropTransform, AverageMeter
from util import adjust_learning_rate, warmup_learning_rate
from util import set_optimizer, save_model
from timm.data import resolve_data_config
from timm.data.transforms_factory import create_transform
from networks.vit import SupConVit
from losses import SupConLoss
import wandb
from calculator import calculate1, calculate3, calculate4
from quiz_master import quiz3, quiz1, quiz2,  quiz4, group2, group4_question1, group4_question2
from encoder_models.vit import RankVit, ClusterVit
from encoder_models.resnet import Resnet
from tqdm import tqdm

try:
    import apex
    from apex import amp, optimizers
except ImportError:
    pass

MODELS = {
    'vit': ClusterVit,
    'rankvit': RankVit
}

CALCULATOR = {
    1: calculate1,
    2: calculate1,
    3: calculate3,
    4: calculate4,
    5: calculate1,
    7: calculate1,
    8: calculate4,
}

QUIZ_OPTIONS = {
    1: quiz1,
    2: quiz2,
    3: quiz3,
    4: quiz4,
    5: group2,
    7: group4_question1,
    8: group4_question2,
}

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=1,
                        help='save frequency')
    parser.add_argument('--batch_size', type=int, default=256,
                        help='batch_size')
    parser.add_argument('--num_workers', type=int, default=16,
                        help='num of workers to use')
    parser.add_argument('--epochs', type=int, default=1000,
                        help='number of training epochs')
    parser.add_argument('--use_parallel', type=bool, default=True,
                        help='Use parallel trainer')
    parser.add_argument('--gpus', type=str, default="2,3",
                        help='Use parallel trainer')

    # optimization
    parser.add_argument('--learning_rate', type=float, default=0.05,
                        help='learning rate')
    parser.add_argument('--lr_decay_epochs', type=str, default='700,800,900',
                        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', 'path'], help='dataset')
    parser.add_argument('--group_num', type=str, default='group1')
    parser.add_argument('--mean', type=str, help='mean of dataset in path in form of str tuple')
    parser.add_argument('--std', type=str, help='std of dataset in path in form of str tuple')
    parser.add_argument('--data_folder', type=str, default=None, help='path to custom dataset')
    parser.add_argument('--size', type=int, default=32, help='parameter for RandomResizedCrop')
    parser.add_argument('--pretrained', type=bool, default=False)
    parser.add_argument('--wandb_id', type=str, default="")
    parser.add_argument('--wandb', type=bool, default=False)
    parser.add_argument('--wandb_pn', type=str, default="")
    
    # method
    parser.add_argument('--method', type=str, default='SupCon',
                        choices=['SupCon', 'SimCLR'], help='choose method')

    # temperature
    parser.add_argument('--temp', type=float, default=0.07,
                        help='temperature for loss function')

    # other setting
    parser.add_argument('--cosine', action='store_true',
                        help='using cosine annealing')
    parser.add_argument('--syncBN', action='store_true',
                        help='using synchronized batch normalization')
    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()

    # check if dataset is path that passed required arguments
    if opt.dataset == 'path':
        assert opt.data_folder is not None \
            and opt.mean is not None \
            and opt.std is not None

    # set the path according to the environment
    if opt.data_folder is None:
        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 = '{}_{}_{}_lr_{}_decay_{}_bsz_{}_temp_{}_trial_{}'.\
        format(opt.method, opt.dataset, opt.model, opt.learning_rate,
               opt.weight_decay, opt.batch_size, opt.temp, 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 + opt.group_num)
    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)

    return opt


def set_loader(opt):
    # construct data loader
    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 == 'path':
        mean = eval(opt.mean)
        std = eval(opt.std)
    else:
        raise ValueError('dataset not supported: {}'.format(opt.dataset))

    config = resolve_data_config({}, model=timm.create_model("vit_base_patch16_224", pretrained=False, num_classes=0))
    config['std'] = std
    config['mean'] = mean
    train_transform = create_transform(**config, is_training=True)

    if opt.dataset == 'cifar10':
        train_dataset = datasets.CIFAR10(root=opt.data_folder,
                                         transform=TwoCropTransform(train_transform),
                                         download=True)
    elif opt.dataset == 'cifar100':
        train_dataset = datasets.CIFAR100(root=opt.data_folder,
                                          transform=TwoCropTransform(train_transform),
                                          download=True)
    elif opt.dataset == 'path':

        if opt.group_num == 'group1':
            train_dataset = [(Question1Dataset(root=f'{opt.data_folder}/question1',
                                           transform=TwoCropTransform(train_transform)), opt.batch_size // 6),
                         (Question2Dataset(root=f'{opt.data_folder}/question2',
                                           transform=TwoCropTransform(train_transform)), opt.batch_size // 3),
                         (Question3Dataset(root=f'{opt.data_folder}/question3',
                                           transform=TwoCropTransform(train_transform)), opt.batch_size // 6),
                         (Question4Dataset(root=f'{opt.data_folder}/question4',
                                           transform=TwoCropTransform(train_transform)), opt.batch_size // 5),
                         ]
        elif opt.group_num == 'group2':
            train_dataset = [(
        
                            Group2Dataset(root=f'{opt.data_folder}/quiz_1_v2',
                                            transform=TwoCropTransform(train_transform)), 10),
                            ]
        elif opt.group_num == 'group3':
            train_dataset = [(
                Group3Dataset(root=f'{opt.data_folder}/question1/train',
                              transform=TwoCropTransform(train_transform)), opt.batch_size // 2),
            ]
        elif opt.group_num == 'group4':
            train_dataset = [
                             (Question2Dataset(root=f'{opt.data_folder}/question1',
                                               transform=TwoCropTransform(train_transform)), opt.batch_size // 3),
                             (Question4Dataset(root=f'{opt.data_folder}/question2',
                                               transform=TwoCropTransform(train_transform)), opt.batch_size // 5),
                            ]

    else:
        raise ValueError(opt.dataset)

    train_sampler = None
    train_loaders = []
    for dataset, batch_size in train_dataset:
        train_loaders.append(torch.utils.data.DataLoader(
            dataset, batch_size=batch_size, shuffle=(train_sampler is None),
            num_workers=opt.num_workers, pin_memory=True, sampler=train_sampler))        

    return train_loaders

def set_loader_category(opt):
    # construct data loader
    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 == 'path':
        mean = eval(opt.mean)
        std = eval(opt.std)
    else:
        raise ValueError('dataset not supported: {}'.format(opt.dataset))

    config = resolve_data_config({}, model=timm.create_model("vit_base_patch16_224", pretrained=False, num_classes=0))
    config['std'] = std
    config['mean'] = mean
    train_transform = create_transform(**config, is_training=True)

    if opt.dataset == 'cifar10':
        train_dataset = datasets.CIFAR10(root=opt.data_folder,
                                         transform=TwoCropTransform(train_transform),
                                         download=True)
    elif opt.dataset == 'cifar100':
        train_dataset = datasets.CIFAR100(root=opt.data_folder,
                                          transform=TwoCropTransform(train_transform),
                                          download=True)
    elif opt.dataset == 'path':
        train_dataset = datasets.ImageFolder(root=opt.data_folder,
                                            transform=TwoCropTransform(train_transform))
    else:
        raise ValueError(opt.dataset)

    train_sampler = None
    train_loader = torch.utils.data.DataLoader(
        train_dataset, batch_size=opt.batch_size, shuffle=(train_sampler is None),
        num_workers=opt.num_workers, pin_memory=True, sampler=train_sampler)

    return train_loader


def set_model(opt):

    model = SupConVit(name=opt.model, pretrained=opt.pretrained)
    criterion = SupConLoss(temperature=opt.temp)

    # enable synchronized Batch Normalization
    if opt.syncBN:
        model = apex.parallel.convert_syncbn_model(model)

    if torch.cuda.is_available():
        if torch.cuda.device_count() > 1 and opt.use_parallel:
           model.encoder = torch.nn.DataParallel(model.encoder)
        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()

    end = time.time()
    idx = 0
    from random import shuffle
    shuffle(train_loader)
    for loader in train_loader:
        print(f"Using loader {loader.dataset.__class__}")
        for images in loader:
            num_cats = images[0].shape[0]
            num_pos = images[0].shape[1]
            labels = []
            # Mark which batch even came from. Every image within a single question sample is a positive pair
            for i in range(num_cats):
                labels = labels + [i + 1] * num_pos
            labels = torch.tensor(labels, dtype=int)
            # Reshape the images from 5D to 4D tensors.
            images = [images[0].reshape([images[0].shape[0] * images[0].shape[1],  *images[0].shape[2:]]),
                    images[1].reshape([images[0].shape[0] * images[0].shape[1],  *images[0].shape[2:]])]

            if labels.shape[0] != images[0].shape[0]:
                print(f'Skipping question {labels.shape[0]} != {images[0].shape[0]}')
                continue
            else:
                idx += 1
            data_time.update(time.time() - end)

            images = torch.cat([images[0], images[1]], dim=0)
            if torch.cuda.is_available():
                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
            features = model(images)
            f1, f2 = torch.split(features, [bsz, bsz], dim=0)
            features = torch.cat([f1.unsqueeze(1), f2.unsqueeze(1)], dim=1)
            if opt.method == 'SupCon':
                loss = criterion(features, labels)
            elif opt.method == 'SimCLR':
                loss = criterion(features)
            else:
                raise ValueError('contrastive method not supported: {}'.
                                format(opt.method))

            # update metric
            losses.update(loss.item(), 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})'.format(
                    epoch, idx + 1, len(train_loader), batch_time=batch_time,
                    data_time=data_time, loss=losses))
                sys.stdout.flush()
                
    return losses.avg


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

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

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

        images = torch.cat([images[0], images[1]], dim=0)
        if torch.cuda.is_available():
            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
        features = model(images)
        f1, f2 = torch.split(features, [bsz, bsz], dim=0)
        features = torch.cat([f1.unsqueeze(1), f2.unsqueeze(1)], dim=1)
        if opt.method == 'SimCLR':
            loss = criterion(features)
        else:
            raise ValueError('contrastive method not supported: {}'.
                             format(opt.method))

        # update metric
        losses.update(loss.item(), 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})'.format(
                   epoch, idx + 1, len(train_loader), batch_time=batch_time,
                   data_time=data_time, loss=losses))
            sys.stdout.flush()

    return losses.avg

def main():
    opt = parse_option()
    if opt.wandb:
        wandb.init(project=opt.wandb_pn, entity=opt.wandb_id)
    # build data loader
    if opt.method == 'SimCLR':
        train_loader = set_loader_category(opt)
    else:
        train_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()
        if opt.method == 'SimCLR':
            loss = train_category(train_loader, model, criterion, optimizer, epoch, opt)
        else:
            loss = train(train_loader, model, criterion, optimizer, epoch, opt)
        if opt.wandb:
            wandb.log({'loss':loss}, step=epoch)
        time2 = time.time()
        # tensorboard logger
        logger.log_value('loss', loss, epoch)
        logger.log_value('learning_rate', optimizer.param_groups[0]['lr'], epoch)
        print(f'epoch {epoch}, loss: {loss} ')
        weights_name = f'ckpt_{opt.method}_pretrained_{opt.pretrained}_{opt.group_num}_epoch_{epoch}.pth'

        if epoch % opt.save_freq == 0:

            weights_path = os.path.join(
                opt.save_folder,weights_name)
            save_model(model, optimizer, opt, epoch, weights_path)
            model_name = 'vit'
            valid_path = "valid"
            test_path = "test"

            root_dir = "/".join(opt.data_folder.split("/")[:-1]) 

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


if __name__ == '__main__':
    main()