main.py

# --------------------------------------------------------
# InternImage
# Copyright (c) 2022 OpenGVLab
# Licensed under The MIT License [see LICENSE for details]
# --------------------------------------------------------

import argparse
import datetime
import os
import random
import subprocess
import time
from contextlib import suppress

import numpy as np
import torch
import torch.backends.cudnn as cudnn
import torch.distributed as dist
from config import get_config
from dataset import build_loader
from ddp_hooks import fp16_compress_hook
from logger import create_logger
from lr_scheduler import build_scheduler
from models import build_model
from optimizer import build_optimizer
from timm.loss import LabelSmoothingCrossEntropy, SoftTargetCrossEntropy
from timm.utils import ApexScaler, AverageMeter, ModelEma, accuracy
from utils import MyAverageMeter
from utils import NativeScalerWithGradNormCount as NativeScaler
from utils import (auto_resume_helper, get_grad_norm, load_checkpoint,
                   load_ema_checkpoint, load_pretrained, reduce_tensor,
                   save_checkpoint)

try:
    from apex import amp
    has_apex = True
except ImportError:
    has_apex = False
# assert not has_apex, "The code is modified based on native amp"

has_native_amp = False
try:
    if getattr(torch.cuda.amp, 'autocast') is not None:
        has_native_amp = True
except AttributeError:
    pass

TORCH_VERSION = tuple(int(x) for x in torch.__version__.split('.')[:2])


def obsolete_torch_version(torch_version, version_threshold):
    return torch_version == 'parrots' or torch_version <= version_threshold


def parse_option():
    parser = argparse.ArgumentParser(
        'InternImage training and evaluation script', add_help=False)
    parser.add_argument('--cfg',
                        type=str,
                        required=True,
                        metavar='FILE',
                        help='path to config file')
    parser.add_argument(
        '--opts',
        help="Modify config options by adding 'KEY VALUE' pairs. ",
        default=None,
        nargs='+')

    # easy config modification
    parser.add_argument('--batch-size',
                        type=int,
                        help='batch size for single GPU')
    parser.add_argument('--dataset',
                        type=str,
                        help='dataset name',
                        default=None)
    parser.add_argument('--data-path', type=str, help='path to dataset')
    parser.add_argument('--zip',
                        action='store_true',
                        help='use zipped dataset instead of folder dataset')
    parser.add_argument(
        '--cache-mode',
        type=str,
        default='part',
        choices=['no', 'full', 'part'],
        help='no: no cache, '
        'full: cache all data, '
        'part: sharding the dataset into nonoverlapping pieces and only cache one piece'
    )
    parser.add_argument(
        '--pretrained',
        help=
        'pretrained weight from checkpoint, could be imagenet22k pretrained weight'
    )
    parser.add_argument('--resume', help='resume from checkpoint')
    parser.add_argument('--accumulation-steps',
                        type=int,
                        default=1,
                        help='gradient accumulation steps')
    parser.add_argument(
        '--use-checkpoint',
        action='store_true',
        help='whether to use gradient checkpointing to save memory')
    parser.add_argument(
        '--amp-opt-level',
        type=str,
        default='O1',
        choices=['O0', 'O1', 'O2'],
        help='mixed precision opt level, if O0, no amp is used')
    parser.add_argument(
        '--output',
        default='output',
        type=str,
        metavar='PATH',
        help=
        'root of output folder, the full path is <output>/<model_name>/<tag> (default: output)'
    )
    parser.add_argument('--tag', help='tag of experiment')
    parser.add_argument('--eval',
                        action='store_true',
                        help='Perform evaluation only')
    parser.add_argument('--throughput',
                        action='store_true',
                        help='Test throughput only')
    parser.add_argument('--save-ckpt-num', default=1, type=int)
    parser.add_argument(
        '--use-zero',
        action='store_true',
        help='whether to use ZeroRedundancyOptimizer (ZeRO) to save memory')

    # distributed training
    parser.add_argument('--local-rank',
                        type=int,
                        required=True,
                        help='local rank for DistributedDataParallel')

    args, unparsed = parser.parse_known_args()
    config = get_config(args)

    return args, config


