utils/common.py

# Copyright (c) Alibaba, Inc. and its affiliates.
import os.path as osp
import numpy as np

import torch
from torch.utils.data import DataLoader, Subset
from torchvision.datasets import ImageFolder
import torchvision.transforms as transforms

from datasets.ImbalanceCIFAR import IMBALANCECIFAR10, IMBALANCECIFAR100
from datasets.ImbalanceImageNet import LT_Dataset
from datasets.tinyimages_300k import TinyImages
from models.base import BaseModel
from models.resnet import ResNet34, ResNet18
from models.resnet_imagenet import ResNet50# , ResNet18
from models.feat_pool import IDFeatPool

from utils.utils import TwoCropTransform, de_parallel


def build_dataset(args, ngpus_per_node, is_training=True):
    # get batch size:
    train_batch_size = args.batch_size if not args.ddp else int(args.batch_size/ngpus_per_node/args.num_nodes)
    num_workers = args.num_workers if not args.ddp else int((args.num_workers+ngpus_per_node)/ngpus_per_node)

    # data:
    if args.dataset in ['cifar10', 'cifar100']:
        train_transform = transforms.Compose([
            transforms.RandomCrop(32, padding=4),
            transforms.RandomHorizontalFlip(),
            transforms.RandomApply([
                transforms.ColorJitter(0.4, 0.4, 0.4, 0.1)
            ], p=0.8),
            transforms.RandomGrayscale(p=0.2),
            transforms.ToTensor(),
        ])
        test_transform = transforms.Compose([
            transforms.ToTensor(),
        ])
    elif args.dataset == 'imagenet' or args.dataset == 'waterbird':
        mean = [0.485, 0.456, 0.406]
        std = [0.229, 0.224, 0.225]
        train_transform = transforms.Compose([
            transforms.RandomResizedCrop(224, scale=(0.2, 1.0)),
            transforms.RandomHorizontalFlip(),
            transforms.RandomApply([
                transforms.ColorJitter(0.4, 0.4, 0.4, 0.1)
            ], p=0.8),
            transforms.RandomGrayscale(p=0.2),
            transforms.ToTensor(),
            transforms.Normalize(mean, std)
        ])
        test_transform = transforms.Compose([
            transforms.Resize(256),
            transforms.CenterCrop(224),
            transforms.ToTensor(),
            transforms.Normalize(mean, std)
        ])
    if args.dataset == 'cifar10':
        num_classes = 10
        train_set = IMBALANCECIFAR10(train=True, transform=TwoCropTransform(train_transform), imbalance_ratio=args.imbalance_ratio, root=args.data_root_path)
        test_set = IMBALANCECIFAR10(train=False, transform=test_transform, imbalance_ratio=args.imbalance_ratio, root=args.data_root_path)
    elif args.dataset == 'cifar100':
        num_classes = 100
        train_set = IMBALANCECIFAR100(train=True, transform=TwoCropTransform(train_transform), imbalance_ratio=args.imbalance_ratio, root=args.data_root_path)
        test_set = IMBALANCECIFAR100(train=False, transform=test_transform, imbalance_ratio=args.imbalance_ratio, root=args.data_root_path)
    elif args.dataset == 'imagenet':
        num_classes = args.id_class_number
        train_set = LT_Dataset(
            osp.join(args.data_root_path, 'imagenet'), './datasets/ImageNet_LT/ImageNet_LT_train.txt', transform=TwoCropTransform(train_transform), 
            subset_class_idx=np.arange(0,args.id_class_number))
        if is_training:
            test_set = LT_Dataset(
                osp.join(args.data_root_path, 'imagenet'), './datasets/ImageNet_LT/ImageNet_LT_val.txt', transform=test_transform,
                subset_class_idx=np.arange(0,args.id_class_number))
        else:
            test_set = ImageFolder(osp.join(args.data_root_path, 'imagenet', 'val'), transform=test_transform)
    elif args.dataset == 'waterbird':
        num_classes = 2
        train_set = ImageFolder(osp.join(args.data_root_path, 'waterbird_LT', 'train'), transform=TwoCropTransform(train_transform))
        setattr(train_set, 'img_num_per_cls', [363, 3699])
        if is_training:
            test_set = ImageFolder(osp.join(args.data_root_path, 'waterbird_LT', 'val'), transform=test_transform)
        else:
            test_set = ImageFolder(osp.join(args.data_root_path, 'waterbird_LT', 'test'), transform=test_transform)
    if args.ddp:
        train_sampler = torch.utils.data.distributed.DistributedSampler(train_set)
    else:
        train_sampler = None
    train_loader = DataLoader(train_set, batch_size=train_batch_size, shuffle=not args.ddp, num_workers=num_workers,
                                drop_last=True, pin_memory=True, sampler=train_sampler)
    test_loader = DataLoader(test_set, batch_size=args.test_batch_size, shuffle=False, num_workers=num_workers, 
                                drop_last=False, pin_memory=True)
    
    if is_training:
        if args.ood_aux_dataset in ['TinyImages', 'ImExtra', 'Texture']:
            if args.dataset in ['cifar10', 'cifar100']:
                ood_set = Subset(TinyImages(args.data_root_path, transform=train_transform), list(range(args.num_ood_samples)))
            elif args.dataset == 'imagenet':
                ood_set = ImageFolder(osp.join(args.data_root_path, 'imagenet/extra_1k'), transform=train_transform)
            elif args.dataset == 'waterbird':
                ood_set = ImageFolder(osp.join(args.data_root_path, 'texture'), transform=train_transform)
            else:
                raise NotImplementedError(args.dataset, args.ood_aux_dataset)
            
            if args.ddp:
                ood_sampler = torch.utils.data.distributed.DistributedSampler(ood_set)
            else:
                ood_sampler = None
            ood_loader = DataLoader(ood_set, batch_size=train_batch_size, shuffle=not args.ddp, num_workers=num_workers,
                                        drop_last=True, pin_memory=True, sampler=ood_sampler)
            ood_num = len(ood_set)
            print('Training on %s with %d images and %d validation images | %d OOD training images.' % (args.dataset, len(train_set), len(test_set), len(ood_set)))
        elif args.ood_aux_dataset in ['VOS', 'NPOS']:
            # sample_num = min(train_set.img_num_per_cls) * 10
            sample_num = max(train_set.img_num_per_cls)
            feat_dim = 512 if 'cifar' in args.dataset else 1024  # TODO: check
            mode = args.ood_aux_dataset
            device = 'cuda:0'
            ood_loader = IDFeatPool(num_classes, sample_num=sample_num, feat_dim=feat_dim, mode=mode, device=device)
            ood_num = num_classes * sample_num
        elif args.ood_aux_dataset in ['CIFAR']:
            if args.dataset == 'cifar10':
                dout = 'cifar100'
            elif args.dataset == 'cifar100':
                dout = 'cifar10'
            ood_set = ImageFolder(osp.join(args.data_root_path, f'SCOOD/data/images/{dout}/test'), transform=train_transform)
            if args.ddp:
                ood_sampler = torch.utils.data.distributed.DistributedSampler(ood_set)
            else:
                ood_sampler = None
            ood_loader = DataLoader(ood_set, batch_size=train_batch_size, shuffle=not args.ddp, num_workers=num_workers,
                                        drop_last=True, pin_memory=True, sampler=ood_sampler)
            ood_num = len(ood_set)
            print('Training on %s with %d images and %d validation images | %d OOD training images.' % (args.dataset, len(train_set), len(test_set), len(ood_set)))
        else:
            raise NotImplementedError(f'{args.dataset} v.s. {args.ood_aux_dataset}')

        img_num_per_cls_and_ood = np.array(train_set.img_num_per_cls + [ood_num])

        return num_classes, train_loader, test_loader, ood_loader, train_sampler, img_num_per_cls_and_ood
    else:

        img_num_per_cls_and_ood = np.array(train_set.img_num_per_cls + [args.num_ood_samples])

        return num_classes, test_loader, img_num_per_cls_and_ood


