path: root/supervised/baseline.py

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 tqdm import tqdm

from datautils import color_distortion, Clip, RandomGaussianBlur
from models import CIFARResNet50, ImageNetResNet50
from optimizers import LARS
from schedulers import LinearWarmupAndCosineAnneal, LinearLR


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()

TENSORBOARD_PATH = os.path.join('runs', args.codename)
CHECKPOINT_PATH = os.path.join('checkpoints', args.codename)

random.seed(args.seed)
torch.manual_seed(args.seed)
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

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)

    resnet = CIFARResNet50()
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)

    resnet = ImageNetResNet50()
else:
    raise NotImplementedError(f"Dataset '{args.dataset}' is not implemented.")

resnet = resnet.to(device)

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)

num_train_batches = len(train_loader)
num_test_batches = len(test_loader)


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 resnet.named_parameters()
                   if not exclude_from_wd_and_adaptation(name)],
        'weight_decay': args.weight_decay,
        'layer_adaptation': True,
    },
    {
        'params': [p for name, p in resnet.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.")

# Restore checkpoint
if args.restore_epoch > 0:
    checkpoint_path = os.path.join(CHECKPOINT_PATH, f'{args.restore_epoch:04d}.pt')
    checkpoint = torch.load(checkpoint_path)
    resnet.load_state_dict(checkpoint['resnet_state_dict'])
    optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
    print(f'[RESTORED][{args.restore_epoch}/{args.n_epochs}]\t'
          f'Train loss: {checkpoint["train_loss"]:.4f}\t'
          f'Test loss: {checkpoint["test_loss"]:.4f}')

if args.sched == 'warmup-anneal':
    scheduler = LinearWarmupAndCosineAnneal(
        optimizer,
        warm_up=args.warmup_epochs / args.n_epochs,
        T_max=args.n_epochs * num_train_batches,
        last_epoch=args.restore_epoch * num_train_batches - 1
    )
elif args.sched == 'linear':
    scheduler = LinearLR(
        optimizer,
        num_epochs=args.n_epochs * num_train_batches,
        last_epoch=args.restore_epoch * num_train_batches - 1
    )
elif args.sched is None or args.sched == '' or args.sched == 'const':
    scheduler = None
else:
    raise NotImplementedError(f"Scheduler '{args.sched}' is not implemented.")

if args.optim == 'lars':
    optimizer = LARS(optimizer)

criterion = CrossEntropyLoss()

if not os.path.exists(CHECKPOINT_PATH):
    os.makedirs(CHECKPOINT_PATH)
writer = SummaryWriter(TENSORBOARD_PATH)

curr_train_iters = args.restore_epoch * num_train_batches
curr_test_iters = args.restore_epoch * num_test_batches
for epoch in range(args.restore_epoch, args.n_epochs):
    train_loss = 0
    training_progress = tqdm(
        enumerate(train_loader), desc='Train loss: ', total=num_train_batches
    )

    resnet.train()
    for batch, (images, targets) in training_progress:
        images, targets = images.to(device), targets.to(device)

        resnet.zero_grad()
        output = resnet(images)
        loss = criterion(output, targets)
        loss.backward()
        optimizer.step()
        if args.sched:
            scheduler.step()

        train_loss += loss.item()
        train_loss_mean = train_loss / (batch + 1)
        training_progress.set_description(f'Train loss: {train_loss_mean:.4f}')
        writer.add_scalar('Loss/train', loss, curr_train_iters + 1)
        curr_train_iters += 1

    test_loss = 0
    test_acc = 0
    test_progress = tqdm(
        enumerate(test_loader), desc='Test loss: ', total=num_test_batches
    )

    resnet.eval()
    with torch.no_grad():
        for batch, (images, targets) in test_progress:
            images, targets = images.to(device), targets.to(device)

            output = resnet(images)
            loss = criterion(output, targets)
            _, prediction = output.max(-1)

            test_loss += loss
            test_loss_mean = test_loss / (batch + 1)
            test_progress.set_description(f'Test loss: {test_loss_mean:.4f}')
            test_acc += (prediction == targets).float().mean()
            test_acc_mean = test_acc / (batch + 1)
            writer.add_scalar('Loss/test', loss, curr_test_iters + 1)
            curr_test_iters += 1

    train_loss_mean = train_loss / num_train_batches
    test_loss_mean = test_loss / num_test_batches
    test_acc_mean = test_acc / num_test_batches
    print(f'[{epoch + 1}/{args.n_epochs}]\t'
          f'Train loss: {train_loss_mean:.4f}\t'
          f'Test loss: {test_loss_mean:.4f}\t',
          f'Test acc: {test_acc_mean:.4f}')

    writer.add_scalar('Acc', test_acc_mean, epoch + 1)

    torch.save({'epoch': epoch,
                'resnet_state_dict': resnet.state_dict(),
                'optimizer_state_dict': optimizer.state_dict(),
                'train_loss': train_loss_mean, 'test_loss': test_loss_mean,
                }, os.path.join(CHECKPOINT_PATH, f'{epoch + 1:04d}.pt'))