@torch.no_grad()
def throughput(data_loader, model, logger):
    model.eval()

    for idx, (images, _) in enumerate(data_loader):
        if type(images) == list:
            images = [item.cuda(non_blocking=True) for item in images]
        else:
            images = images.cuda(non_blocking=True)
        batch_size = images.shape[0]
        for i in range(50):
            model(images)
        torch.cuda.synchronize()
        logger.info(f'throughput averaged with 30 times')
        tic1 = time.time()
        for i in range(30):
            model(images)
        torch.cuda.synchronize()
        tic2 = time.time()
        logger.info(
            f'batch_size {batch_size} throughput {30 * batch_size / (tic2 - tic1)}'
        )
        return


def main(config):
    # prepare data loaders
    dataset_train, dataset_val, dataset_test, data_loader_train, \
        data_loader_val, data_loader_test, mixup_fn = build_loader(config)

    # build runner
    logger.info(f'Creating model:{config.MODEL.TYPE}/{config.MODEL.NAME}')
    model = build_model(config)
    model.cuda()
    logger.info(str(model))

    # build optimizer
    optimizer = build_optimizer(config, model)

    if config.AMP_OPT_LEVEL != 'O0':
        config.defrost()
        if has_native_amp:
            config.native_amp = True
            use_amp = 'native'
        elif has_apex:
            config.apex_amp = True
            use_amp = 'apex'
        else:
            use_amp = None
            logger.warning(
                'Neither APEX or native Torch AMP is available, using float32. '
                'Install NVIDA apex or upgrade to PyTorch 1.6')
        config.freeze()

    # setup automatic mixed-precision (AMP) loss scaling and op casting
    amp_autocast = suppress  # do nothing
    loss_scaler = None
    if config.AMP_OPT_LEVEL != 'O0':
        if use_amp == 'apex':
            model, optimizer = amp.initialize(model,
                                              optimizer,
                                              opt_level=config.AMP_OPT_LEVEL)
            loss_scaler = ApexScaler()
            if config.LOCAL_RANK == 0:
                logger.info(
                    'Using NVIDIA APEX AMP. Training in mixed precision.')
        if use_amp == 'native':
            amp_autocast = torch.cuda.amp.autocast
            loss_scaler = NativeScaler()
            if config.LOCAL_RANK == 0:
                logger.info(
                    'Using native Torch AMP. Training in mixed precision.')
        else:
            if config.LOCAL_RANK == 0:
                logger.info('AMP not enabled. Training in float32.')

    # put model on gpus
    model = torch.nn.parallel.DistributedDataParallel(
        model, device_ids=[config.LOCAL_RANK], broadcast_buffers=False)

    try:
        model.register_comm_hook(state=None, hook=fp16_compress_hook)
        logger.info('using fp16_compress_hook!')
    except:
        logger.info('cannot register fp16_compress_hook!')

    model_without_ddp = model.module

    n_parameters = sum(p.numel() for p in model.parameters()
                       if p.requires_grad)
    logger.info(f'number of params: {n_parameters}')
    if hasattr(model_without_ddp, 'flops'):
        flops = model_without_ddp.flops()
        logger.info(f'number of GFLOPs: {flops / 1e9}')

    # build learning rate scheduler
    lr_scheduler = build_scheduler(config, optimizer, len(data_loader_train)) \
        if not config.EVAL_MODE else None

    # build criterion
    if config.AUG.MIXUP > 0.:
        # smoothing is handled with mixup label transform
        criterion = SoftTargetCrossEntropy()
    elif config.MODEL.LABEL_SMOOTHING > 0.:
        criterion = LabelSmoothingCrossEntropy(
            smoothing=config.MODEL.LABEL_SMOOTHING)
    else:
        criterion = torch.nn.CrossEntropyLoss()

    max_accuracy = 0.0
    max_ema_accuracy = 0.0
    # set auto resume
    if config.MODEL.RESUME == '' and config.TRAIN.AUTO_RESUME:
        resume_file = auto_resume_helper(config.OUTPUT)
        if resume_file:
            if config.MODEL.RESUME:
                logger.warning(
                    f'auto-resume changing resume file from {config.MODEL.RESUME} to {resume_file}'
                )
            config.defrost()
            config.MODEL.RESUME = resume_file
            config.freeze()
            logger.info(f'auto resuming from {resume_file}')
        else:
            logger.info(
                f'no checkpoint found in {config.OUTPUT}, ignoring auto resume'
            )

    # set resume and pretrain
    if config.MODEL.RESUME:
        max_accuracy = load_checkpoint(config, model_without_ddp, optimizer,
                                       lr_scheduler, loss_scaler, logger)
        if data_loader_val is not None:
            acc1, acc5, loss = validate(config, data_loader_val, model)
            logger.info(
                f'Accuracy of the network on the {len(dataset_val)} test images: {acc1:.1f}%'
            )
    elif config.MODEL.PRETRAINED:
        load_pretrained(config, model_without_ddp, logger)
        if data_loader_val is not None:
            acc1, acc5, loss = validate(config, data_loader_val, model)
            logger.info(
                f'Accuracy of the network on the {len(dataset_val)} test images: {acc1:.1f}%'
            )

    # evaluate EMA
    model_ema = None
    if config.TRAIN.EMA.ENABLE:
        # Important to create EMA model after cuda(), DP wrapper, and AMP but before SyncBN and DDP wrapper
        model_ema = ModelEma(model, decay=config.TRAIN.EMA.DECAY)
        print('Using EMA with decay = %.8f' % config.TRAIN.EMA.DECAY)
        if config.MODEL.RESUME:
            load_ema_checkpoint(config, model_ema, logger)
            acc1, acc5, loss = validate(config, data_loader_val, model_ema.ema)
            logger.info(
                f'Accuracy of the ema network on the {len(dataset_val)} test images: {acc1:.1f}%'
            )

    if config.THROUGHPUT_MODE:
        throughput(data_loader_val, model, logger)

    if config.EVAL_MODE:
        return

    # train
    logger.info('Start training')
    start_time = time.time()
    for epoch in range(config.TRAIN.START_EPOCH, config.TRAIN.EPOCHS):
        data_loader_train.sampler.set_epoch(epoch)

        train_one_epoch(config,
                        model,
                        criterion,
                        data_loader_train,
                        optimizer,
                        epoch,
                        mixup_fn,
                        lr_scheduler,
                        amp_autocast,
                        loss_scaler,
                        model_ema=model_ema)
        if (epoch % config.SAVE_FREQ == 0 or epoch == (config.TRAIN.EPOCHS - 1)) and \
                config.TRAIN.OPTIMIZER.USE_ZERO:
            optimizer.consolidate_state_dict(to=0)
        if dist.get_rank() == 0 and (epoch % config.SAVE_FREQ == 0
                                     or epoch == (config.TRAIN.EPOCHS - 1)):
            save_checkpoint(config,
                            epoch,
                            model_without_ddp,
                            max_accuracy,
                            optimizer,
                            lr_scheduler,
                            loss_scaler,
                            logger,
                            model_ema=model_ema)
        if data_loader_val is not None and epoch % config.EVAL_FREQ == 0:
            acc1, acc5, loss = validate(config, data_loader_val, model, epoch)
            logger.info(
                f'Accuracy of the network on the {len(dataset_val)} test images: {acc1:.1f}%'
            )
            if dist.get_rank() == 0 and acc1 > max_accuracy:
                save_checkpoint(config,
                                epoch,
                                model_without_ddp,
                                max_accuracy,
                                optimizer,
                                lr_scheduler,
                                loss_scaler,
                                logger,
                                model_ema=model_ema,
                                best='best')
            max_accuracy = max(max_accuracy, acc1)
            logger.info(f'Max accuracy: {max_accuracy:.2f}%')

            if config.TRAIN.EMA.ENABLE:
                acc1, acc5, loss = validate(config, data_loader_val,
                                            model_ema.ema, epoch)
                logger.info(
                    f'Accuracy of the ema network on the {len(dataset_val)} test images: {acc1:.1f}%'
                )
                if dist.get_rank() == 0 and acc1 > max_ema_accuracy:
                    save_checkpoint(config,
                                    epoch,
                                    model_without_ddp,
                                    max_accuracy,
                                    optimizer,
                                    lr_scheduler,
                                    loss_scaler,
                                    logger,
                                    model_ema=model_ema,
                                    best='ema_best')
                max_ema_accuracy = max(max_ema_accuracy, acc1)
                logger.info(f'Max ema accuracy: {max_ema_accuracy:.2f}%')

    total_time = time.time() - start_time
    total_time_str = str(datetime.timedelta(seconds=int(total_time)))
    logger.info('Training time {}'.format(total_time_str))


