imagenet_train_subclasses.py

# adapted from https://github.com/pytorch/examples/blob/main/imagenet/main.py
import argparse
import math
import os
import random
import shutil
import time
import warnings
from enum import Enum
from pathlib import Path

import numpy as np
import torch
import torch.backends.cudnn as cudnn
import torch.distributed as dist
import torch.multiprocessing as mp
import torch.nn as nn
import torch.nn.parallel
import torch.optim
import torch.utils.data
import torch.utils.data.distributed
import torchvision.datasets as datasets
import torchvision.transforms as transforms
from torch.optim.lr_scheduler import MultiStepLR
from torch.utils.data import Dataset, Sampler, Subset
from tqdm import tqdm

from openood.networks import ResNet18_224x224

SUB_CLASSES = {
    '10': [404, 751, 94, 281, 62, 207, 31, 949, 814, 569],
    'dogs': list(range(151, 251)),
}


def cosine_annealing(step, total_steps, lr_max, lr_min):
    return lr_min + (lr_max - lr_min) * 0.5 * \
                (1 + np.cos(step / total_steps * np.pi))


# ######################################
# sampler
class TwoSourceBatchSampler(Sampler):
    def __init__(self, real_inds, syn_inds, batch_size, real_ratio=0.5):
        # assert len(real_inds) == 1281167
        self.real_inds = real_inds
        self.syn_inds = syn_inds
        self.batch_size = batch_size
        self.real_batch_size = int(self.batch_size * real_ratio)
        self.syn_batch_size = self.batch_size - self.real_batch_size
        if real_ratio == 0:
            assert self.real_batch_size == 0
        elif real_ratio == 1:
            assert self.syn_batch_size == 0

        self.num_batches = math.ceil(len(self.real_inds) / self.batch_size)
        super().__init__(None)

    def __iter__(self):
        batch_counter = 0
        real_inds_shuffled = [
            self.real_inds[i] for i in torch.randperm(len(self.real_inds))
        ]
        syn_inds_shuffled = [
            self.syn_inds[i] for i in torch.randperm(len(self.syn_inds))
        ]

        real_offset = 0
        syn_offset = 0
        while batch_counter < self.num_batches:
            real_batch = real_inds_shuffled[
                real_offset:min(real_offset +
                                self.real_batch_size, len(real_inds_shuffled))]
            real_offset += self.real_batch_size

            syn_batch = syn_inds_shuffled[
                syn_offset:min(syn_offset +
                               self.syn_batch_size, len(syn_inds_shuffled))]
            syn_offset += self.syn_batch_size

            batch = real_batch + syn_batch
            np.random.shuffle(batch)
            yield batch
            batch_counter += 1

    def __len__(self):
        return self.num_batches


class DistributedSampler(Sampler):
    def __init__(self,
                 indices,
                 num_replicas=None,
                 rank=None,
                 shuffle=True,
                 seed=0,
                 drop_last=False):
        if num_replicas is None:
            if not dist.is_available():
                raise RuntimeError(
                    'Requires distributed package to be available')
            num_replicas = dist.get_world_size()
        if rank is None:
            if not dist.is_available():
                raise RuntimeError(
                    'Requires distributed package to be available')
            rank = dist.get_rank()
        if rank >= num_replicas or rank < 0:
            raise ValueError('Invalid rank {}, rank should be in the interval'
                             ' [0, {}]'.format(rank, num_replicas - 1))
        self.indices = indices
        self.num_replicas = num_replicas
        self.rank = rank
        self.epoch = 0
        self.drop_last = drop_last
        # If the dataset length is evenly divisible by # of replicas, then there
        # is no need to drop any data, since the dataset will be split equally.
        if self.drop_last and len(
                self.indices
        ) % self.num_replicas != 0:  # type: ignore[arg-type]
            # Split to nearest available length that is evenly divisible.
            # This is to ensure each rank receives the same amount of data when
            # using this Sampler.
            self.num_samples = math.ceil(
                (len(self.indices) - self.num_replicas) /
                self.num_replicas  # type: ignore[arg-type]
            )
        else:
            self.num_samples = math.ceil(
                len(self.indices) /
                self.num_replicas)  # type: ignore[arg-type]
        self.total_size = self.num_samples * self.num_replicas
        self.shuffle = shuffle
        self.seed = seed

    def __iter__(self):
        if self.shuffle:
            # deterministically shuffle based on epoch and seed
            g = torch.Generator()
            g.manual_seed(self.seed + self.epoch)
            indices = torch.tensor(self.indices)[torch.randperm(
                len(self.indices), generator=g)].tolist()
        else:
            indices = self.indices

        if not self.drop_last:
            # add extra samples to make it evenly divisible
            padding_size = self.total_size - len(indices)
            if padding_size <= len(indices):
                indices += indices[:padding_size]
            else:
                indices += (
                    indices *
                    math.ceil(padding_size / len(indices)))[:padding_size]
        else:
            # remove tail of data to make it evenly divisible.
            indices = indices[:self.total_size]
        assert len(indices) == self.total_size

        # subsample
        indices = indices[self.rank:self.total_size:self.num_replicas]
        assert len(indices) == self.num_samples

        return iter(indices)

    def __len__(self):
        return self.num_samples

    def set_epoch(self, epoch):
        self.epoch = epoch