def build_plain_train_loader(args):  # for statistic during test
    if args.dataset in ['cifar10', 'cifar100']:
        train_transform = transforms.Compose([
            transforms.RandomCrop(32, padding=4),
            transforms.RandomHorizontalFlip(),
            transforms.RandomApply([
                transforms.ColorJitter(0.4, 0.4, 0.4, 0.1)
            ], p=0.8),
            transforms.RandomGrayscale(p=0.2),
            transforms.ToTensor(),
        ])
        test_transform = transforms.Compose([
            transforms.ToTensor(),
        ])
    elif args.dataset == 'imagenet':
        mean = [0.485, 0.456, 0.406]
        std = [0.229, 0.224, 0.225]
        train_transform = transforms.Compose([
            transforms.RandomResizedCrop(224, scale=(0.2, 1.0)),
            transforms.RandomHorizontalFlip(),
            transforms.RandomApply([
                transforms.ColorJitter(0.4, 0.4, 0.4, 0.1)
            ], p=0.8),
            transforms.RandomGrayscale(p=0.2),
            transforms.ToTensor(),
            transforms.Normalize(mean, std)
        ])
        test_transform = transforms.Compose([
            transforms.Resize(256),
            transforms.CenterCrop(224),
            transforms.ToTensor(),
            transforms.Normalize(mean, std)
        ])
    if args.dataset == 'cifar10':
        num_classes = 10
        train_set = IMBALANCECIFAR10(train=True, transform=test_transform, imbalance_ratio=args.imbalance_ratio, root=args.data_root_path)
    elif args.dataset == 'cifar100':
        num_classes = 100
        train_set = IMBALANCECIFAR100(train=True, transform=test_transform, imbalance_ratio=args.imbalance_ratio, root=args.data_root_path)
    elif args.dataset == 'imagenet':
        num_classes = args.id_class_number
        train_set = LT_Dataset(
            osp.join(args.data_root_path, 'imagenet'), './datasets/ImageNet_LT/ImageNet_LT_train.txt', transform=test_transform, 
            subset_class_idx=np.arange(0,args.id_class_number))
        
    train_loader = DataLoader(train_set, batch_size=1, shuffle=False, num_workers=4,
                                drop_last=True, pin_memory=True)
    
    return train_loader