def train_one_epoch(config,
                    model,
                    criterion,
                    data_loader,
                    optimizer,
                    epoch,
                    mixup_fn,
                    lr_scheduler,
                    amp_autocast=suppress,
                    loss_scaler=None,
                    model_ema=None):
    model.train()
    optimizer.zero_grad()

    num_steps = len(data_loader)
    batch_time = AverageMeter()
    model_time = AverageMeter()
    loss_meter = AverageMeter()
    norm_meter = MyAverageMeter(300)

    start = time.time()
    end = time.time()

    amp_type = torch.float16 if config.AMP_TYPE == 'float16' else torch.bfloat16
    for idx, (samples, targets) in enumerate(data_loader):
        iter_begin_time = time.time()
        if type(samples) == list:
            samples = [item.cuda(non_blocking=True) for item in samples]
        else:
            samples = samples.cuda(non_blocking=True)
        targets = targets.cuda(non_blocking=True)

        if mixup_fn is not None:
            samples, targets = mixup_fn(samples, targets)

        if not obsolete_torch_version(TORCH_VERSION,
                                      (1, 9)) and config.AMP_OPT_LEVEL != 'O0':
            with amp_autocast(dtype=amp_type):
                outputs = model(samples)
        else:
            with amp_autocast():
                outputs = model(samples)

        if config.TRAIN.ACCUMULATION_STEPS > 1:
            if not obsolete_torch_version(
                    TORCH_VERSION, (1, 9)) and config.AMP_OPT_LEVEL != 'O0':
                with amp_autocast(dtype=amp_type):
                    loss = criterion(outputs, targets)
                    loss = loss / config.TRAIN.ACCUMULATION_STEPS
            else:
                with amp_autocast():
                    loss = criterion(outputs, targets)
                    loss = loss / config.TRAIN.ACCUMULATION_STEPS
            if config.AMP_OPT_LEVEL != 'O0':
                is_second_order = hasattr(optimizer, 'is_second_order') and \
                    optimizer.is_second_order
                grad_norm = loss_scaler(loss,
                                        optimizer,
                                        clip_grad=config.TRAIN.CLIP_GRAD,
                                        parameters=model.parameters(),
                                        create_graph=is_second_order,
                                        update_grad=(idx + 1) %
                                        config.TRAIN.ACCUMULATION_STEPS == 0)
                if (idx + 1) % config.TRAIN.ACCUMULATION_STEPS == 0:
                    optimizer.zero_grad()
                    if model_ema is not None:
                        model_ema.update(model)
            else:
                loss.backward()
                if config.TRAIN.CLIP_GRAD:
                    grad_norm = torch.nn.utils.clip_grad_norm_(
                        model.parameters(), config.TRAIN.CLIP_GRAD)
                else:
                    grad_norm = get_grad_norm(model.parameters())
                if (idx + 1) % config.TRAIN.ACCUMULATION_STEPS == 0:
                    optimizer.step()
                    optimizer.zero_grad()
                    if model_ema is not None:
                        model_ema.update(model)
            if (idx + 1) % config.TRAIN.ACCUMULATION_STEPS == 0:
                lr_scheduler.step_update(epoch * num_steps + idx)
        else:
            if not obsolete_torch_version(
                    TORCH_VERSION, (1, 9)) and config.AMP_OPT_LEVEL != 'O0':
                with amp_autocast(dtype=amp_type):
                    loss = criterion(outputs, targets)
            else:
                with amp_autocast():
                    loss = criterion(outputs, targets)
            optimizer.zero_grad()
            if config.AMP_OPT_LEVEL != 'O0':
                is_second_order = hasattr(optimizer, 'is_second_order') and \
                    optimizer.is_second_order
                grad_norm = loss_scaler(loss,
                                        optimizer,
                                        clip_grad=config.TRAIN.CLIP_GRAD,
                                        parameters=model.parameters(),
                                        create_graph=is_second_order,
                                        update_grad=(idx + 1) %
                                        config.TRAIN.ACCUMULATION_STEPS == 0)
                if model_ema is not None:
                    model_ema.update(model)
            else:
                loss.backward()
                if config.TRAIN.CLIP_GRAD:
                    grad_norm = torch.nn.utils.clip_grad_norm_(
                        model.parameters(), config.TRAIN.CLIP_GRAD)
                else:
                    grad_norm = get_grad_norm(model.parameters())
                optimizer.step()
                if model_ema is not None:
                    model_ema.update(model)

            lr_scheduler.step_update(epoch * num_steps + idx)

        torch.cuda.synchronize()

        loss_meter.update(loss.item(), targets.size(0))
        if grad_norm is not None:
            norm_meter.update(grad_norm.item())
        batch_time.update(time.time() - end)
        model_time.update(time.time() - iter_begin_time)
        end = time.time()

        if idx % config.PRINT_FREQ == 0:
            lr = optimizer.param_groups[0]['lr']
            memory_used = torch.cuda.max_memory_allocated() / (1024.0 * 1024.0)
            etas = batch_time.avg * (num_steps - idx)
            logger.info(
                f'Train: [{epoch}/{config.TRAIN.EPOCHS}][{idx}/{num_steps}]\t'
                f'eta {datetime.timedelta(seconds=int(etas))} lr {lr:.6f}\t'
                f'time {batch_time.val:.4f} ({batch_time.avg:.4f})\t'
                f'model_time {model_time.val:.4f} ({model_time.avg:.4f})\t'
                f'loss {loss_meter.val:.4f} ({loss_meter.avg:.4f})\t'
                f'grad_norm {norm_meter.val:.4f} ({norm_meter.avg:.4f}/{norm_meter.var:.4f})\t'
                f'mem {memory_used:.0f}MB')
    epoch_time = time.time() - start
    logger.info(
        f'EPOCH {epoch} training takes {datetime.timedelta(seconds=int(epoch_time))}'
    )


