main_deepspeed.py

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

import os
import time
import random
import argparse
import datetime
import numpy as np
import subprocess

import torch
import torch.backends.cudnn as cudnn
import torch.distributed as dist
import deepspeed
from timm.loss import LabelSmoothingCrossEntropy, SoftTargetCrossEntropy
from timm.utils import accuracy, AverageMeter

from config import get_config
from models import build_model
from dataset import build_loader
from lr_scheduler import build_scheduler
from optimizer import set_weight_decay_and_lr
from logger import create_logger
from utils import load_pretrained, reduce_tensor, MyAverageMeter
from ddp_hooks import fp16_compress_hook
from ema_deepspeed import EMADeepspeed

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('--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('--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('--accumulation-steps', type=int, default=1, help="gradient accumulation steps")

    # distributed training
    parser.add_argument("--local-rank", type=int, required=True, help='local rank for DistributedDataParallel')
    parser.add_argument('--disable-grad-scalar', action='store_true', help='disable Grad Scalar')

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

    return args, config


def seed_everything(seed, rank):
    seed = seed + rank
    torch.manual_seed(seed)
    torch.cuda.manual_seed(seed)
    np.random.seed(seed)
    random.seed(seed)
    cudnn.benchmark = True


def save_config(config):
    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}")


def build_criterion(config):
    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()
    return criterion


def scale_learning_rate(config, num_processes):
    # 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 * num_processes / 512.0
    linear_scaled_warmup_lr = config.TRAIN.WARMUP_LR * \
        config.DATA.BATCH_SIZE * num_processes / 512.0
    linear_scaled_min_lr = config.TRAIN.MIN_LR * \
        config.DATA.BATCH_SIZE * num_processes / 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
    config.freeze()

    logger.info('BASE_LR={}'.format(config.TRAIN.BASE_LR))
    logger.info('WARMUP_LR={}'.format(config.TRAIN.WARMUP_LR))
    logger.info('MIN_LR={}'.format(config.TRAIN.MIN_LR))


def log_model_statistic(model_wo_ddp):
    n_parameters = sum(p.numel() for p in model_wo_ddp.parameters()
                       if p.requires_grad)
    logger.info(f"number of params: {n_parameters/1e6} M")
    if hasattr(model_wo_ddp, 'flops'):
        flops = model_wo_ddp.flops()
        logger.info(f"number of GFLOPs: {flops / 1e9}")


def get_parameter_groups(model, config):
    skip = {}
    skip_keywords = {}
    if hasattr(model, 'no_weight_decay'):
        skip = model.no_weight_decay()
    if hasattr(model, 'no_weight_decay_keywords'):
        skip_keywords = model.no_weight_decay_keywords()

    parameters = set_weight_decay_and_lr(
        model,
        config.TRAIN.WEIGHT_DECAY,
        config.TRAIN.BASE_LR,
        skip,
        skip_keywords,
        lr_layer_decay=config.TRAIN.LR_LAYER_DECAY,
        lr_layer_decay_ratio=config.TRAIN.LR_LAYER_DECAY_RATIO,
        freeze_backbone=config.TRAIN.OPTIMIZER.FREEZE_BACKBONE,
        dcn_lr_mul=config.TRAIN.OPTIMIZER.DCN_LR_MUL,
    )
    return parameters


def get_optimizer_state_str(optimizer):
    states = []
    for param_group in optimizer.param_groups:
        states.append(f'name={param_group["name"]} lr={param_group["lr"]} weight_decay={param_group["weight_decay"]}')
    return '\n'.join(states)