def build_model(args, num_classes, device, gpu_id, return_features=False, is_training=True):
    # model:
    num_outputs = num_classes
    if any(x in args.ood_metric for x in ['bkg_c', 'bin_disc']):
        num_outputs += 1
    elif any(x in args.ood_metric for x in ['mc_disc']):
        num_outputs += 2
    # 'ResNet18', 'ResNet34', or 'ResNet50'
    model: BaseModel = eval(args.model)(num_classes=num_classes, num_outputs=num_outputs, return_features=return_features).to(device)
    if args.ddp:
        model = torch.nn.parallel.DistributedDataParallel(
            model, device_ids=[gpu_id], broadcast_buffers=False,
            find_unused_parameters=True
            )
    else:
        # model = torch.nn.DataParallel(model)
        pass
    # print('Model Done.')

    if is_training:
        # optimizer:
        if args.opt == 'adam':
            optimizer = torch.optim.Adam(model.parameters(), lr=args.lr, weight_decay=args.wd)
        elif args.opt == 'sgd':
            optimizer = torch.optim.SGD(model.parameters(), lr=args.lr, weight_decay=args.wd, momentum=args.momentum, nesterov=True)
        if args.decay == 'cos':
            scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=args.epochs)
        elif args.decay == 'multisteps':
            scheduler = torch.optim.lr_scheduler.MultiStepLR(optimizer, args.decay_epochs, gamma=0.1)
        # print('Optimizer Done.')

        return model, optimizer, scheduler, num_outputs
    else:
        return model, num_outputs


def build_prior(args, model, img_num_per_cls, num_classes, num_outputs, device):
    img_num_per_cls = torch.from_numpy(img_num_per_cls).to(device)
    if args.logit_adjust > 0:
        adjustments = img_num_per_cls / img_num_per_cls.sum()
        adjustments = args.logit_adjust * torch.log(adjustments + 1e-12)[None, :]
        if args.ood_metric in ['bkg_c'] and adjustments.shape[1] != num_outputs:
            placeholder = torch.zeros_like(adjustments[:, :num_outputs - num_classes])
            adjustments = torch.cat((adjustments, placeholder), dim=1)
    else:
        if args.ood_metric in ['bkg_c']:
            adjustments = torch.zeros((1, num_outputs), device=device)
        else:
            adjustments = torch.zeros((1, num_classes), device=device)

    return adjustments