class DistributedTwoSourceSampler(
        torch.utils.data.distributed.DistributedSampler):
    def __init__(self,
                 dataset,
                 num_replicas=None,
                 rank=None,
                 shuffle=True,
                 seed=0,
                 drop_last=False):

        if num_replicas is None:
            if not dist.is_available():
                raise RuntimeError(
                    'Requires distributed package to be available')
            num_replicas = dist.get_world_size()
        if rank is None:
            if not dist.is_available():
                raise RuntimeError(
                    'Requires distributed package to be available')
            rank = dist.get_rank()
        if rank >= num_replicas or rank < 0:
            raise ValueError('Invalid rank {}, rank should be in the interval'
                             ' [0, {}]'.format(rank, num_replicas - 1))

        self.real_sampler = DistributedSampler(dataset.real_inds, num_replicas,
                                               rank, shuffle, seed, drop_last)
        self.syn_sampler = DistributedSampler(dataset.syn_inds, num_replicas,
                                              rank, shuffle, seed, drop_last)

        self.num_replicas = num_replicas
        self.rank = rank
        self.drop_last = drop_last
        self.num_samples = self.real_sampler.num_samples + self.syn_sampler.num_samples
        self.total_size = self.real_sampler.total_size + self.syn_sampler.total_size
        assert self.total_size == self.num_samples * self.num_replicas
        self.shuffle = shuffle
        self.seed = seed

    def __iter__(self):
        return [
            list(self.real_sampler.__iter__()),
            list(self.syn_sampler.__iter__())
        ]

    def __len__(self):
        return self.num_samples

    def set_epoch(self, epoch):
        self.real_sampler.set_epoch(epoch)
        self.syn_sampler.set_epoch(epoch)


class DistributedTwoSourceBatchSampler(DistributedTwoSourceSampler):
    def __init__(self,
                 dataset,
                 num_replicas=None,
                 rank=None,
                 shuffle=True,
                 seed=0,
                 drop_last=False,
                 batch_size=256,
                 real_ratio=0.9):
        super().__init__(dataset=dataset,
                         num_replicas=num_replicas,
                         rank=rank,
                         shuffle=shuffle,
                         seed=seed,
                         drop_last=drop_last)
        self.batch_size = batch_size
        self.real_ratio = real_ratio

    def __iter__(self):
        real_inds, syn_inds = super().__iter__()
        batch_sampler = TwoSourceBatchSampler(real_inds, syn_inds,
                                              self.batch_size, self.real_ratio)
        return iter(batch_sampler)

    def __len__(self):
        return math.ceil(self.real_sampler.num_samples / self.batch_size)


# ######################################

parser = argparse.ArgumentParser(description='PyTorch ImageNet Training')
parser.add_argument('--data',
                    metavar='DIR',
                    nargs='?',
                    default='imagenet',
                    help='path to dataset (default: imagenet)')
parser.add_argument('--sub-classes-split',
                    type=str,
                    default='10',
                    choices=['10', 'dogs'])
parser.add_argument('--extra-data',
                    type=str,
                    default=None,
                    help='path to extra data folder')
parser.add_argument('--real-ratio',
                    type=float,
                    default=0.9,
                    help='ratio of real data within a batch')