def build_ds_config(config, args):
    opt_lower = config.TRAIN.OPTIMIZER.NAME.lower()
    if opt_lower == 'adamw':
        optimizer = {
            "type": "AdamW",
            "params": {
                "lr": config.TRAIN.BASE_LR,
                "eps": config.TRAIN.OPTIMIZER.EPS,
                "betas": config.TRAIN.OPTIMIZER.BETAS,
                "weight_decay": config.TRAIN.WEIGHT_DECAY
            }
        }
    else:
        return NotImplemented

    ds_config = {
        "train_micro_batch_size_per_gpu": config.DATA.BATCH_SIZE,
        "optimizer": optimizer,
        "fp16": {
            "enabled": True,
            "auto_cast": True,
            "loss_scale": 1 if args.disable_grad_scalar else 0
        },
        "zero_optimization": {
            "stage": 1,
        },
        "steps_per_print": 1e10,
        "gradient_accumulation_steps": config.TRAIN.ACCUMULATION_STEPS,
        "gradient_clipping": config.TRAIN.CLIP_GRAD,
    }
    return ds_config


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

    for idx, (images, _) in enumerate(data_loader):
        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 train_epoch(config, model, criterion, data_loader, optimizer, epoch, mixup_fn, lr_scheduler, model_ema=None):
    model.train()

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

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

    for idx, (samples, targets) in enumerate(data_loader):
        iter_begin_time = time.time()
        samples = samples.cuda(non_blocking=True)
        targets = targets.cuda(non_blocking=True)

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

        outputs = model(samples)
        loss = criterion(outputs, targets)

        model.backward(loss)
        model.step()

        if model_ema is not None:
            model_ema(model)

        if (idx + 1) % config.TRAIN.ACCUMULATION_STEPS == 0:
            lr_scheduler.step_update(epoch * num_steps + idx)

        torch.cuda.synchronize()
        loss_meter.update(loss.item(), targets.size(0))
        norm_meter.update(optimizer._global_grad_norm)
        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 eval_epoch(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):
        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


