train_sample.py

# -*- coding: utf-8 -*-
# @Author  : DevinYang(pistonyang@gmail.com)

import argparse, time, logging, os
import models
import torch
import warnings
import apex
from scripts.utils import get_model, set_model

from torchtoolbox import metric
from torchtoolbox.nn import LabelSmoothingLoss
from torchtoolbox.optimizer import CosineWarmupLr, Lookahead
from torchtoolbox.nn.init import KaimingInitializer
from torchtoolbox.tools import no_decay_bias, \
    mixup_data, mixup_criterion, check_dir, summary

from torchvision import transforms
from torchvision.datasets import CIFAR10, CIFAR100, MNIST, FashionMNIST, ImageFolder
from torch.utils.data import DataLoader
from torch import nn
from torch.nn import functional as F
from torch import optim
from apex import amp

os.environ['CUDA_VISIBLE_DEVICES'] = '0'

parser = argparse.ArgumentParser(description='Train a model on ImageNet.')
parser.add_argument('--data-path', type=str, required=True,
                    help='training and validation dataset.')
parser.add_argument('--dataset', type=str, default='mnist',
                    help='Which dataset to train, default is mnist')
parser.add_argument('--batch-size', type=int, default=32,
                    help='training batch size per device (CPU/GPU).')
parser.add_argument('--dtype', type=str, default='float32',
                    help='data type for training. default is float32')
parser.add_argument('--devices', type=str, default='0',
                    help='gpus to use.')
parser.add_argument('-j', '--num-data-workers', dest='num_workers', default=4, type=int,
                    help='number of preprocessing workers')
parser.add_argument('--epochs', type=int, default=1,
                    help='number of training epochs.')
parser.add_argument('--lr', type=float, default=0,
                    help='learning rate. default is 0.')
parser.add_argument('--momentum', type=float, default=0.9,
                    help='momentum value for optimizer, default is 0.9.')
parser.add_argument('--wd', type=float, default=0.0001,
                    help='weight decay rate. default is 0.0001.')
parser.add_argument('--dropout', type=float, default=0.,
                    help='model dropout rate.')
parser.add_argument('--sync-bn', action='store_true',
                    help='use Apex Sync-BN.')
parser.add_argument('--lookahead', action='store_true',
                    help='use lookahead optimizer.')
parser.add_argument('--warmup-lr', type=float, default=0.0,
                    help='starting warmup learning rate. default is 0.0.')
parser.add_argument('--warmup-epochs', type=int, default=0,
                    help='number of warmup epochs.')
parser.add_argument('--model', type=str, required=True,
                    help='type of model to use. see vision_model for options.')
parser.add_argument('--alpha', type=float, default=0,
                    help='model param.')
parser.add_argument('--input-size', type=int, default=32,
                    help='size of the input image size. default is 224')
parser.add_argument('--padding', type=int, default=4,
                    help='pad input image')
parser.add_argument('--norm-layer', type=str, default='',
                    help='Norm layer to use.')
parser.add_argument('--activation', type=str, default='',
                    help='activation to use.')
parser.add_argument('--mixup', action='store_true',
                    help='whether train the model with mix-up. default is false.')
parser.add_argument('--mixup-alpha', type=float, default=0.2,
                    help='beta distribution parameter for mixup sampling, default is 0.2.')
parser.add_argument('--mixup-off-epoch', type=int, default=0,
                    help='how many epochs to train without mixup, default is 0.')
parser.add_argument('--label-smoothing', action='store_true',
                    help='use label smoothing or not in training. default is false.')
parser.add_argument('--no-wd', action='store_true',
                    help='whether to remove weight decay on bias, and beta/gamma for batchnorm layers.')
parser.add_argument('--save-dir', type=str, default='params',
                    help='directory of saved models')
parser.add_argument('--log-interval', type=int, default=50,
                    help='Number of batches to wait before logging.')
parser.add_argument('--logging-file', type=str, default='train_sample.log',
                    help='name of training log file')
parser.add_argument('--resume-epoch', type=int, default=0,
                    help='epoch to resume training from.')
parser.add_argument('--resume-param', type=str, default='',
                    help='resume training param path.')
parser.add_argument("--local_rank", default=0, type=int)
args = parser.parse_args()


