rm ori FAA2

-# 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
-