@torch.no_grad()
def validate(config, data_loader, model, epoch=None):
    criterion = torch.nn.CrossEntropyLoss()
    model.eval()

    batch_time = AverageMeter()
    loss_meter = AverageMeter()
    acc1_meter = AverageMeter()
    acc5_meter = AverageMeter()

    end = time.time()
    for idx, (images, target) in enumerate(data_loader):
        if type(images) == list:
            images = [item.cuda(non_blocking=True) for item in images]
        else:
            images = images.cuda(non_blocking=True)
        target = target.cuda(non_blocking=True)
        output = model(images)

        # convert 22k to 1k to evaluate
        if output.size(-1) == 21841:
            convert_file = './meta_data/map22kto1k.txt'
            with open(convert_file, 'r') as f:
                convert_list = [int(line) for line in f.readlines()]
            output = output[:, convert_list]

        # measure accuracy and record loss
        loss = criterion(output, target)
        acc1, acc5 = accuracy(output, target, topk=(1, 5))

        acc1 = reduce_tensor(acc1)
        acc5 = reduce_tensor(acc5)
        loss = reduce_tensor(loss)

        loss_meter.update(loss.item(), target.size(0))
        acc1_meter.update(acc1.item(), target.size(0))
        acc5_meter.update(acc5.item(), target.size(0))

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

        if idx % config.PRINT_FREQ == 0:
            memory_used = torch.cuda.max_memory_allocated() / (1024.0 * 1024.0)
            logger.info(f'Test: [{idx}/{len(data_loader)}]\t'
                        f'Time {batch_time.val:.3f} ({batch_time.avg:.3f})\t'
                        f'Loss {loss_meter.val:.4f} ({loss_meter.avg:.4f})\t'
                        f'Acc@1 {acc1_meter.val:.3f} ({acc1_meter.avg:.3f})\t'
                        f'Acc@5 {acc5_meter.val:.3f} ({acc5_meter.avg:.3f})\t'
                        f'Mem {memory_used:.0f}MB')
    if epoch is not None:
        logger.info(
            f'[Epoch:{epoch}] * Acc@1 {acc1_meter.avg:.3f} Acc@5 {acc5_meter.avg:.3f}'
        )
    else:
        logger.info(
            f' * Acc@1 {acc1_meter.avg:.3f} Acc@5 {acc5_meter.avg:.3f}')

    return acc1_meter.avg, acc5_meter.avg, loss_meter.avg