parser.add_argument('-j',
                    '--workers',
                    default=4,
                    type=int,
                    metavar='N',
                    help='number of data loading workers (default: 4)')
parser.add_argument('--epochs',
                    default=90,
                    type=int,
                    metavar='N',
                    help='number of total epochs to run')
parser.add_argument('--start-epoch',
                    default=0,
                    type=int,
                    metavar='N',
                    help='manual epoch number (useful on restarts)')
parser.add_argument('-b',
                    '--batch-size',
                    default=256,
                    type=int,
                    metavar='N',
                    help='mini-batch size (default: 256), this is the total '
                    'batch size of all GPUs on the current node when '
                    'using Data Parallel or Distributed Data Parallel')
parser.add_argument('--lr',
                    '--learning-rate',
                    default=0.1,
                    type=float,
                    metavar='LR',
                    help='initial learning rate',
                    dest='lr')
parser.add_argument('--momentum',
                    default=0.9,
                    type=float,
                    metavar='M',
                    help='momentum')
parser.add_argument('--wd',
                    '--weight-decay',
                    default=1e-4,
                    type=float,
                    metavar='W',
                    help='weight decay (default: 1e-4)',
                    dest='weight_decay')
parser.add_argument('-p',
                    '--print-freq',
                    default=10,
                    type=int,
                    metavar='N',
                    help='print frequency (default: 10)')
parser.add_argument('--resume',
                    default='',
                    type=str,
                    metavar='PATH',
                    help='path to latest checkpoint (default: none)')
parser.add_argument('-e',
                    '--evaluate',
                    dest='evaluate',
                    action='store_true',
                    help='evaluate model on validation set')
parser.add_argument('--world-size',
                    default=-1,
                    type=int,
                    help='number of nodes for distributed training')
parser.add_argument('--rank',
                    default=-1,
                    type=int,
                    help='node rank for distributed training')
parser.add_argument('--dist-url',
                    default='tcp://224.66.41.62:23456',
                    type=str,
                    help='url used to set up distributed training')
parser.add_argument('--dist-backend',
                    default='nccl',
                    type=str,
                    help='distributed backend')
parser.add_argument('--seed',
                    default=None,
                    type=int,
                    help='seed for initializing training. ')
parser.add_argument('--gpu', default=None, type=int, help='GPU id to use.')
parser.add_argument('--multiprocessing-distributed',
                    action='store_true',
                    help='Use multi-processing distributed training to launch '
                    'N processes per node, which has N GPUs. This is the '
                    'fastest way to use PyTorch for either single node or '
                    'multi node data parallel training')
parser.add_argument('--dummy',
                    action='store_true',
                    help='use fake data to benchmark')

best_acc1 = 0
best_epoch = -1


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

    if args.seed is not None:
        random.seed(args.seed)
        torch.manual_seed(args.seed)

    if args.gpu is not None:
        warnings.warn('You have chosen a specific GPU. This will completely '
                      'disable data parallelism.')

    if args.dist_url == 'env://' and args.world_size == -1:
        args.world_size = int(os.environ['WORLD_SIZE'])

    args.distributed = args.world_size > 1 or args.multiprocessing_distributed

    if torch.cuda.is_available():
        args.ngpus_per_node = torch.cuda.device_count()
    else:
        args.ngpus_per_node = 1
    if args.multiprocessing_distributed:
        # Since we have ngpus_per_node processes per node, the total world_size
        # needs to be adjusted accordingly
        args.world_size = args.ngpus_per_node * args.world_size
        # Use torch.multiprocessing.spawn to launch distributed processes: the
        # main_worker process function
        mp.spawn(main_worker,
                 nprocs=args.ngpus_per_node,
                 args=(args.ngpus_per_node, args))
    else:
        # Simply call main_worker function
        main_worker(args.gpu, args.ngpus_per_node, args)


