Pikaliov
1b3206b6a7
- Add enhanced README with project structure and quick start guide - Initialize repository with DCNv4 CUDA extension (PyTorch module) - Include classification, detection, and segmentation subdirectories - Reference upstream OpenGVLab DCNv4 implementation Co-Authored-By: Claude Haiku 4.5 <noreply@anthropic.com>
672 lines
26 KiB
Python
672 lines
26 KiB
Python
# --------------------------------------------------------
|
|
# DCNv4
|
|
# Copyright (c) 2024 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
|
|
from timm.utils import ModelEma, ApexScaler
|
|
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 build_optimizer
|
|
from logger import create_logger
|
|
from utils import NativeScalerWithGradNormCount as NativeScaler
|
|
from utils import (load_checkpoint, load_pretrained, save_checkpoint,
|
|
get_grad_norm, auto_resume_helper, reduce_tensor,
|
|
load_ema_checkpoint, MyAverageMeter)
|
|
|
|
from contextlib import suppress
|
|
from ddp_hooks import fp16_compress_hook
|
|
|
|
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,
|
|
default=0,
|
|
help='local rank for DistributedDataParallel')
|
|
|
|
args, unparsed = parser.parse_known_args()
|
|
|
|
if 'LOCAL_RANK' not in os.environ:
|
|
os.environ['LOCAL_RANK'] = str(args.local_rank)
|
|
|
|
|
|
config = get_config(args)
|
|
|
|
config.defrost()
|
|
config.LOCAL_RANK = int(os.environ['LOCAL_RANK'])
|
|
config.freeze()
|
|
return args, 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 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, amp_autocast=amp_autocast)
|
|
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, amp_autocast=amp_autocast)
|
|
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, amp_autocast=amp_autocast)
|
|
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, amp_autocast)
|
|
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, amp_autocast)
|
|
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()
|
|
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, amp_autocast=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)
|
|
with amp_autocast():
|
|
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
|
|
if 'SLURM_PROCID' in os.environ and int(os.environ['SLURM_NNODES']) != 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)
|