def get_dataset(name):
    path = args.data_path
    download = not os.path.exists(path) or not os.listdir(path)
    if name == 'cifar10':
        train_ = CIFAR10(path, train=True, transform=train_transform, download=download)
        val_ = CIFAR10(path, train=False, transform=val_transform, download=download)
    elif name == 'cifar100':
        train_ = CIFAR100(path, train=True, transform=train_transform, download=download)
        val_ = CIFAR100(path, train=False, transform=val_transform, download=download)
    elif name == 'mnist':
        train_ = MNIST(path, train=True, transform=train_transform, download=download)
        val_ = MNIST(path, train=False, transform=val_transform, download=download)
    elif name == 'fashion_mnist':
        train_ = FashionMNIST(path, train=True, transform=train_transform, download=download)
        val_ = FashionMNIST(path, train=False, transform=val_transform, download=download)
    else:
        train_ = ImageFolder(os.path.join(path, 'train'), transform=train_transform)
        val_ = ImageFolder(os.path.join(path, 'val'), transform=val_transform)
    return train_, val_


filehandler = logging.FileHandler(args.logging_file)
streamhandler = logging.StreamHandler()

logger = logging.getLogger('')
logger.setLevel(logging.INFO)
logger.addHandler(filehandler)
logger.addHandler(streamhandler)

logger.info(args)
warnings.filterwarnings("ignore", "(Possibly )?corrupt EXIF data", UserWarning)

torch.backends.cudnn.benchmark = True

train_transform = transforms.Compose([
    transforms.Pad(args.padding),
    transforms.RandomCrop(args.input_size),
    # Cutout(),
    transforms.RandomHorizontalFlip(),
    transforms.ToTensor(),
])

val_transform = transforms.ToTensor()
train_set, val_set = get_dataset(args.dataset)

classes = len(train_set.classes)
num_training_samples = len(train_set)

check_dir(args.save_dir)
assert torch.cuda.is_available(), \
    "Please don't waste of your time,it's impossible to train on CPU."

device = torch.device("cuda:0")
device_ids = args.devices.strip().split(',')
device_ids = [int(device) for device in device_ids]