def main_worker(gpu, ngpus_per_node, args):
    global best_acc1
    global best_epoch
    args.gpu = gpu

    if args.gpu is not None:
        print('Use GPU: {} for training'.format(args.gpu))

    if args.distributed:
        if args.dist_url == 'env://' and args.rank == -1:
            args.rank = int(os.environ['RANK'])
        if args.multiprocessing_distributed:
            # For multiprocessing distributed training, rank needs to be the
            # global rank among all the processes
            args.rank = args.rank * args.ngpus_per_node + gpu
        dist.init_process_group(backend=args.dist_backend,
                                init_method=args.dist_url,
                                world_size=args.world_size,
                                rank=args.rank)

    sub_classes = SUB_CLASSES[args.sub_classes_split]
    data_name = f'imagenet-{args.sub_classes_split}'
    expr_name = f'{data_name}_resnet18_224x224' \
        f'_e{args.epochs}_bs{args.batch_size}_lr{args.lr:g}'
    if args.extra_data is None:
        expr_name += '_default'
    else:
        #extra_data_name = '-'.join(
        #    args.extra_data.split('/')[-1].split('-')[2:])
        expr_name += \
            f'_extradata-biggan-ratio{args.real_ratio:.2f}'
    expr_name += f'/s{args.seed}'

    output_dir = os.path.join('results', expr_name)
    if not args.multiprocessing_distributed or (
            args.multiprocessing_distributed
            and args.rank % ngpus_per_node == 0):
        if not os.path.exists(output_dir):
            os.makedirs(output_dir)

    # create model
    print("=> creating model 'ResNet18'")
    model = ResNet18_224x224(num_classes=len(sub_classes))

    if not torch.cuda.is_available() and not torch.backends.mps.is_available():
        print('using CPU, this will be slow')
    elif args.distributed:
        # For multiprocessing distributed, DistributedDataParallel constructor
        # should always set the single device scope, otherwise,
        # DistributedDataParallel will use all available devices.
        if torch.cuda.is_available():
            if args.gpu is not None:
                torch.cuda.set_device(args.gpu)
                model.cuda(args.gpu)
                # When using a single GPU per process and per
                # DistributedDataParallel, we need to divide the batch size
                # ourselves based on the total number of GPUs of the current node.
                args.batch_size = int(args.batch_size / ngpus_per_node)
                args.workers = int(
                    (args.workers + ngpus_per_node - 1) / ngpus_per_node)
                model = torch.nn.parallel.DistributedDataParallel(
                    model, device_ids=[args.gpu])
            else:
                model.cuda()
                # DistributedDataParallel will divide and allocate batch_size to all
                # available GPUs if device_ids are not set
                model = torch.nn.parallel.DistributedDataParallel(model)
    elif args.gpu is not None and torch.cuda.is_available():
        torch.cuda.set_device(args.gpu)
        model = model.cuda(args.gpu)
    elif torch.backends.mps.is_available():
        device = torch.device('mps')
        model = model.to(device)
    else:
        # DataParallel will divide and allocate batch_size to all available GPUs
        if args.arch.startswith('alexnet') or args.arch.startswith('vgg'):
            model.features = torch.nn.DataParallel(model.features)
            model.cuda()
        else:
            model = torch.nn.DataParallel(model).cuda()

    if torch.cuda.is_available():
        if args.gpu:
            device = torch.device('cuda:{}'.format(args.gpu))
        else:
            device = torch.device('cuda')
    elif torch.backends.mps.is_available():
        device = torch.device('mps')
    else:
        device = torch.device('cpu')

    # Data loading code
    if args.dummy:
        print('=> Dummy data is used!')
        train_dataset = datasets.FakeData(1281167, (3, 224, 224), 1000,
                                          transforms.ToTensor())
        val_dataset = datasets.FakeData(50000, (3, 224, 224), 1000,
                                        transforms.ToTensor())
    else:
        traindir = os.path.join(args.data, 'train')
        valdir = os.path.join(args.data, 'val')
        normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406],
                                         std=[0.229, 0.224, 0.225])

        train_dataset = datasets.ImageFolder(
            traindir,
            transforms.Compose([
                transforms.RandomResizedCrop(224),
                transforms.RandomHorizontalFlip(),
                transforms.ToTensor(),
                normalize,
            ]))

        val_dataset = datasets.ImageFolder(
            valdir,
            transforms.Compose([
                transforms.Resize(256),
                transforms.CenterCrop(224),
                transforms.ToTensor(),
                normalize,
            ]))

        LABEL_MAPPING = {v: i for i, v in enumerate(sorted(sub_classes))}

        def get_subset_dataset(dset):
            indices = np.in1d(dset.targets, sub_classes).nonzero()[0]
            new_samples = []
            new_targets = []
            for i in indices:
                new_samples.append(
                    (dset.samples[i][0], LABEL_MAPPING[dset.samples[i][1]]))
                new_targets.append(LABEL_MAPPING[dset.samples[i][1]])
            dset.samples = new_samples
            dset.targets = new_targets
            return dset

        train_dataset = get_subset_dataset(train_dataset)
        val_dataset = get_subset_dataset(val_dataset)

        if args.extra_data:
            extra_train_dataset = datasets.ImageFolder(args.extra_data,
                                                       transform=None)
            # sanity check
            for k, v in extra_train_dataset.class_to_idx.items():
                assert LABEL_MAPPING[int(k)] == v

            # include extra data into train_dataset
            train_dataset.real_inds = list(range(len(train_dataset)))
            train_dataset.samples.extend(extra_train_dataset.samples)
            train_dataset.targets.extend(extra_train_dataset.targets)
            train_dataset.syn_inds = list(
                set(range(len(train_dataset))) - set(train_dataset.real_inds))

    if args.distributed:
        if args.extra_data is None:
            train_sampler = torch.utils.data.distributed.DistributedSampler(
                train_dataset)
        else:
            train_batch_sampler = DistributedTwoSourceBatchSampler(
                train_dataset,
                batch_size=args.batch_size,
                real_ratio=args.real_ratio)

        val_sampler = torch.utils.data.distributed.DistributedSampler(
            val_dataset, shuffle=False, drop_last=True)
    else:
        train_sampler = None
        val_sampler = None

    if args.extra_data:
        train_loader = torch.utils.data.DataLoader(
            train_dataset,
            batch_sampler=train_batch_sampler,
            num_workers=args.workers,
            pin_memory=True)
    else:
        train_loader = torch.utils.data.DataLoader(
            train_dataset,
            batch_size=args.batch_size,
            shuffle=(train_sampler is None),
            num_workers=args.workers,
            pin_memory=True,
            sampler=train_sampler)

    val_loader = torch.utils.data.DataLoader(val_dataset,
                                             batch_size=args.batch_size,
                                             shuffle=False,
                                             num_workers=args.workers,
                                             pin_memory=True,
                                             sampler=val_sampler)

    # define loss function (criterion), optimizer, and learning rate scheduler
    criterion = nn.CrossEntropyLoss().to(device)

    optimizer = torch.optim.SGD(model.parameters(),
                                args.lr,
                                momentum=args.momentum,
                                weight_decay=args.weight_decay)

    scheduler = torch.optim.lr_scheduler.LambdaLR(
        optimizer,
        lr_lambda=lambda step: cosine_annealing(
            step,
            args.epochs * len(train_loader),
            1,
            1e-6 / args.lr,
        ),
    )

    # optionally resume from a checkpoint
    if args.resume:
        if os.path.isfile(args.resume):
            print("=> loading checkpoint '{}'".format(args.resume))
            if args.gpu is None:
                checkpoint = torch.load(args.resume)
            elif torch.cuda.is_available():
                # Map model to be loaded to specified single gpu.
                loc = 'cuda:{}'.format(args.gpu)
                checkpoint = torch.load(args.resume, map_location=loc)
            args.start_epoch = checkpoint['epoch']
            best_acc1 = checkpoint['best_acc1']
            if args.gpu is not None:
                # best_acc1 may be from a checkpoint from a different GPU
                best_acc1 = best_acc1.to(args.gpu)
            model.load_state_dict(checkpoint['state_dict'])
            optimizer.load_state_dict(checkpoint['optimizer'])
            scheduler.load_state_dict(checkpoint['scheduler'])
            print("=> loaded checkpoint '{}' (epoch {})".format(
                args.resume, checkpoint['epoch']))
        else:
            print("=> no checkpoint found at '{}'".format(args.resume))

    if args.evaluate:
        validate(val_loader, model, criterion, args)
        return

    for epoch in tqdm(range(args.start_epoch, args.epochs),
                      total=args.epochs,
                      desc='Epochs',
                      position=0,
                      leave=True,
                      disable=(args.multiprocessing_distributed
                               and args.rank % ngpus_per_node != 0)):
        if args.distributed:
            if args.extra_data:
                train_batch_sampler.set_epoch(epoch)
            else:
                train_sampler.set_epoch(epoch)

        # train for one epoch
        train(train_loader, model, criterion, optimizer, scheduler, epoch,
              device, args)

        # evaluate on validation set
        acc1 = validate(val_loader, model, criterion, args)

        # scheduler.step()

        # remember best acc@1 and save checkpoint
        is_best = acc1 > best_acc1

        if is_best:
            if not args.multiprocessing_distributed or (
                    args.multiprocessing_distributed
                    and args.rank % ngpus_per_node == 0):

                old_fname = 'best_epoch{}_acc{:.2f}.ckpt'.format(
                    best_epoch, best_acc1)
                old_pth = os.path.join(output_dir, old_fname)
                Path(old_pth).unlink(missing_ok=True)

                save_fname = 'best_epoch{}_acc{:.2f}.ckpt'.format(epoch, acc1)
                save_pth = os.path.join(output_dir, save_fname)

                try:
                    torch.save(model.module.state_dict(), save_pth)
                except:
                    torch.save(model.state_dict(), save_pth)

        best_acc1 = max(acc1, best_acc1)
        if is_best:
            best_epoch = epoch