def train(config, ds_config):
    # -------------- build ---------------- #
    
    _, dataset_val, _, data_loader_train, data_loader_val, _, mixup_fn = build_loader(config)
    model = build_model(config)
    model.cuda()

    if config.MODEL.PRETRAINED:
        load_pretrained(config, model, logger)

    logger.info(ds_config)
    model, optimizer, _, _ = deepspeed.initialize(
        config=ds_config,
        model=model,
        model_parameters=get_parameter_groups(model, config),
        dist_init_required=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

    lr_scheduler = build_scheduler(config, optimizer, len(data_loader_train))
    criterion = build_criterion(config)

    model_ema = None
    if config.TRAIN.EMA.ENABLE:
        model_ema = EMADeepspeed(model, config.TRAIN.EMA.DECAY)

    # -------------- resume ---------------- #
    
    max_accuracy = 0.0
    max_accuracy_ema = 0.0
    client_state = {}
    if config.MODEL.RESUME == '' and config.TRAIN.AUTO_RESUME:
        if os.path.exists(os.path.join(config.OUTPUT, 'latest')):
            config.defrost()
            config.MODEL.RESUME = config.OUTPUT
            config.freeze()
            tag = None
    elif config.MODEL.RESUME:
        config.MODEL.RESUME = os.path.dirname(config.MODEL.RESUME)
        tag = os.path.basename(config.MODEL.RESUME)
    if config.MODEL.RESUME:
        logger.info('loading checkpoint from {}'.format(config.MODEL.RESUME))
        _, client_state = model.load_checkpoint(load_dir=config.MODEL.RESUME, tag=tag)
        logger.info(f'client_state={client_state.keys()}')
        lr_scheduler.load_state_dict(client_state['custom_lr_scheduler'])
        max_accuracy = client_state['max_accuracy']

        if model_ema is not None:
            max_accuracy_ema = client_state.get('max_accuracy_ema', 0.0)
            model_ema.load_state_dict((client_state['model_ema']))
        
    # -------------- training ---------------- #
    
    logger.info(f"Creating model:{config.MODEL.TYPE}/{config.MODEL.NAME}")
    logger.info(str(model))
    logger.info(get_optimizer_state_str(optimizer))
    logger.info("Start training")
    logger.info('max_accuracy: {}'.format(max_accuracy))
    log_model_statistic(model_without_ddp)

    start_time = time.time()
    start_epoch = client_state['epoch'] + 1 if 'epoch' in client_state else config.TRAIN.START_EPOCH
    for epoch in range(start_epoch, config.TRAIN.EPOCHS):
        data_loader_train.sampler.set_epoch(epoch)
        train_epoch(config, model, criterion, data_loader_train, optimizer, epoch, mixup_fn, lr_scheduler,
                    model_ema=model_ema)

        if epoch % config.SAVE_FREQ == 0 or epoch == config.TRAIN.EPOCHS - 1:
            model.save_checkpoint(
                save_dir=config.OUTPUT,
                tag=f'epoch{epoch}',
                client_state={
                    'custom_lr_scheduler': lr_scheduler.state_dict(),
                    'max_accuracy': max_accuracy,
                    'epoch': epoch,
                    'config': config,
                    'max_accuracy_ema': max_accuracy_ema if model_ema is not None else 0.0,
                    'model_ema': model_ema.state_dict() if model_ema is not None else None,
                }
            )

        if epoch % config.EVAL_FREQ == 0:
            acc1, _, _ = eval_epoch(config, data_loader_val, model, epoch)
            logger.info(f"Accuracy of the network on the {len(dataset_val)} test images: {acc1:.1f}%")

            if acc1 > max_accuracy:
                model.save_checkpoint(
                    save_dir=config.OUTPUT,
                    tag='best',
                    client_state={
                        'custom_lr_scheduler': lr_scheduler.state_dict(),
                        'max_accuracy': max_accuracy,
                        'epoch': epoch,
                        'config': config,
                        'max_accuracy_ema': max_accuracy_ema if model_ema is not None else 0.0,
                        'model_ema': model_ema.state_dict() if model_ema is not None else None,
                    }
                )

            max_accuracy = max(max_accuracy, acc1)
            logger.info(f'Max accuracy: {max_accuracy:.2f}%')

            if model_ema is not None:
                with model_ema.activate(model):
                    acc1_ema, _, _ = eval_epoch(config, data_loader_val, model, epoch)
                    logger.info(f"[EMA] Accuracy of the network on the {len(dataset_val)} test images: {acc1_ema:.1f}%")
                    max_accuracy_ema = max(max_accuracy_ema, acc1_ema)
                    logger.info(f'[EMA] Max accuracy: {max_accuracy_ema:.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 eval(config):
    _, _, _, _, data_loader_val, _, _ = build_loader(config)
    model = build_model(config)
    model.cuda()
    model = torch.nn.parallel.DistributedDataParallel(
        model, device_ids=[config.LOCAL_RANK], broadcast_buffers=False)

    model_wo_ddp = model.module
    if config.MODEL.RESUME:
        try:
            checkpoint = torch.load(config.MODEL.RESUME, map_location='cpu')
            msg = model_wo_ddp.load_state_dict(checkpoint['model'], strict=False)
            logger.info(msg)
        except:
            try:
                from deepspeed.utils.zero_to_fp32 import get_fp32_state_dict_from_zero_checkpoint
                ckpt_dir = os.path.dirname(config.MODEL.RESUME)
                tag = os.path.basename(config.MODEL.RESUME)
                state_dict = get_fp32_state_dict_from_zero_checkpoint(checkpoint_dir=ckpt_dir, tag=tag) 
                model_wo_ddp.load_state_dict(state_dict)
            except:
                checkpoint = torch.load(os.path.join(config.MODEL.RESUME, 'mp_rank_00_model_states.pt'), map_location='cpu')
                model_wo_ddp.load_state_dict(checkpoint['module'])
    elif config.MODEL.PRETRAINED:
        load_pretrained(config, model_wo_ddp, logger)

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

    eval_epoch(config, data_loader_val, model)


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

    # init distributed env
    if 'SLURM_PROCID' in os.environ and int(os.environ['SLURM_TASKS_PER_NODE']) != 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()

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

    if dist.get_rank() == 0: save_config(config)
    scale_learning_rate(config, dist.get_world_size())
    seed_everything(config.SEED, dist.get_rank())

    if config.EVAL_MODE:
        eval(config)
    else:
        train(config, build_ds_config(config, args))