474 lines
19 KiB
Python
474 lines
19 KiB
Python
#!/usr/bin/env python
|
|
# Copyright (c) Facebook, Inc. and its affiliates.
|
|
# All rights reserved.
|
|
|
|
# This source code is licensed under the license found in the
|
|
# LICENSE file in the root directory of this source tree.
|
|
|
|
import argparse
|
|
import builtins
|
|
import math
|
|
import os
|
|
import random
|
|
import shutil
|
|
import time
|
|
import warnings
|
|
|
|
# ===== SLURM, to delete =====
|
|
import signal
|
|
# ===== SLURM, to delete =====
|
|
|
|
import torch
|
|
import torch.nn as nn
|
|
import torch.nn.parallel
|
|
import torch.backends.cudnn as cudnn
|
|
import torch.distributed as dist
|
|
import torch.optim
|
|
import torch.multiprocessing as mp
|
|
import torch.utils.data
|
|
import torch.utils.data.distributed
|
|
import torchvision.transforms as transforms
|
|
import torchvision.datasets as datasets
|
|
import torchvision.models as torchvision_models
|
|
|
|
from functools import partial
|
|
import vits
|
|
|
|
import moco.builder
|
|
import moco.loader
|
|
import moco.optimizer
|
|
|
|
# ===== SLURM, to delete =====
|
|
def signalHandler(a, b):
|
|
if a == signal.SIGUSR1:
|
|
logger.info('Got SIGUSR1.')
|
|
elif a == signal.SIGTERM:
|
|
logger.info('Got SIGTERM.')
|
|
|
|
signal.signal(signal.SIGUSR1, signalHandler)
|
|
signal.signal(signal.SIGTERM, signalHandler)
|
|
# ===== SLURM, to delete =====
|
|
|
|
torchvision_model_names = sorted(name for name in torchvision_models.__dict__
|
|
if name.islower() and not name.startswith("__")
|
|
and callable(torchvision_models.__dict__[name]))
|
|
|
|
model_names = ['vit_small', 'vit_base', 'vit_large', 'vit_huge'] + torchvision_model_names
|
|
|
|
parser = argparse.ArgumentParser(description='PyTorch ImageNet Training')
|
|
parser.add_argument('data', metavar='DIR',
|
|
help='path to dataset')
|
|
parser.add_argument('-a', '--arch', metavar='ARCH', default='resnet50',
|
|
choices=model_names,
|
|
help='model architecture: ' +
|
|
' | '.join(model_names) +
|
|
' (default: resnet50)')
|
|
parser.add_argument('-j', '--workers', default=32, type=int, metavar='N',
|
|
help='number of data loading workers (default: 32)')
|
|
parser.add_argument('--epochs', default=100, 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=4096, type=int,
|
|
metavar='N',
|
|
help='mini-batch size (default: 4096), 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.45, type=float,
|
|
metavar='LR', help='initial (base) 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-6, type=float,
|
|
metavar='W', help='weight decay (default: 1e-6)',
|
|
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('--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')
|
|
|
|
# moco specific configs:
|
|
parser.add_argument('--moco-dim', default=256, type=int,
|
|
help='feature dimension (default: 256)')
|
|
parser.add_argument('--moco-mlp-dim', default=4096, type=int,
|
|
help='hidden dimension in MLPs (default: 4096)')
|
|
parser.add_argument('--moco-m', default=0.99, type=float,
|
|
help='moco (initial) momentum of updating momentum encoder '
|
|
'the value will gradually increase to 1 with a '
|
|
'half-cycle cosine schedule (default: 0.99)')
|
|
parser.add_argument('--moco-t', default=1.0, type=float,
|
|
help='softmax temperature (default: 1.0)')
|
|
|
|
# vit specific configs:
|
|
parser.add_argument('--stop-grad-conv1', action='store_true',
|
|
help='stop-grad after first conv, or patch embedding')
|
|
|
|
# other upgrades
|
|
parser.add_argument('--optimizer', default='lars', type=str,
|
|
choices=['lars', 'adamw'],
|
|
help='optimizer used (default: lars)')
|
|
parser.add_argument('--warmup-epochs', default=10, type=int, metavar='N',
|
|
help='number of warmup epochs')
|
|
|
|
# ===== OTHERS, to delete =====
|
|
parser.add_argument('--checkpoint-folder', default='.', type=str, metavar='PATH',
|
|
help='path to save the checkpoints (default: .)')
|
|
# ===== OTHERS, to delete =====
|
|
|
|
def main():
|
|
args = parser.parse_args()
|
|
|
|
if args.seed is not None:
|
|
random.seed(args.seed)
|
|
torch.manual_seed(args.seed)
|
|
cudnn.deterministic = True
|
|
warnings.warn('You have chosen to seed training. '
|
|
'This will turn on the CUDNN deterministic setting, '
|
|
'which can slow down your training considerably! '
|
|
'You may see unexpected behavior when restarting '
|
|
'from checkpoints.')
|
|
|
|
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
|
|
|
|
# ===== PATH, to delete =====
|
|
if not os.path.exists(args.checkpoint_folder):
|
|
os.makedirs(args.checkpoint_folder)
|
|
# ===== PATH, to delete =====
|
|
|
|
ngpus_per_node = torch.cuda.device_count()
|
|
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 = ngpus_per_node * args.world_size
|
|
# Use torch.multiprocessing.spawn to launch distributed processes: the
|
|
# main_worker process function
|
|
mp.spawn(main_worker, nprocs=ngpus_per_node, args=(ngpus_per_node, args))
|
|
else:
|
|
# Simply call main_worker function
|
|
main_worker(args.gpu, ngpus_per_node, args)
|
|
|
|
|
|
def main_worker(gpu, ngpus_per_node, args):
|
|
args.gpu = gpu
|
|
|
|
# suppress printing if not first GPU on each node
|
|
if args.multiprocessing_distributed and (args.gpu != 0 or args.rank != 0):
|
|
def print_pass(*args):
|
|
pass
|
|
builtins.print = print_pass
|
|
|
|
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 * 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)
|
|
torch.distributed.barrier()
|
|
# create model
|
|
print("=> creating model '{}'".format(args.arch))
|
|
if args.arch.startswith('vit'):
|
|
model = moco.builder.MoCo(
|
|
partial(vits.__dict__[args.arch], stop_grad_conv1=args.stop_grad_conv1),
|
|
True, # with vit setup
|
|
args.moco_dim, args.moco_mlp_dim, args.moco_t)
|
|
else:
|
|
model = moco.builder.MoCo(
|
|
partial(torchvision_models.__dict__[args.arch], zero_init_residual=True),
|
|
False, # with resnet setup
|
|
args.moco_dim, args.moco_mlp_dim, args.moco_t)
|
|
|
|
# infer learning rate before changing batch size
|
|
init_lr = args.lr * args.batch_size / 256
|
|
|
|
if not torch.cuda.is_available():
|
|
print('using CPU, this will be slow')
|
|
elif args.distributed:
|
|
# Apply SyncBN
|
|
model = torch.nn.SyncBatchNorm.convert_sync_batchnorm(model)
|
|
# For multiprocessing distributed, DistributedDataParallel constructor
|
|
# should always set the single device scope, otherwise,
|
|
# DistributedDataParallel will use all available devices.
|
|
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 we have
|
|
args.batch_size = int(args.batch_size / args.world_size)
|
|
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:
|
|
torch.cuda.set_device(args.gpu)
|
|
model = model.cuda(args.gpu)
|
|
# comment out the following line for debugging
|
|
# raise NotImplementedError("Only DistributedDataParallel is supported.")
|
|
else:
|
|
# AllGather/rank implementation in this code only supports DistributedDataParallel.
|
|
raise NotImplementedError("Only DistributedDataParallel is supported.")
|
|
print(model) # print model after SyncBatchNorm
|
|
|
|
# define loss function (criterion) and optimizer
|
|
criterion = nn.CrossEntropyLoss().cuda(args.gpu)
|
|
|
|
if args.optimizer == 'lars':
|
|
optimizer = moco.optimizer.LARS(model.parameters(), init_lr,
|
|
weight_decay=args.weight_decay,
|
|
momentum=args.momentum)
|
|
elif args.optimizer == 'adamw':
|
|
optimizer = moco.optimizer.AdamW(model.parameters(), init_lr,
|
|
weight_decay=args.weight_decay)
|
|
|
|
scaler = torch.cuda.amp.GradScaler()
|
|
|
|
# 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)
|
|
else:
|
|
# 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']
|
|
model.load_state_dict(checkpoint['state_dict'])
|
|
optimizer.load_state_dict(checkpoint['optimizer'])
|
|
scaler.load_state_dict(checkpoint['scaler'])
|
|
print("=> loaded checkpoint '{}' (epoch {})"
|
|
.format(args.resume, checkpoint['epoch']))
|
|
else:
|
|
print("=> no checkpoint found at '{}'".format(args.resume))
|
|
|
|
cudnn.benchmark = True
|
|
|
|
# Data loading code
|
|
traindir = os.path.join(args.data, 'train')
|
|
normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406],
|
|
std=[0.229, 0.224, 0.225])
|
|
|
|
# BYOL's augmentation recipe: https://arxiv.org/abs/2006.07733
|
|
# except min-scale kept as 0.2
|
|
augmentation1 = [
|
|
transforms.RandomResizedCrop(224, scale=(0.2, 1.)),
|
|
transforms.RandomApply([
|
|
transforms.ColorJitter(0.4, 0.4, 0.2, 0.1) # not strengthened
|
|
], p=0.8),
|
|
transforms.RandomGrayscale(p=0.2),
|
|
transforms.RandomApply([moco.loader.GaussianBlur([.1, 2.])], p=1.0),
|
|
transforms.RandomHorizontalFlip(),
|
|
transforms.ToTensor(),
|
|
normalize
|
|
]
|
|
|
|
augmentation2 = [
|
|
transforms.RandomResizedCrop(224, scale=(0.2, 1.)),
|
|
transforms.RandomApply([
|
|
transforms.ColorJitter(0.4, 0.4, 0.2, 0.1) # not strengthened
|
|
], p=0.8),
|
|
transforms.RandomGrayscale(p=0.2),
|
|
transforms.RandomApply([moco.loader.GaussianBlur([.1, 2.])], p=0.1),
|
|
transforms.RandomApply([moco.loader.Solarize()], p=0.2),
|
|
transforms.RandomHorizontalFlip(),
|
|
transforms.ToTensor(),
|
|
normalize
|
|
]
|
|
|
|
train_dataset = datasets.ImageFolder(
|
|
traindir,
|
|
moco.loader.TwoCropsTransform(transforms.Compose(augmentation1),
|
|
transforms.Compose(augmentation2)))
|
|
|
|
if args.distributed:
|
|
train_sampler = torch.utils.data.distributed.DistributedSampler(train_dataset)
|
|
else:
|
|
train_sampler = None
|
|
|
|
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, drop_last=True)
|
|
|
|
for epoch in range(args.start_epoch, args.epochs):
|
|
if args.distributed:
|
|
train_sampler.set_epoch(epoch)
|
|
adjust_learning_rate(optimizer, init_lr, epoch, args)
|
|
|
|
# train for one epoch
|
|
train(train_loader, model, criterion, optimizer, scaler, epoch, args)
|
|
|
|
if not args.multiprocessing_distributed or (args.multiprocessing_distributed
|
|
and args.rank == 0): # only the first GPU saves checkpoint
|
|
save_checkpoint({
|
|
'epoch': epoch + 1,
|
|
'arch': args.arch,
|
|
'state_dict': model.state_dict(),
|
|
'optimizer' : optimizer.state_dict(),
|
|
'scaler': scaler.state_dict(),
|
|
}, is_best=False, filename='%s/checkpoint_%04d.pth.tar' % (args.checkpoint_folder, epoch))
|
|
|
|
|
|
def train(train_loader, model, criterion, optimizer, scaler, epoch, 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()
|
|
moco_m = adjust_moco_momentum(epoch, args)
|
|
for i, (images, _) in enumerate(train_loader):
|
|
# measure data loading time
|
|
data_time.update(time.time() - end)
|
|
|
|
if args.gpu is not None:
|
|
images[0] = images[0].cuda(args.gpu, non_blocking=True)
|
|
images[1] = images[1].cuda(args.gpu, non_blocking=True)
|
|
|
|
# compute output
|
|
with torch.cuda.amp.autocast(True):
|
|
output1, output2, target = model(im1=images[0], im2=images[1], m=moco_m)
|
|
loss = (criterion(output1, target) + criterion(output2, target)) * (args.moco_t * 2.)
|
|
|
|
# acc1/acc5 are N-way contrast classifier accuracy
|
|
# measure accuracy and record loss
|
|
acc1, acc5 = accuracy(output1, target, topk=(1, 5))
|
|
losses.update(loss.item(), images[0].size(0))
|
|
top1.update(acc1[0], images[0].size(0))
|
|
top5.update(acc5[0], images[0].size(0))
|
|
|
|
# compute gradient and do SGD step
|
|
optimizer.zero_grad()
|
|
scaler.scale(loss).backward()
|
|
scaler.step(optimizer)
|
|
scaler.update()
|
|
|
|
# measure elapsed time
|
|
batch_time.update(time.time() - end)
|
|
end = time.time()
|
|
|
|
if i % args.print_freq == 0:
|
|
progress.display(i)
|
|
|
|
|
|
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 AverageMeter(object):
|
|
"""Computes and stores the average and current value"""
|
|
def __init__(self, name, fmt=':f'):
|
|
self.name = name
|
|
self.fmt = fmt
|
|
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 __str__(self):
|
|
fmtstr = '{name} {val' + self.fmt + '} ({avg' + self.fmt + '})'
|
|
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]
|
|
print('\t'.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 adjust_learning_rate(optimizer, init_lr, epoch, args):
|
|
"""Decays the learning rate with half-cycle cosine after warmup"""
|
|
if epoch < args.warmup_epochs:
|
|
lr = init_lr / (args.warmup_epochs + 1) * (epoch + 1)
|
|
else:
|
|
lr = init_lr * 0.5 * (1. + math.cos(math.pi * (epoch - args.warmup_epochs) / (args.epochs - args.warmup_epochs)))
|
|
for param_group in optimizer.param_groups:
|
|
param_group['lr'] = lr
|
|
|
|
|
|
def adjust_moco_momentum(epoch, args):
|
|
"""Adjust moco momentum based on current epoch"""
|
|
return 1. - 0.5 * (1. + math.cos(math.pi * epoch / args.epochs)) * (1. - args.moco_m)
|
|
|
|
|
|
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()
|