def train(train_loader, model, criterion, optimizer, scheduler, epoch, device,
          args):
    batch_time = AverageMeter('Time', ':6.3f')
    data_time = AverageMeter('Data', ':6.3f')
    losses = AverageMeter('Loss', ':.4e')
    top1 = AverageMeter('Acc@1', ':6.2f')
    top5 = AverageMeter('Acc@5', ':6.2f')
    progress = ProgressMeter(len(train_loader),
                             [batch_time, data_time, losses, top1, top5],
                             prefix='Epoch: [{}]'.format(epoch))

    # switch to train mode
    model.train()

    end = time.time()
    for i, (images, target) in tqdm(
            enumerate(train_loader),
            total=len(train_loader),
            desc='Train',
            position=1,
            leave=False,
            disable=(args.multiprocessing_distributed
                     and args.rank % args.ngpus_per_node != 0)):
        # measure data loading time
        data_time.update(time.time() - end)

        # move data to the same device as model
        images = images.to(device, non_blocking=True)
        target = target.to(device, non_blocking=True)

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

        # measure accuracy and record loss
        acc1, acc5 = accuracy(output, target, topk=(1, 5))
        losses.update(loss.item(), images.size(0))
        top1.update(acc1[0], images.size(0))
        top5.update(acc5[0], images.size(0))

        # compute gradient and do SGD step
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        scheduler.step()

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

        # if not args.multiprocessing_distributed or (
        #    args.multiprocessing_distributed
        #    and args.rank % args.ngpus_per_node == 0):
        #    if i % args.print_freq == 0:
        #        progress.display(i + 1)