if __name__ == '__main__':
    _, config = parse_option()

    if config.AMP_OPT_LEVEL != 'O0':
        assert has_native_amp, 'Please update pytorch(1.6+) to support amp!'

    # init distributed env
    # In the newer versions of Slurm, the format of `SLURM_TASKS_PER_NODE` has changed from a single
    # numeric string to a format like `8(xn)`, which represents n nodes is used in the training.
    if 'SLURM_PROCID' in os.environ and int(os.environ['SLURM_TASKS_PER_NODE'][0]) != 1:
        print('\nDist init: SLURM')
        rank = int(os.environ['SLURM_PROCID'])
        gpu = rank % torch.cuda.device_count()
        config.defrost()
        config.LOCAL_RANK = gpu
        config.freeze()

        world_size = int(os.environ['SLURM_NTASKS'])
        if 'MASTER_PORT' not in os.environ:
            os.environ['MASTER_PORT'] = '29501'
        node_list = os.environ['SLURM_NODELIST']
        addr = subprocess.getoutput(
            f'scontrol show hostname {node_list} | head -n1')
        if 'MASTER_ADDR' not in os.environ:
            os.environ['MASTER_ADDR'] = addr

        os.environ['RANK'] = str(rank)
        os.environ['LOCAL_RANK'] = str(gpu)
        os.environ['LOCAL_SIZE'] = str(torch.cuda.device_count())
        os.environ['WORLD_SIZE'] = str(world_size)
    if 'RANK' in os.environ and 'WORLD_SIZE' in os.environ:
        rank = int(os.environ['RANK'])
        world_size = int(os.environ['WORLD_SIZE'])
        print(f'RANK and WORLD_SIZE in environ: {rank}/{world_size}')
    else:
        rank = -1
        world_size = -1
    torch.cuda.set_device(config.LOCAL_RANK)
    torch.distributed.init_process_group(backend='nccl',
                                         init_method='env://',
                                         world_size=world_size,
                                         rank=rank)
    torch.distributed.barrier()

    seed = config.SEED + dist.get_rank()
    torch.manual_seed(seed)
    torch.cuda.manual_seed(seed)
    np.random.seed(seed)
    random.seed(seed)
    cudnn.benchmark = True

    # linear scale the learning rate according to total batch size, may not be optimal
    linear_scaled_lr = config.TRAIN.BASE_LR * \
        config.DATA.BATCH_SIZE * dist.get_world_size() / 512.0
    linear_scaled_warmup_lr = config.TRAIN.WARMUP_LR * \
        config.DATA.BATCH_SIZE * dist.get_world_size() / 512.0
    linear_scaled_min_lr = config.TRAIN.MIN_LR * \
        config.DATA.BATCH_SIZE * dist.get_world_size() / 512.0
    # gradient accumulation also need to scale the learning rate
    if config.TRAIN.ACCUMULATION_STEPS > 1:
        linear_scaled_lr = linear_scaled_lr * config.TRAIN.ACCUMULATION_STEPS
        linear_scaled_warmup_lr = linear_scaled_warmup_lr * config.TRAIN.ACCUMULATION_STEPS
        linear_scaled_min_lr = linear_scaled_min_lr * config.TRAIN.ACCUMULATION_STEPS
    config.defrost()
    config.TRAIN.BASE_LR = linear_scaled_lr
    config.TRAIN.WARMUP_LR = linear_scaled_warmup_lr
    config.TRAIN.MIN_LR = linear_scaled_min_lr
    print(config.AMP_OPT_LEVEL, _.amp_opt_level)

    config.freeze()

    os.makedirs(config.OUTPUT, exist_ok=True)
    logger = create_logger(output_dir=config.OUTPUT,
                           dist_rank=dist.get_rank(),
                           name=f'{config.MODEL.NAME}')

    if dist.get_rank() == 0:
        path = os.path.join(config.OUTPUT, 'config.json')
        with open(path, 'w') as f:
            f.write(config.dump())
        logger.info(f'Full config saved to {path}')

    # print config
    logger.info(config.dump())

    main(config)