FAA ver2.0 for colab

+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+# C extensions
+*.so
+
+# Distribution / packaging
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+
+# PyInstaller
+#  Usually these files are written by a python script from a template
+#  before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+.hypothesis/
+.pytest_cache/
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+local_settings.py
+db.sqlite3
+
+# Flask stuff:
+instance/
+.webassets-cache
+
+# Scrapy stuff:
+.scrapy
+
+# Sphinx documentation
+docs/_build/
+
+# PyBuilder
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# pyenv
+.python-version
+
+# celery beat schedule file
+celerybeat-schedule
+
+# SageMath parsed files
+*.sage.py
+
+# Environments
+.env
+.venv
+env/
+venv/
+ENV/
+env.bak/
+venv.bak/
+
+# Spyder project settings
+.spyderproject
+.spyproject
+
+# Rope project settings
+.ropeproject
+
+# mkdocs documentation
+/site
+
+# mypy
+.mypy_cache/
+

+# Fast Autoaugment
+<img src="figures/faa.png" width=800px>
+
+A Pytorch Implementation of [Fast AutoAugment](https://arxiv.org/pdf/1905.00397.pdf) and [EfficientNet](https://arxiv.org/abs/1905.11946).
+
+## Prerequisite
+* torch==1.1.0
+* torchvision==0.2.2
+* hyperopt==0.1.2
+* future==0.17.1
+* tb-nightly==1.15.0a20190622
+
+## Usage
+### Training
+#### CIFAR10
+```bash
+# ResNet20 (w/o FastAutoAugment)
+python train.py --seed=24 --scale=3 --optimizer=sgd --fast_auto_augment=False
+
+# ResNet20 (w/ FastAutoAugment)
+python train.py --seed=24 --scale=3 --optimizer=sgd --fast_auto_augment=True
+
+# ResNet20 (w/ FastAutoAugment, Pre-found policy)
+python train.py --seed=24 --scale=3 --optimizer=sgd --fast_auto_augment=True \
+                --augment_path=runs/ResNet_Scale3_FastAutoAugment/augmentation.cp
+
+# ResNet32 (w/o FastAutoAugment)
+python train.py --seed=24 --scale=5 --optimizer=sgd --fast_auto_augment=False
+
+# ResNet32 (w/ FastAutoAugment)
+python train.py --seed=24 --scale=5 --optimizer=sgd --fast_auto_augment=True
+
+# EfficientNet (w/ FastAutoAugment)
+python train.py --seed=24 --pi=0 --optimizer=adam --fast_auto_augment=True \
+                --network=efficientnet_cifar10 --activation=swish
+```
+
+#### ImageNet (You can use any backbone networks in [torchvision.models](https://pytorch.org/docs/stable/torchvision/models.html))
+```bash
+
+# BaseNet (w/o FastAutoAugment)
+python train.py --seed=24 --dataset=imagenet --optimizer=adam --network=resnet50
+
+# EfficientNet (w/ FastAutoAugment) (UnderConstruction)
+python train.py --seed=24 --dataset=imagenet --pi=0 --optimizer=adam --fast_auto_augment=True \
+                --network=efficientnet --activation=swish
+```
+
+### Eval
+```bash
+# Single Image testing
+python eval.py --model_path=runs/ResNet_Scale3_Basline
+
+# 5-crops testing
+python eval.py --model_path=runs/ResNet_Scale3_Basline --five_crops=True
+```
+
+## Experiments
+### Fast AutoAugment
+#### ResNet20 (CIFAR10)
+* Pre-trained model [[Download](https://drive.google.com/file/d/12D8050yGGiKWGt8_R8QTlkoQ6wq_icBn/view?usp=sharing)]
+* Validation Curve
+<img src="figures/resnet20_valid.png">
+
+* Evaluation (Acc @1)
+
+|                | Valid | Test(Single) |
+|----------------|-------|-------------|
+| ResNet20       | 90.70 | **91.45**   |
+| ResNet20 + FAA |**92.46**| **91.45** |
+
+#### ResNet34 (CIFAR10)
+* Validation Curve
+<img src="figures/resnet34_valid.png">
+
+* Evaluation (Acc @1)
+
+|                | Valid | Test(Single) |
+|----------------|-------|-------------|
+| ResNet34       | 91.54 | 91.47       |
+| ResNet34 + FAA |**92.76**| **91.99** |
+
+### Found Policy [[Download](https://drive.google.com/file/d/1Ia_IxPY3-T7m8biyl3QpxV1s5EA5gRDF/view?usp=sharing)]
+<img src="figures/pm.png">
+
+### Augmented images
+<img src="figures/augmented_images.png">
+<img src="figures/augmented_images2.png">

+import os
+import fire
+import json
+from pprint import pprint
+
+import torch
+import torch.nn as nn
+
+from utils import *
+
+
+def eval(model_path):
+    print('\n[+] Parse arguments')
+    kwargs_path = os.path.join(model_path, 'kwargs.json')
+    kwargs = json.loads(open(kwargs_path).read())
+    args, kwargs = parse_args(kwargs)
+    pprint(args)
+    device = torch.device('cuda' if args.use_cuda else 'cpu')
+
+    print('\n[+] Create network')
+    model = select_model(args)
+    optimizer = select_optimizer(args, model)
+    criterion = nn.CrossEntropyLoss()
+    if args.use_cuda:
+        model = model.cuda()
+        criterion = criterion.cuda()
+
+    print('\n[+] Load model')
+    weight_path = os.path.join(model_path, 'model', 'model.pt')
+    model.load_state_dict(torch.load(weight_path))
+
+    print('\n[+] Load dataset')
+    test_transform = get_valid_transform(args, model)
+    test_dataset = get_dataset(args, test_transform, 'test')
+    test_loader = iter(get_dataloader(args, test_dataset))
+
+    print('\n[+] Start testing')
+    _test_res = validate(args, model, criterion, test_loader, step=0, writer=None)
+
+    print('\n[+] Valid results')
+    print('  Acc@1 : {:.3f}%'.format(_test_res[0].data.cpu().numpy()[0]*100))
+    print('  Acc@5 : {:.3f}%'.format(_test_res[1].data.cpu().numpy()[0]*100))
+    print('  Loss : {:.3f}'.format(_test_res[2].data))
+    print('  Infer Time(per image) : {:.3f}ms'.format(_test_res[3]*1000 / len(test_dataset)))
+
+
+if __name__ == '__main__':
+    fire.Fire(eval)

+import copy
+import json
+import time
+import torch
+import random
+import torchvision.transforms as transforms
+
+from torch.utils.data import Subset
+from sklearn.model_selection import StratifiedShuffleSplit
+from concurrent.futures import ProcessPoolExecutor
+
+from transforms import *
+from hyperopt import fmin, tpe, hp, STATUS_OK, Trials
+from utils import *
+
+
+DEFALUT_CANDIDATES = [
+    ShearXY,
+    TranslateXY,
+    Rotate,
+    AutoContrast,
+    Invert,
+    Equalize,
+    Solarize,
+    Posterize,
+    Contrast,
+    Color,
+    Brightness,
+    Sharpness,
+    Cutout,
+#     SamplePairing,
+]
+
+
+def train_child(args, model, dataset, subset_indx, device=None):
+    optimizer = select_optimizer(args, model)
+    scheduler = select_scheduler(args, optimizer)
+    criterion = nn.CrossEntropyLoss()
+
+    dataset.transform = transforms.Compose([
+        transforms.Resize(32),
+        transforms.ToTensor()])
+    subset = Subset(dataset, subset_indx)
+    data_loader = get_inf_dataloader(args, subset)
+
+    if device:
+        model = model.to(device)
+        criterion = criterion.to(device)
+
+    elif args.use_cuda:
+        model = model.cuda()
+        criterion = criterion.cuda()
+
+        if torch.cuda.device_count() > 1:
+            print('\n[+] Use {} GPUs'.format(torch.cuda.device_count()))
+            model = nn.DataParallel(model)
+
+    start_t = time.time()
+    for step in range(args.start_step, args.max_step):
+        batch = next(data_loader)
+        _train_res = train_step(args, model, optimizer, scheduler, criterion, batch, step, None, device)
+
+        if step % args.print_step == 0:
+            print('\n[+] Training step: {}/{}\tElapsed time: {:.2f}min\tLearning rate: {}\tDevice: {}'.format(
+                step, args.max_step,(time.time()-start_t)/60, optimizer.param_groups[0]['lr'], device))
+
+            print('  Acc@1 : {:.3f}%'.format(_train_res[0].data.cpu().numpy()[0]*100))
+            print('  Acc@5 : {:.3f}%'.format(_train_res[1].data.cpu().numpy()[0]*100))
+            print('  Loss : {}'.format(_train_res[2].data))
+
+    return _train_res
+
+
+def validate_child(args, model, dataset, subset_indx, transform, device=None):
+    criterion = nn.CrossEntropyLoss()
+
+    if device:
+        model = model.to(device)
+        criterion = criterion.to(device)
+
+    elif args.use_cuda:
+        model = model.cuda()
+        criterion = criterion.cuda()
+
+    dataset.transform = transform
+    subset = Subset(dataset, subset_indx)
+    data_loader = get_dataloader(args, subset, pin_memory=False)
+
+    return validate(args, model, criterion, data_loader, 0, None, device)
+
+
+def get_next_subpolicy(transform_candidates, op_per_subpolicy=2):
+    n_candidates = len(transform_candidates)
+    subpolicy = []
+
+    for i in range(op_per_subpolicy):
+        indx = random.randrange(n_candidates)
+        prob = random.random()
+        mag = random.random()
+        subpolicy.append(transform_candidates[indx](prob, mag))
+
+    subpolicy = transforms.Compose([
+        *subpolicy,
+        transforms.Resize(32),
+        transforms.ToTensor()])
+
+    return subpolicy
+
+
+def search_subpolicies(args, transform_candidates, child_model, dataset, Da_indx, B, device):
+    subpolicies = []
+
+    for b in range(B):
+        subpolicy = get_next_subpolicy(transform_candidates)
+        val_res = validate_child(args, child_model, dataset, Da_indx, subpolicy, device)
+        subpolicies.append((subpolicy, val_res[2]))
+
+    return subpolicies
+
+
+def search_subpolicies_hyperopt(args, transform_candidates, child_model, dataset, Da_indx, B, device):
+
+    def _objective(sampled):
+        subpolicy = [transform(prob, mag)
+                     for transform, prob, mag in sampled]
+
+        subpolicy = transforms.Compose([
+            transforms.Resize(32),
+            *subpolicy,
+            transforms.ToTensor()])
+
+        val_res = validate_child(args, child_model, dataset, Da_indx, subpolicy, device)
+        loss = val_res[2].cpu().numpy()
+        return {'loss': loss, 'status': STATUS_OK }
+
+    space = [(hp.choice('transform1', transform_candidates), hp.uniform('prob1', 0, 1.0), hp.uniform('mag1', 0, 1.0)),
+             (hp.choice('transform2', transform_candidates), hp.uniform('prob2', 0, 1.0), hp.uniform('mag2', 0, 1.0))]
+
+    trials = Trials()
+    best = fmin(_objective,
+                space=space,
+                algo=tpe.suggest,
+                max_evals=B,
+                trials=trials)
+
+    subpolicies = []
+    for t in trials.trials:
+        vals = t['misc']['vals']
+        subpolicy = [transform_candidates[vals['transform1'][0]](vals['prob1'][0], vals['mag1'][0]),
+                     transform_candidates[vals['transform2'][0]](vals['prob2'][0], vals['mag2'][0])]
+        subpolicy = transforms.Compose([
+            ## baseline augmentation
+            transforms.Pad(4),
+            transforms.RandomCrop(32),
+            transforms.RandomHorizontalFlip(),
+            ## policy
+            *subpolicy,
+            ## to tensor
+            transforms.ToTensor()])
+        subpolicies.append((subpolicy, t['result']['loss']))
+
+    return subpolicies
+
+
+def get_topn_subpolicies(subpolicies, N=10):
+    return sorted(subpolicies, key=lambda subpolicy: subpolicy[1])[:N]
+
+
+def process_fn(args_str, model, dataset, Dm_indx, Da_indx, T, transform_candidates, B, N, k):
+    kwargs = json.loads(args_str)
+    args, kwargs = parse_args(kwargs)
+    device_id = k % torch.cuda.device_count()
+    device = torch.device('cuda:%d' % device_id)
+    _transform = []
+
+    print('[+] Child %d training strated (GPU: %d)' % (k, device_id))
+
+    # train child model
+    child_model = copy.deepcopy(model)
+    train_res = train_child(args, child_model, dataset, Dm_indx, device)
+
+    # search sub policy
+    for t in range(T):
+        #subpolicies = search_subpolicies(args, transform_candidates, child_model, dataset, Da_indx, B, device)
+        subpolicies = search_subpolicies_hyperopt(args, transform_candidates, child_model, dataset, Da_indx, B, device)
+        subpolicies = get_topn_subpolicies(subpolicies, N)
+        _transform.extend([subpolicy[0] for subpolicy in subpolicies])
+
+    return _transform
+
+
+def fast_auto_augment(args, model, transform_candidates=None, K=5, B=100, T=2, N=10, num_process=5):
+    args_str = json.dumps(args._asdict())
+    dataset = get_dataset(args, None, 'trainval')
+    num_process = min(torch.cuda.device_count(), num_process)
+    transform, futures = [], []
+
+    torch.multiprocessing.set_start_method('spawn', force=True)
+
+    if not transform_candidates:
+        transform_candidates = DEFALUT_CANDIDATES
+
+    # split
+    Dm_indexes, Da_indexes = split_dataset(args, dataset, K)
+
+    with ProcessPoolExecutor(max_workers=num_process) as executor:
+        for k, (Dm_indx, Da_indx) in enumerate(zip(Dm_indexes, Da_indexes)):
+            future = executor.submit(process_fn,
+                    args_str, model, dataset, Dm_indx, Da_indx, T, transform_candidates, B, N, k)
+            futures.append(future)
+
+        for future in futures:
+            transform.extend(future.result())
+
+    transform = transforms.RandomChoice(transform)
+
+    return transform

+from .basenet import BaseNet

+import torch.nn as nn
+
+class BaseNet(nn.Module):
+    def __init__(self, backbone, args):
+        super(BaseNet, self).__init__()
+
+        # Separate layers
+        self.first = nn.Sequential(*list(backbone.children())[:1])
+        self.after = nn.Sequential(*list(backbone.children())[1:-1])
+        self.fc = list(backbone.children())[-1]
+
+        self.img_size = (224, 224)
+
+    def forward(self, x):
+        f = self.first(x)
+        x = self.after(f)
+        x = x.reshape(x.size(0), -1)
+        x = self.fc(x)
+        return x, f

+import math
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+def round_fn(orig, multiplier):
+    if not multiplier:
+        return orig
+
+    return int(math.ceil(multiplier * orig))
+
+
+def get_activation_fn(activation):
+    if activation == "swish":
+        return Swish
+
+    elif activation == "relu":
+        return nn.ReLU
+
+    else:
+        raise Exception('Unkown activation %s' % activation)
+
+
+class Swish(nn.Module):
+    """ Swish activation function, s(x) = x * sigmoid(x) """
+
+    def __init__(self, inplace=False):
+        super().__init__()
+        self.inplace = True
+
+    def forward(self, x):
+        if self.inplace:
+            x.mul_(F.sigmoid(x))
+            return x
+        else:
+            return x * F.sigmoid(x)
+
+
+class ConvBlock(nn.Module):
+    """ Conv + BatchNorm + Activation """
+
+    def __init__(self, in_channel, out_channel, kernel_size,
+                 padding=0, stride=1, activation="swish"):
+        super().__init__()
+        self.fw = nn.Sequential(
+                nn.Conv2d(in_channel, out_channel, kernel_size,
+                          padding=padding, stride=stride, bias=False),
+                nn.BatchNorm2d(out_channel),
+                get_activation_fn(activation)())
+
+    def forward(self, x):
+        return self.fw(x)
+
+
+class DepthwiseConvBlock(nn.Module):
+    """ DepthwiseConv2D + BatchNorm + Activation """
+
+    def __init__(self, in_channel, kernel_size,
+                 padding=0, stride=1, activation="swish"):
+        super().__init__()
+        self.fw = nn.Sequential(
+                nn.Conv2d(in_channel, in_channel, kernel_size,
+                          padding=padding, stride=stride, groups=in_channel, bias=False),
+                nn.BatchNorm2d(in_channel),
+                get_activation_fn(activation)())
+
+    def forward(self, x):
+        return self.fw(x)
+
+
+class MBConv(nn.Module):
+    """ Inverted residual block """
+
+    def __init__(self, in_channel, out_channel, kernel_size,
+                 stride=1, expand_ratio=1, activation="swish"):
+        super().__init__()
+        self.in_channel = in_channel
+        self.out_channel = out_channel
+        self.expand_ratio = expand_ratio
+        self.stride = stride
+
+        if expand_ratio != 1:
+            self.expand = ConvBlock(in_channel, in_channel*expand_ratio, 1,
+                                    activation=activation)
+
+        self.dw_conv = DepthwiseConvBlock(in_channel*expand_ratio, kernel_size,
+                                          padding=(kernel_size-1)//2,
+                                          stride=stride, activation=activation)
+
+        self.pw_conv = ConvBlock(in_channel*expand_ratio, out_channel, 1,
+                                 activation=activation)
+
+    def forward(self, inputs):
+        if self.expand_ratio != 1:
+            x = self.expand(inputs)
+        else:
+            x = inputs
+
+        x = self.dw_conv(x)
+        x = self.pw_conv(x)
+
+        if self.in_channel == self.out_channel and \
+                self.stride == 1:
+            x = x + inputs
+
+        return x
+
+
+class Net(nn.Module):
+    """ EfficientNet """
+
+    def __init__(self, args):
+        super(Net, self).__init__()
+        pi = args.pi
+        activation = args.activation
+        num_classes = args.num_classes
+
+        self.d = 1.2 ** pi
+        self.w = 1.1 ** pi
+        self.r = 1.15 ** pi
+        self.img_size = (round_fn(224, self.r), round_fn(224, self.r))
+
+        self.stage1 = ConvBlock(3, round_fn(32, self.w),
+                                kernel_size=3, padding=1, stride=2, activation=activation)
+
+        self.stage2 = self.make_layers(round_fn(32, self.w), round_fn(16, self.w),
+                                       depth=round_fn(1, self.d), kernel_size=3,
+                                       half_resolution=False, expand_ratio=1, activation=activation)
+
+        self.stage3 = self.make_layers(round_fn(16, self.w), round_fn(24, self.w),
+                                       depth=round_fn(2, self.d), kernel_size=3,
+                                       half_resolution=True, expand_ratio=6, activation=activation)
+
+        self.stage4 = self.make_layers(round_fn(24, self.w), round_fn(40, self.w),
+                                       depth=round_fn(2, self.d), kernel_size=5,
+                                       half_resolution=True, expand_ratio=6, activation=activation)
+
+        self.stage5 = self.make_layers(round_fn(40, self.w), round_fn(80, self.w),
+                                       depth=round_fn(3, self.d), kernel_size=3,
+                                       half_resolution=True, expand_ratio=6, activation=activation)
+
+        self.stage6 = self.make_layers(round_fn(80, self.w), round_fn(112, self.w),
+                                       depth=round_fn(3, self.d), kernel_size=5,
+                                       half_resolution=False, expand_ratio=6, activation=activation)
+
+        self.stage7 = self.make_layers(round_fn(112, self.w), round_fn(192, self.w),
+                                       depth=round_fn(4, self.d), kernel_size=5,
+                                       half_resolution=True, expand_ratio=6, activation=activation)
+
+        self.stage8 = self.make_layers(round_fn(192, self.w), round_fn(320, self.w),
+                                       depth=round_fn(1, self.d), kernel_size=3,
+                                       half_resolution=False, expand_ratio=6, activation=activation)
+
+        self.stage9 = ConvBlock(round_fn(320, self.w), round_fn(1280, self.w),
+                                kernel_size=1, activation=activation)
+
+        self.fc = nn.Linear(round_fn(7*7*1280, self.w), num_classes)
+
+    def make_layers(self, in_channel, out_channel, depth, kernel_size,
+                    half_resolution=False, expand_ratio=1, activation="swish"):
+        blocks = []
+        for i in range(depth):
+            stride = 2 if half_resolution and i==0 else 1
+            blocks.append(
+                    MBConv(in_channel, out_channel, kernel_size,
+                           stride=stride, expand_ratio=expand_ratio, activation=activation))
+            in_channel = out_channel
+
+        return nn.Sequential(*blocks)
+
+    def forward(self, x):
+        assert x.size()[-2:] == self.img_size, \
+                'Image size must be %r, but %r given' % (self.img_size, x.size()[-2])
+
+        x = self.stage1(x)
+        x = self.stage2(x)
+        x = self.stage3(x)
+        x = self.stage4(x)
+        x = self.stage5(x)
+        x = self.stage6(x)
+        x = self.stage7(x)
+        x = self.stage8(x)
+        x = self.stage9(x)
+        x = x.reshape(x.size(0), -1)
+        x = self.fc(x)
+        return x, x

+import math
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+def round_fn(orig, multiplier):
+    if not multiplier:
+        return orig
+
+    return int(math.ceil(multiplier * orig))
+
+
+def get_activation_fn(activation):
+    if activation == "swish":
+        return Swish
+
+    elif activation == "relu":
+        return nn.ReLU
+
+    else:
+        raise Exception('Unkown activation %s' % activation)
+
+
+class Swish(nn.Module):
+    """ Swish activation function, s(x) = x * sigmoid(x) """
+
+    def __init__(self, inplace=False):
+        super().__init__()
+        self.inplace = True
+
+    def forward(self, x):
+        if self.inplace:
+            x.mul_(F.sigmoid(x))
+            return x
+        else:
+            return x * F.sigmoid(x)
+
+
+class ConvBlock(nn.Module):
+    """ Conv + BatchNorm + Activation """
+
+    def __init__(self, in_channel, out_channel, kernel_size,
+                 padding=0, stride=1, activation="swish"):
+        super().__init__()
+        self.fw = nn.Sequential(
+                nn.Conv2d(in_channel, out_channel, kernel_size,
+                          padding=padding, stride=stride, bias=False),
+                nn.BatchNorm2d(out_channel),
+                get_activation_fn(activation)())
+
+    def forward(self, x):
+        return self.fw(x)
+
+
+class DepthwiseConvBlock(nn.Module):
+    """ DepthwiseConv2D + BatchNorm + Activation """
+
+    def __init__(self, in_channel, kernel_size,
+                 padding=0, stride=1, activation="swish"):
+        super().__init__()
+        self.fw = nn.Sequential(
+                nn.Conv2d(in_channel, in_channel, kernel_size,
+                          padding=padding, stride=stride, groups=in_channel, bias=False),
+                nn.BatchNorm2d(in_channel),
+                get_activation_fn(activation)())
+
+    def forward(self, x):
+        return self.fw(x)
+
+
+class SEBlock(nn.Module):
+    """ Squeeze and Excitation Block """
+
+    def __init__(self, in_channel, se_ratio=16):
+        super().__init__()
+        self.global_avgpool = nn.AdaptiveAvgPool2d((1,1))
+        inter_channel = in_channel // se_ratio
+
+        self.reduce = nn.Sequential(
+                nn.Conv2d(in_channel, inter_channel,
+                          kernel_size=1, padding=0, stride=1),
+                nn.ReLU())
+
+        self.expand = nn.Sequential(
+                nn.Conv2d(inter_channel, in_channel,
+                          kernel_size=1, padding=0, stride=1),
+                nn.Sigmoid())
+
+
+    def forward(self, x):
+        s = self.global_avgpool(x)
+        s = self.reduce(s)
+        s = self.expand(s)
+        return x * s
+
+
+class MBConv(nn.Module):
+    """ Inverted residual block """
+
+    def __init__(self, in_channel, out_channel, kernel_size,
+                 stride=1, expand_ratio=1, activation="swish", use_seblock=False):
+        super().__init__()
+        self.in_channel = in_channel
+        self.out_channel = out_channel
+        self.expand_ratio = expand_ratio
+        self.stride = stride
+        self.use_seblock = use_seblock
+
+        if expand_ratio != 1:
+            self.expand = ConvBlock(in_channel, in_channel*expand_ratio, 1,
+                                    activation=activation)
+
+        self.dw_conv = DepthwiseConvBlock(in_channel*expand_ratio, kernel_size,
+                                          padding=(kernel_size-1)//2,
+                                          stride=stride, activation=activation)
+
+        if use_seblock:
+            self.seblock = SEBlock(in_channel*expand_ratio)
+
+        self.pw_conv = ConvBlock(in_channel*expand_ratio, out_channel, 1,
+                                 activation=activation)
+
+    def forward(self, inputs):
+        if self.expand_ratio != 1:
+            x = self.expand(inputs)
+        else:
+            x = inputs
+
+        x = self.dw_conv(x)
+
+        if self.use_seblock:
+            x = self.seblock(x)
+
+        x = self.pw_conv(x)
+
+        if self.in_channel == self.out_channel and \
+                self.stride == 1:
+            x = x + inputs
+
+        return x
+
+
+class Net(nn.Module):
+    """ EfficientNet """
+
+    def __init__(self, args):
+        super(Net, self).__init__()
+        pi = args.pi
+        activation = args.activation
+        num_classes = 10
+
+        self.d = 1.2 ** pi
+        self.w = 1.1 ** pi
+        self.r = 1.15 ** pi
+        self.img_size = (round_fn(32, self.r), round_fn(32, self.r))
+        self.use_seblock = args.use_seblock
+
+        self.stage1 = ConvBlock(3, round_fn(32, self.w),
+                                kernel_size=3, padding=1, stride=2, activation=activation)
+
+        self.stage2 = self.make_layers(round_fn(32, self.w), round_fn(16, self.w),
+                                       depth=round_fn(1, self.d), kernel_size=3,
+                                       half_resolution=False, expand_ratio=1, activation=activation)
+
+        self.stage3 = self.make_layers(round_fn(16, self.w), round_fn(24, self.w),
+                                       depth=round_fn(2, self.d), kernel_size=3,
+                                       half_resolution=True, expand_ratio=6, activation=activation)
+
+        self.stage4 = self.make_layers(round_fn(24, self.w), round_fn(40, self.w),
+                                       depth=round_fn(2, self.d), kernel_size=5,
+                                       half_resolution=True, expand_ratio=6, activation=activation)
+
+        self.stage5 = self.make_layers(round_fn(40, self.w), round_fn(80, self.w),
+                                       depth=round_fn(3, self.d), kernel_size=3,
+                                       half_resolution=True, expand_ratio=6, activation=activation)
+
+        self.stage6 = self.make_layers(round_fn(80, self.w), round_fn(112, self.w),
+                                       depth=round_fn(3, self.d), kernel_size=5,
+                                       half_resolution=False, expand_ratio=6, activation=activation)
+
+        self.stage7 = self.make_layers(round_fn(112, self.w), round_fn(192, self.w),
+                                       depth=round_fn(4, self.d), kernel_size=5,
+                                       half_resolution=True, expand_ratio=6, activation=activation)
+
+        self.stage8 = self.make_layers(round_fn(192, self.w), round_fn(320, self.w),
+                                       depth=round_fn(1, self.d), kernel_size=3,
+                                       half_resolution=False, expand_ratio=6, activation=activation)
+
+        self.stage9 = ConvBlock(round_fn(320, self.w), round_fn(1280, self.w),
+                                kernel_size=1, activation=activation)
+
+        self.fc = nn.Linear(round_fn(1280, self.w), num_classes)
+
+    def make_layers(self, in_channel, out_channel, depth, kernel_size,
+                    half_resolution=False, expand_ratio=1, activation="swish"):
+        blocks = []
+        for i in range(depth):
+            stride = 2 if half_resolution and i==0 else 1
+            blocks.append(
+                    MBConv(in_channel, out_channel, kernel_size,
+                           stride=stride, expand_ratio=expand_ratio, activation=activation, use_seblock=self.use_seblock))
+            in_channel = out_channel
+
+        return nn.Sequential(*blocks)
+
+    def forward(self, x):
+        assert x.size()[-2:] == self.img_size, \
+                'Image size must be %r, but %r given' % (self.img_size, x.size()[-2])
+
+        s = self.stage1(x)
+        x = self.stage2(s)
+        x = self.stage3(x)
+        x = self.stage4(x)
+        x = self.stage5(x)
+        x = self.stage6(x)
+        x = self.stage7(x)
+        x = self.stage8(x)
+        x = self.stage9(x)
+        x = x.reshape(x.size(0), -1)
+        x = self.fc(x)
+        return x, s

+import torch.nn as nn
+
+
+class ResidualBlock(nn.Module):
+    def __init__(self, in_channel, out_channel, stride):
+        super(ResidualBlock, self).__init__()
+        self.in_channel = in_channel
+        self.out_channel = out_channel
+        self.stride = stride
+
+        self.conv1 = nn.Sequential(
+                nn.Conv2d(in_channel, out_channel,
+                          kernel_size=3, padding=1, stride=stride),
+                nn.BatchNorm2d(out_channel))
+
+        self.relu = nn.ReLU(inplace=True)
+
+        self.conv2 = nn.Sequential(
+                nn.Conv2d(out_channel, out_channel,
+                          kernel_size=3, padding=1),
+                nn.BatchNorm2d(out_channel))
+
+        if self.in_channel != self.out_channel or \
+                self.stride != 1:
+            self.down = nn.Sequential(
+                    nn.Conv2d(in_channel, out_channel,
+                              kernel_size=1, stride=stride),
+                    nn.BatchNorm2d(out_channel))
+
+    def forward(self, b):
+        t = self.conv1(b)
+        t = self.relu(t)
+        t = self.conv2(t)
+
+        if self.in_channel != self.out_channel or \
+                self.stride != 1:
+            b = self.down(b)
+
+        t += b
+        t = self.relu(t)
+
+        return t
+
+
+class Net(nn.Module):
+    def __init__(self, args):
+        super(Net, self).__init__()
+        scale = args.scale
+
+        self.stem = nn.Sequential(
+                nn.Conv2d(3, 16,
+                          kernel_size=3, padding=1),
+                nn.BatchNorm2d(16),
+                nn.ReLU(inplace=True))
+
+        self.layer1 = nn.Sequential(*[
+            ResidualBlock(16, 16, 1) for _ in range(2*scale)])
+
+        self.layer2 = nn.Sequential(*[
+            ResidualBlock(in_channel=(16 if i==0 else 32),
+                          out_channel=32,
+                          stride=(2 if i==0 else 1)) for i in range(2*scale)])
+
+        self.layer3 = nn.Sequential(*[
+            ResidualBlock(in_channel=(32 if i==0 else 64),
+                          out_channel=64,
+                          stride=(2 if i==0 else 1)) for i in range(2*scale)])
+
+        self.avg_pool = nn.AdaptiveAvgPool2d((1,1))
+
+        self.fc = nn.Linear(64, 10)
+
+        for m in self.modules():
+            if isinstance(m, nn.Conv2d):
+                nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
+
+            elif isinstance(m, (nn.BatchNorm2d, nn.GroupNorm)):
+                nn.init.constant_(m.weight, 1)
+                nn.init.constant_(m.bias, 0)
+
+    def forward(self, x):
+        s = self.stem(x)
+        x = self.layer1(s)
+        x = self.layer2(x)
+        x = self.layer3(x)
+        x = self.avg_pool(x)
+        x = x.reshape(x.size(0), -1)
+        x = self.fc(x)
+
+        return x, s

+future
+tb-nightly
+torchvision
+torch
+hyperopt

+import os
+import fire
+import time
+import json
+import random
+from pprint import pprint
+
+import torch.nn as nn
+import torch.backends.cudnn as cudnn
+from torch.utils.tensorboard import SummaryWriter
+
+from networks import *
+from utils import *
+
+
+def train(**kwargs):
+    print('\n[+] Parse arguments')
+    args, kwargs = parse_args(kwargs)
+    pprint(args)
+    device = torch.device('cuda' if args.use_cuda else 'cpu')
+
+    print('\n[+] Create log dir')
+    model_name = get_model_name(args)
+    log_dir = os.path.join('./runs', model_name)
+    os.makedirs(os.path.join(log_dir, 'model'))
+    json.dump(kwargs, open(os.path.join(log_dir, 'kwargs.json'), 'w'))
+    writer = SummaryWriter(log_dir=log_dir)
+
+    if args.seed is not None:
+        random.seed(args.seed)
+        torch.manual_seed(args.seed)
+        cudnn.deterministic = True
+
+    print('\n[+] Create network')
+    model = select_model(args)
+    optimizer = select_optimizer(args, model)
+    scheduler = select_scheduler(args, optimizer)
+    criterion = nn.CrossEntropyLoss()
+    if args.use_cuda:
+        model = model.cuda()
+        criterion = criterion.cuda()
+    #writer.add_graph(model)
+
+    print('\n[+] Load dataset')
+    transform = get_train_transform(args, model, log_dir)
+    val_transform = get_valid_transform(args, model)
+    train_dataset = get_dataset(args, transform, 'train')
+    valid_dataset = get_dataset(args, val_transform, 'val')
+    train_loader = iter(get_inf_dataloader(args, train_dataset))
+    max_epoch = len(train_dataset) // args.batch_size
+    best_acc = -1
+
+    print('\n[+] Start training')
+    if torch.cuda.device_count() > 1:
+        print('\n[+] Use {} GPUs'.format(torch.cuda.device_count()))
+        model = nn.DataParallel(model)
+
+    start_t = time.time()
+    for step in range(args.start_step, args.max_step):
+        batch = next(train_loader)
+        _train_res = train_step(args, model, optimizer, scheduler, criterion, batch, step, writer)
+
+        if step % args.print_step == 0:
+            print('\n[+] Training step: {}/{}\tTraining epoch: {}/{}\tElapsed time: {:.2f}min\tLearning rate: {}'.format(
+                step, args.max_step, current_epoch, max_epoch, (time.time()-start_t)/60, optimizer.param_groups[0]['lr']))
+            writer.add_scalar('train/learning_rate', optimizer.param_groups[0]['lr'], global_step=step)
+            writer.add_scalar('train/acc1', _train_res[0], global_step=step)
+            writer.add_scalar('train/acc5', _train_res[1], global_step=step)
+            writer.add_scalar('train/loss', _train_res[2], global_step=step)
+            writer.add_scalar('train/forward_time', _train_res[3], global_step=step)
+            writer.add_scalar('train/backward_time', _train_res[4], global_step=step)
+            print('  Acc@1 : {:.3f}%'.format(_train_res[0].data.cpu().numpy()[0]*100))
+            print('  Acc@5 : {:.3f}%'.format(_train_res[1].data.cpu().numpy()[0]*100))
+            print('  Loss : {}'.format(_train_res[2].data))
+            print('  FW Time : {:.3f}ms'.format(_train_res[3]*1000))
+            print('  BW Time : {:.3f}ms'.format(_train_res[4]*1000))
+
+        if step % args.val_step == args.val_step-1:
+            valid_loader = iter(get_dataloader(args, valid_dataset))
+            _valid_res = validate(args, model, criterion, valid_loader, step, writer)
+            print('\n[+] Valid results')
+            writer.add_scalar('valid/acc1', _valid_res[0], global_step=step)
+            writer.add_scalar('valid/acc5', _valid_res[1], global_step=step)
+            writer.add_scalar('valid/loss', _valid_res[2], global_step=step)
+            print('  Acc@1 : {:.3f}%'.format(_valid_res[0].data.cpu().numpy()[0]*100))
+            print('  Acc@5 : {:.3f}%'.format(_valid_res[1].data.cpu().numpy()[0]*100))
+            print('  Loss : {}'.format(_valid_res[2].data))
+
+            if _valid_res[0] > best_acc:
+                best_acc = _valid_res[0]
+                torch.save(model.state_dict(), os.path.join(log_dir, "model","model.pt"))
+                print('\n[+] Model saved')
+
+    writer.close()
+
+
+if __name__ == '__main__':
+    fire.Fire(train)

+import numpy as np
+import torch.nn as nn
+import torchvision.transforms as transforms
+
+from abc import ABC, abstractmethod
+from PIL import Image, ImageOps, ImageEnhance
+
+
+class BaseTransform(ABC):
+
+    def __init__(self, prob, mag):
+        self.prob = prob
+        self.mag = mag
+
+    def __call__(self, img):
+        return transforms.RandomApply([self.transform], self.prob)(img)
+
+    def __repr__(self):
+        return '%s(prob=%.2f, magnitude=%.2f)' % \
+                (self.__class__.__name__, self.prob, self.mag)
+
+    @abstractmethod
+    def transform(self, img):
+        pass
+
+
+class ShearXY(BaseTransform):
+
+    def transform(self, img):
+        degrees = self.mag * 360
+        t = transforms.RandomAffine(0, shear=degrees, resample=Image.BILINEAR)
+        return t(img)
+
+
+class TranslateXY(BaseTransform):
+
+    def transform(self, img):
+        translate = (self.mag, self.mag)
+        t = transforms.RandomAffine(0, translate=translate, resample=Image.BILINEAR)
+        return t(img)
+
+
+class Rotate(BaseTransform):
+
+    def transform(self, img):
+        degrees = self.mag * 360
+        t = transforms.RandomRotation(degrees, Image.BILINEAR)
+        return t(img)
+
+
+class AutoContrast(BaseTransform):
+
+    def transform(self, img):
+        cutoff = int(self.mag * 49)
+        return ImageOps.autocontrast(img, cutoff=cutoff)
+
+
+class Invert(BaseTransform):
+
+    def transform(self, img):
+        return ImageOps.invert(img)
+
+
+class Equalize(BaseTransform):
+
+    def transform(self, img):
+        return ImageOps.equalize(img)
+
+
+class Solarize(BaseTransform):
+
+    def transform(self, img):
+        threshold = (1-self.mag) * 255
+        return ImageOps.solarize(img, threshold)
+
+
+class Posterize(BaseTransform):
+
+    def transform(self, img):
+        bits = int((1-self.mag) * 8)
+        return ImageOps.posterize(img, bits=bits)
+
+
+class Contrast(BaseTransform):
+
+    def transform(self, img):
+        factor = self.mag * 10
+        return ImageEnhance.Contrast(img).enhance(factor)
+
+
+class Color(BaseTransform):
+
+    def transform(self, img):
+        factor = self.mag * 10
+        return ImageEnhance.Color(img).enhance(factor)
+
+
+class Brightness(BaseTransform):
+
+    def transform(self, img):
+        factor = self.mag * 10
+        return ImageEnhance.Brightness(img).enhance(factor)
+
+
+class Sharpness(BaseTransform):
+
+    def transform(self, img):
+        factor = self.mag * 10
+        return ImageEnhance.Sharpness(img).enhance(factor)
+
+
+class Cutout(BaseTransform):
+
+    def transform(self, img):
+        n_holes = 1
+        length = 24 * self.mag
+        cutout_op = CutoutOp(n_holes=n_holes, length=length)
+        return cutout_op(img)
+
+
+class CutoutOp(object):
+    """
+    https://github.com/uoguelph-mlrg/Cutout
+
+    Randomly mask out one or more patches from an image.
+
+    Args:
+        n_holes (int): Number of patches to cut out of each image.
+        length (int): The length (in pixels) of each square patch.
+    """
+    def __init__(self, n_holes, length):
+        self.n_holes = n_holes
+        self.length = length
+
+    def __call__(self, img):
+        """
+        Args:
+            img (Tensor): Tensor image of size (C, H, W).
+        Returns:
+            Tensor: Image with n_holes of dimension length x length cut out of it.
+        """
+        w, h = img.size
+
+        mask = np.ones((h, w, 1), np.uint8)
+
+        for n in range(self.n_holes):
+            y = np.random.randint(h)
+            x = np.random.randint(w)
+
+            y1 = np.clip(y - self.length // 2, 0, h).astype(int)
+            y2 = np.clip(y + self.length // 2, 0, h).astype(int)
+            x1 = np.clip(x - self.length // 2, 0, w).astype(int)
+            x2 = np.clip(x + self.length // 2, 0, w).astype(int)
+
+            mask[y1: y2, x1: x2, :] = 0.
+
+        img = mask*np.asarray(img).astype(np.uint8)
+        img = Image.fromarray(mask*np.asarray(img))
+
+        return img
+