def validate(val_loader, model, criterion, args):
    def run_validate(loader, base_progress=0):
        with torch.no_grad():
            end = time.time()
            for i, (images, target) in enumerate(loader):
                i = base_progress + i
                if args.gpu is not None and torch.cuda.is_available():
                    images = images.cuda(args.gpu, non_blocking=True)
                # if torch.backends.mps.is_available():
                #    images = images.to('mps')
                #    target = target.to('mps')
                if torch.cuda.is_available():
                    target = target.cuda(args.gpu, non_blocking=True)

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

                # measure accuracy and record loss
                acc1, acc5 = accuracy(output, target, topk=(1, 5))
                losses.update(loss.item(), images.size(0))
                top1.update(acc1[0], images.size(0))
                top5.update(acc5[0], images.size(0))

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

                if not args.multiprocessing_distributed or (
                        args.multiprocessing_distributed
                        and args.rank % args.ngpus_per_node == 0):
                    if i % args.print_freq == 0:
                        progress.display(i + 1)

    batch_time = AverageMeter('Time', ':6.3f', Summary.NONE)
    losses = AverageMeter('Loss', ':.4e', Summary.NONE)
    top1 = AverageMeter('Acc@1', ':6.2f', Summary.AVERAGE)
    top5 = AverageMeter('Acc@5', ':6.2f', Summary.AVERAGE)
    progress = ProgressMeter(len(val_loader) + (
        args.distributed and
        (len(val_loader.sampler) * args.world_size < len(val_loader.dataset))),
                             [batch_time, losses, top1, top5],
                             prefix='Test: ')

    # switch to evaluate mode
    model.eval()

    run_validate(val_loader)
    if args.distributed:
        top1.all_reduce()
        top5.all_reduce()

    if args.distributed and (len(val_loader.sampler) * args.world_size < len(
            val_loader.dataset)):
        aux_val_dataset = Subset(
            val_loader.dataset,
            range(
                len(val_loader.sampler) * args.world_size,
                len(val_loader.dataset)))
        aux_val_loader = torch.utils.data.DataLoader(
            aux_val_dataset,
            batch_size=args.batch_size,
            shuffle=False,
            num_workers=args.workers,
            pin_memory=True)
        run_validate(aux_val_loader, len(val_loader))

    if not args.multiprocessing_distributed or (
            args.multiprocessing_distributed
            and args.rank % args.ngpus_per_node == 0):
        progress.display_summary()

    return top1.avg