dtype = args.dtype
epochs = args.epochs
resume_epoch = args.resume_epoch
num_workers = args.num_workers
initializer = KaimingInitializer()
batch_size = args.batch_size * len(device_ids)
batches_pre_epoch = num_training_samples // batch_size
lr = 0.1 * (args.batch_size // 32) if args.lr == 0 else args.lr

train_data = DataLoader(train_set, batch_size, False, pin_memory=True, num_workers=num_workers, drop_last=True)
val_data = DataLoader(val_set, batch_size, False, pin_memory=True, num_workers=num_workers, drop_last=False)

model_setting = set_model(args.dropout, args.norm_layer, args.activation)

try:
    model = get_model(models, args.model, alpha=args.alpha, small_input=True,
                      return_feature=True, norm_feature=True, **model_setting)
except TypeError:
    model = get_model(models, args.model, small_input=True,
                      return_feature=True, norm_feature=True, **model_setting)

summary(model, torch.rand((1, 3, args.input_size, args.input_size)))
model.apply(initializer)
model.to(device)
parameters = model.parameters() if not args.no_wd else no_decay_bias(model)
optimizer = optim.SGD(parameters, lr=lr, momentum=args.momentum,
                      weight_decay=args.wd, nesterov=True)

if args.sync_bn:
    logger.info('Use Apex Synced BN.')
    model = apex.parallel.convert_syncbn_model(model)

if dtype == 'float16':
    logger.info('Train with FP16.')
    model, optimizer = amp.initialize(model, optimizer, opt_level='O1')

if args.lookahead:
    logger.info('Use lookahead optimizer.')
    optimizer = Lookahead(optimizer)

model = nn.DataParallel(model)
lr_scheduler = CosineWarmupLr(optimizer, batches_pre_epoch, epochs,
                              base_lr=args.lr, warmup_epochs=args.warmup_epochs)
if resume_epoch > 0:
    checkpoint = torch.load(args.resume_param)
    model.load_state_dict(checkpoint['model'])
    optimizer.load_state_dict(checkpoint['optimizer'])
    amp.load_state_dict(checkpoint['amp'])
    lr_scheduler.load_state_dict(checkpoint['lr_scheduler'])
    print("Finish loading resume param.")

top1_acc = metric.Accuracy(name='Top1 Accuracy')
top5_acc = metric.TopKAccuracy(top=5, name='Top5 Accuracy')
loss_record = metric.NumericalCost(name='Loss')

Loss = nn.CrossEntropyLoss() if not args.label_smoothing else \
    LabelSmoothingLoss(classes, smoothing=0.1)


@torch.no_grad()
def test(epoch=0, save_status=False):
    top1_acc.reset()
    top5_acc.reset()
    loss_record.reset()
    model.eval()
    for data, labels in val_data:
        data = data.to(device, non_blocking=True)
        labels = labels.to(device, non_blocking=True)

        outputs = model(data)
        losses = Loss(outputs, labels)

        top1_acc.update(outputs, labels)
        top5_acc.update(outputs, labels)
        loss_record.update(losses)

    test_msg = 'Test Epoch {}: {}:{:.5}, {}:{:.5}, {}:{:.5}\n'.format(
        epoch, top1_acc.name, top1_acc.get(), top5_acc.name, top5_acc.get(),
        loss_record.name, loss_record.get())
    logger.info(test_msg)
    if save_status:
        checkpoint = {
            'model': model.state_dict(),
            'optimizer': optimizer.state_dict(),
            'amp': amp.state_dict(),
            'lr_scheduler': lr_scheduler.state_dict(),
        }
        torch.save(checkpoint, '{}/{}_{}_{:.5}.pt'.format(
            args.save_dir, args.model, epoch, top1_acc.get()))


def train():
    for epoch in range(resume_epoch, epochs):
        top1_acc.reset()
        loss_record.reset()
        tic = time.time()

        model.train()
        for i, (data, labels) in enumerate(train_data):
            data = data.to(device, non_blocking=True)
            labels = labels.to(device, non_blocking=True)

            optimizer.zero_grad()
            outputs = model(data)
            loss = Loss(outputs, labels)

            with amp.scale_loss(loss, optimizer) as scaled_loss:
                scaled_loss.backward()
            optimizer.step()

            lr_scheduler.step()
            top1_acc.update(outputs, labels)
            loss_record.update(loss)

            if i % args.log_interval == 0 and i != 0:
                logger.info('Epoch {}, Iter {}, {}:{:.5}, {}:{:.5}, {} samples/s. lr: {:.5}.'.format(
                    epoch, i, top1_acc.name, top1_acc.get(),
                    loss_record.name, loss_record.get(),
                    int((i * batch_size) // (time.time() - tic)),
                    lr_scheduler.learning_rate
                ))

        train_speed = int(num_training_samples // (time.time() - tic))
        epoch_msg = 'Train Epoch {}: {}:{:.5}, {}:{:.5}, {} samples/s.'.format(
            epoch, top1_acc.name, top1_acc.get(), loss_record.name, loss_record.get(), train_speed)
        logger.info(epoch_msg)
        test(epoch)


def train_mixup():
    mixup_off_epoch = epochs if args.mixup_off_epoch == 0 else args.mixup_off_epoch
    for epoch in range(resume_epoch, epochs):
        loss_record.reset()
        alpha = args.mixup_alpha if epoch < mixup_off_epoch else 0
        tic = time.time()

        model.train()
        for i, (data, labels) in enumerate(train_data):
            data = data.to(device, non_blocking=True)
            labels = labels.to(device, non_blocking=True)

            data, labels_a, labels_b, lam = mixup_data(data, labels, alpha)
            optimizer.zero_grad()
            outputs = model(data)
            loss = mixup_criterion(Loss, outputs, labels_a, labels_b, lam)

            with amp.scale_loss(loss, optimizer) as scaled_loss:
                scaled_loss.backward()
            optimizer.step()

            loss_record.update(loss)
            lr_scheduler.step()

            if i % args.log_interval == 0 and i != 0:
                logger.info('Epoch {}, Iter {}, {}:{:.5}, {} samples/s.'.format(
                    epoch, i, loss_record.name, loss_record.get(),
                    int((i * batch_size) // (time.time() - tic))
                ))

        train_speed = int(num_training_samples // (time.time() - tic))
        train_msg = 'Train Epoch {}: {}:{:.5}, {} samples/s, lr:{:.5}'.format(
            epoch, loss_record.name, loss_record.get(),
            train_speed, lr_scheduler.learning_rate)
        logger.info(train_msg)
        test(epoch)


if __name__ == '__main__':
    if args.mixup:
        logger.info('Train using Mixup.')
        train_mixup()
    else:
        train()