import argparse import os import random import torch from torch.nn import CrossEntropyLoss from torch.utils.data import DataLoader from torch.utils.tensorboard import SummaryWriter from torchvision.datasets import CIFAR10, ImageNet from torchvision.transforms import transforms, InterpolationMode from datautils import color_distortion, Clip, RandomGaussianBlur from models import CIFARResNet50, ImageNetResNet50 from optimizers import LARS from schedulers import LinearWarmupAndCosineAnneal, LinearLR from utils import training_log def build_parser(): def range_parser(range_string: str): try: range_ = tuple(map(float, range_string.split('-'))) return range_ except: raise argparse.ArgumentTypeError("Range must be 'start-end.'") parser = argparse.ArgumentParser(description='Supervised baseline') parser.add_argument('--codename', default='cifar10-resnet50-256-lars-warmup', type=str, help="Model descriptor (default: " "'cifar10-resnet50-256-lars-warmup')") parser.add_argument('--seed', default=0, type=int, help='Random seed for reproducibility (default: 0)') data_group = parser.add_argument_group('Dataset parameters') data_group.add_argument('--dataset_dir', default='dataset', type=str, help="Path to dataset directory (default: 'dataset')") data_group.add_argument('--dataset', default='cifar10', type=str, help="Name of dataset (default: 'cifar10')") data_group.add_argument('--crop_size', default=32, type=int, help='Random crop size after resize (default: 32)') data_group.add_argument('--crop_scale', default='0.8-1', type=range_parser, help='Random resize scale range (default: 0.8-1)') data_group.add_argument('--hflip_p', default=0.5, type=float, help='Random horizontal flip probability (default: 0.5)') data_group.add_argument('--distort_s', default=0.5, type=float, help='Distortion strength (default: 0.5)') data_group.add_argument('--gaussian_ker_scale', default=10, type=float, help='Gaussian kernel scale factor ' '(equals to img_size / kernel_size) (default: 10)') data_group.add_argument('--gaussian_sigma', default='0.1-2', type=range_parser, help='Random gaussian blur sigma range (default: 0.1-2)') data_group.add_argument('--gaussian_p', default=0.5, type=float, help='Random gaussian blur probability (default: 0.5)') train_group = parser.add_argument_group('Training parameters') train_group.add_argument('--batch_size', default=256, type=int, help='Batch size (default: 256)') train_group.add_argument('--restore_epoch', default=0, type=int, help='Restore epoch, 0 for training from scratch ' '(default: 0)') train_group.add_argument('--n_epochs', default=1000, type=int, help='Number of epochs (default: 1000)') train_group.add_argument('--warmup_epochs', default=10, type=int, help='Epochs for warmup ' '(only for `warmup-anneal` scheduler) (default: 10)') train_group.add_argument('--n_workers', default=2, type=int, help='Number of dataloader processes (default: 2)') train_group.add_argument('--optim', default='lars', type=str, help="Name of optimizer (default: 'lars')") train_group.add_argument('--sched', default='warmup-anneal', type=str, help="Name of scheduler (default: 'warmup-anneal')") train_group.add_argument('--lr', default=1, type=float, help='Learning rate (default: 1)') train_group.add_argument('--momentum', default=0.9, type=float, help='Momentum (default: 0.9') train_group.add_argument('--weight_decay', default=1e-6, type=float, help='Weight decay (l2 regularization) (default: 1e-6)') args = parser.parse_args() args.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') args.checkpoint_root = os.path.join('checkpoints', args.codename) args.tensorboard_root = os.path.join('runs', args.codename) return args def prepare_dataset(args): if args.dataset == 'cifar10' or args.dataset == 'cifar': train_transform = transforms.Compose([ transforms.RandomResizedCrop( args.crop_size, scale=args.crop_scale, interpolation=InterpolationMode.BICUBIC ), transforms.RandomHorizontalFlip(args.hflip_p), color_distortion(args.distort_s), transforms.ToTensor(), Clip() ]) test_transform = transforms.Compose([ transforms.ToTensor() ]) train_set = CIFAR10(args.dataset_dir, train=True, transform=train_transform, download=True) test_set = CIFAR10(args.dataset_dir, train=False, transform=test_transform) elif args.dataset == 'imagenet1k' or args.dataset == 'imagenet1k': train_transform = transforms.Compose([ transforms.RandomResizedCrop( args.crop_size, scale=args.crop_scale, interpolation=InterpolationMode.BICUBIC ), transforms.RandomHorizontalFlip(args.hflip_p), color_distortion(args.distort_s), transforms.ToTensor(), RandomGaussianBlur( kernel_size=args.crop_size // args.gaussian_ker_scale, sigma_range=args.gaussian_sigma, p=args.gaussian_p ), Clip() ]) test_transform = transforms.Compose([ transforms.Resize(256), transforms.CenterCrop(args.crop_size), transforms.ToTensor(), ]) train_set = ImageNet(args.dataset_dir, 'train', transform=train_transform) test_set = ImageNet(args.dataset_dir, 'val', transform=test_transform) else: raise NotImplementedError(f"Dataset '{args.dataset}' is not implemented.") return train_set, test_set def create_dataloader(args, train_set, test_set): train_loader = DataLoader(train_set, batch_size=args.batch_size, shuffle=True, num_workers=args.n_workers) test_loader = DataLoader(test_set, batch_size=args.batch_size, shuffle=False, num_workers=args.n_workers) args.num_train_batches = len(train_loader) args.num_test_batches = len(test_loader) return train_loader, test_loader def init_model(args): if args.dataset == 'cifar10' or args.dataset == 'cifar': model = CIFARResNet50() elif args.dataset == 'imagenet1k' or args.dataset == 'imagenet1k': model = ImageNetResNet50() else: raise NotImplementedError(f"Dataset '{args.dataset}' is not implemented.") return model def configure_optimizer(args, model): def exclude_from_wd_and_adaptation(name): if 'bn' in name: return True if args.optim == 'lars' and 'bias' in name: return True param_groups = [ { 'params': [p for name, p in model.named_parameters() if not exclude_from_wd_and_adaptation(name)], 'weight_decay': args.weight_decay, 'layer_adaptation': True, }, { 'params': [p for name, p in model.named_parameters() if exclude_from_wd_and_adaptation(name)], 'weight_decay': 0., 'layer_adaptation': False, }, ] if args.optim == 'adam': optimizer = torch.optim.Adam( param_groups, lr=args.lr, betas=(args.momentum, 0.999) ) elif args.optim == 'sdg' or args.optim == 'lars': optimizer = torch.optim.SGD( param_groups, lr=args.lr, momentum=args.momentum ) else: raise NotImplementedError(f"Optimizer '{args.optim}' is not implemented.") return optimizer @training_log def load_checkpoint(args, model, optimizer): checkpoint_path = os.path.join(args.checkpoint_root, f'{args.restore_epoch:04d}.pt') checkpoint = torch.load(checkpoint_path) model.load_state_dict(checkpoint.pop('model_state_dict')) optimizer.load_state_dict(checkpoint.pop('optimizer_state_dict')) return checkpoint def configure_scheduler(args, optimizer): n_iters = args.n_epochs * args.num_train_batches last_iter = args.restore_epoch * args.num_train_batches - 1 if args.sched == 'warmup-anneal': scheduler = LinearWarmupAndCosineAnneal( optimizer, warm_up=args.warmup_epochs / args.n_epochs, T_max=n_iters, last_epoch=last_iter ) elif args.sched == 'linear': scheduler = LinearLR( optimizer, num_epochs=n_iters, last_epoch=last_iter ) elif args.sched is None or args.sched == '' or args.sched == 'const': scheduler = None else: raise NotImplementedError(f"Scheduler '{args.sched}' is not implemented.") return scheduler def wrap_lars(args, optimizer): if args.optim == 'lars': return LARS(optimizer) else: return optimizer def train(args, train_loader, model, loss_fn, optimizer): model.train() for batch, (images, targets) in enumerate(train_loader): images, targets = images.to(args.device), targets.to(args.device) model.zero_grad() output = model(images) loss = loss_fn(output, targets) loss.backward() optimizer.step() yield batch, loss.item() def eval(args, test_loader, model, loss_fn): model.eval() with torch.no_grad(): for batch, (images, targets) in enumerate(test_loader): images, targets = images.to(args.device), targets.to(args.device) output = model(images) loss = loss_fn(output, targets) prediction = output.argmax(1) accuracy = (prediction == targets).float().mean() yield batch, loss.item(), accuracy.item() @training_log def batch_logger(args, writer, batch, epoch, loss, lr): global_batch = epoch * args.num_train_batches + batch writer.add_scalar('Batch loss/train', loss, global_batch + 1) writer.add_scalar('Batch lr/train', lr, global_batch + 1) return { 'batch': batch + 1, 'n_batches': args.num_train_batches, 'global_batch': global_batch + 1, 'epoch': epoch + 1, 'n_epochs': args.n_epochs, 'train_loss': loss, 'lr': lr, } @training_log def epoch_logger(args, writer, epoch, train_loss, test_loss, test_accuracy): train_loss_mean = train_loss.mean().item() test_loss_mean = test_loss.mean().item() test_accuracy_mean = test_accuracy.mean().item() writer.add_scalar('Epoch loss/train', train_loss_mean, epoch + 1) writer.add_scalar('Epoch loss/test', test_loss_mean, epoch + 1) writer.add_scalar('Accuracy/test', test_accuracy_mean, epoch + 1) return { 'epoch': epoch + 1, 'n_epochs': args.n_epochs, 'train_loss': train_loss_mean, 'test_loss': test_loss_mean, 'test_accuracy': test_accuracy_mean } def save_checkpoint(args, epoch_log, model, optimizer): if not os.path.exists(args.checkpoint_root): os.makedirs(args.checkpoint_root) torch.save(epoch_log | { 'model_state_dict': model.state_dict(), 'optimizer_state_dict': optimizer.state_dict(), }, os.path.join(args.checkpoint_root, f"{epoch_log['epoch']:04d}.pt")) if __name__ == '__main__': args = build_parser() random.seed(args.seed) torch.manual_seed(args.seed) train_set, test_set = prepare_dataset(args) train_loader, test_loader = create_dataloader(args, train_set, test_set) resnet = init_model(args).to(args.device) xent = CrossEntropyLoss() optimizer = configure_optimizer(args, resnet) if args.restore_epoch > 0: load_checkpoint(args, resnet, optimizer) scheduler = configure_scheduler(args, optimizer) optimizer = wrap_lars(args, optimizer) writer = SummaryWriter(args.tensorboard_root) for epoch in range(args.restore_epoch, args.n_epochs): train_loss = torch.zeros(args.num_train_batches, device=args.device) test_loss = torch.zeros(args.num_test_batches, device=args.device) test_accuracy = torch.zeros(args.num_test_batches, device=args.device) for batch, loss in train(args, train_loader, resnet, xent, optimizer): train_loss[batch] = loss batch_logger(args, writer, batch, epoch, loss, optimizer.param_groups[0]['lr']) if scheduler and batch != args.num_train_batches - 1: scheduler.step() for batch, loss, accuracy in eval(args, test_loader, resnet, xent): test_loss[batch] = loss test_accuracy[batch] = accuracy epoch_log = epoch_logger(args, writer, epoch, train_loss, test_loss, test_accuracy) save_checkpoint(args, epoch_log, resnet, optimizer) # Step after save checkpoint, otherwise the schedular # will one iter ahead after restore if scheduler: scheduler.step()