def save_checkpoint(state, is_best, filename='checkpoint.pth.tar'):
    torch.save(state, filename)
    if is_best:
        shutil.copyfile(filename, 'model_best.pth.tar')


class Summary(Enum):
    NONE = 0
    AVERAGE = 1
    SUM = 2
    COUNT = 3


class AverageMeter(object):
    """Computes and stores the average and current value."""
    def __init__(self, name, fmt=':f', summary_type=Summary.AVERAGE):
        self.name = name
        self.fmt = fmt
        self.summary_type = summary_type
        self.reset()

    def reset(self):
        self.val = 0
        self.avg = 0
        self.sum = 0
        self.count = 0

    def update(self, val, n=1):
        self.val = val
        self.sum += val * n
        self.count += n
        self.avg = self.sum / self.count

    def all_reduce(self):
        if torch.cuda.is_available():
            device = torch.device('cuda')
        elif torch.backends.mps.is_available():
            device = torch.device('mps')
        else:
            device = torch.device('cpu')
        total = torch.tensor([self.sum, self.count],
                             dtype=torch.float32,
                             device=device)
        dist.all_reduce(total, dist.ReduceOp.SUM, async_op=False)
        self.sum, self.count = total.tolist()
        self.avg = self.sum / self.count

    def __str__(self):
        fmtstr = '{name} {val' + self.fmt + '} ({avg' + self.fmt + '})'
        return fmtstr.format(**self.__dict__)

    def summary(self):
        fmtstr = ''
        if self.summary_type is Summary.NONE:
            fmtstr = ''
        elif self.summary_type is Summary.AVERAGE:
            fmtstr = '{name} {avg:.3f}'
        elif self.summary_type is Summary.SUM:
            fmtstr = '{name} {sum:.3f}'
        elif self.summary_type is Summary.COUNT:
            fmtstr = '{name} {count:.3f}'
        else:
            raise ValueError('invalid summary type %r' % self.summary_type)

        return fmtstr.format(**self.__dict__)


class ProgressMeter(object):
    def __init__(self, num_batches, meters, prefix=''):
        self.batch_fmtstr = self._get_batch_fmtstr(num_batches)
        self.meters = meters
        self.prefix = prefix

    def display(self, batch):
        entries = [self.prefix + self.batch_fmtstr.format(batch)]
        entries += [str(meter) for meter in self.meters]
        tqdm.write('\t'.join(entries))

    def display_summary(self):
        entries = [' *']
        entries += [meter.summary() for meter in self.meters]
        tqdm.write(' '.join(entries))

    def _get_batch_fmtstr(self, num_batches):
        num_digits = len(str(num_batches // 1))
        fmt = '{:' + str(num_digits) + 'd}'
        return '[' + fmt + '/' + fmt.format(num_batches) + ']'


def accuracy(output, target, topk=(1, )):
    """Computes the accuracy over the k top predictions for the specified
    values of k."""
    with torch.no_grad():
        maxk = max(topk)
        batch_size = target.size(0)

        _, pred = output.topk(maxk, 1, True, True)
        pred = pred.t()
        correct = pred.eq(target.view(1, -1).expand_as(pred))

        res = []
        for k in topk:
            correct_k = correct[:k].reshape(-1).float().sum(0, keepdim=True)
            res.append(correct_k.mul_(100.0 / batch_size))
        return res


if __name__ == '__main__':
    main()