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
+

+import os
+import time
+import importlib
+import collections
+import pickle as cp
+import glob
+import numpy as np
+
+import torch
+import torchvision
+import torch.nn.functional as F
+import torchvision.models as models
+import torchvision.transforms as transforms
+from torch.utils.data import Subset
+from torch.utils.data import Dataset, DataLoader
+
+from sklearn.model_selection import StratifiedShuffleSplit
+
+
+TRAIN_DATASET_PATH = '/content/drive/My Drive/CD2 Project/data/BraTS_Training/train_frame/'
+VAL_DATASET_PATH = '/content/drive/My Drive/CD2 Project/data/BraTS_Training/val_frame/'
+current_epoch = 0
+
+
+def split_dataset(args, dataset, k):
+    # load dataset
+    X = list(range(len(dataset)))
+    Y = dataset.targets
+
+    # split to k-fold
+    assert len(X) == len(Y)
+
+    def _it_to_list(_it):
+        return list(zip(*list(_it)))
+
+    sss = StratifiedShuffleSplit(n_splits=k, random_state=args.seed, test_size=0.1)
+    Dm_indexes, Da_indexes = _it_to_list(sss.split(X, Y))
+
+    return Dm_indexes, Da_indexes
+
+
+def concat_image_features(image, features, max_features=3):
+    _, h, w = image.shape
+
+    max_features = min(features.size(0), max_features)
+    image_feature = image.clone()
+
+    for i in range(max_features):
+        feature = features[i:i+1]
+        _min, _max = torch.min(feature), torch.max(feature)
+        feature = (feature - _min) / (_max - _min + 1e-6)
+        feature = torch.cat([feature]*3, 0)
+        feature = feature.view(1, 3, feature.size(1), feature.size(2))
+        feature = F.upsample(feature, size=(h,w), mode="bilinear")
+        feature = feature.view(3, h, w)
+        image_feature = torch.cat((image_feature, feature), 2)
+
+    return image_feature
+
+
+def get_model_name(args):
+    from datetime import datetime
+    now = datetime.now()
+    date_time = now.strftime("%B_%d_%H:%M:%S")
+    model_name = '__'.join([date_time, args.network, str(args.seed)])
+    return model_name
+
+
+def dict_to_namedtuple(d):
+    Args = collections.namedtuple('Args', sorted(d.keys()))
+
+    for k,v in d.items():
+        if type(v) is dict:
+            d[k] = dict_to_namedtuple(v)
+
+        elif type(v) is str:
+            try:
+                d[k] = eval(v)
+            except:
+                d[k] = v
+
+    args = Args(**d)
+    return args
+
+
+def parse_args(kwargs):
+    # combine with default args
+    kwargs['dataset'] =  kwargs['dataset'] if 'dataset' in kwargs else 'cifar10'
+    kwargs['network'] =  kwargs['network'] if 'network' in kwargs else 'resnet_cifar10'
+    kwargs['optimizer'] = kwargs['optimizer'] if 'optimizer' in kwargs else 'adam'
+    kwargs['learning_rate'] = kwargs['learning_rate'] if 'learning_rate' in kwargs else 0.1
+    kwargs['seed'] =  kwargs['seed'] if 'seed' in kwargs else None
+    kwargs['use_cuda'] =  kwargs['use_cuda'] if 'use_cuda' in kwargs else True
+    kwargs['use_cuda'] =  kwargs['use_cuda'] and torch.cuda.is_available()
+    kwargs['num_workers'] = kwargs['num_workers'] if 'num_workers' in kwargs else 4
+    kwargs['print_step'] = kwargs['print_step'] if 'print_step' in kwargs else 2000
+    kwargs['val_step'] = kwargs['val_step'] if 'val_step' in kwargs else 2000
+    kwargs['scheduler'] = kwargs['scheduler'] if 'scheduler' in kwargs else 'exp'
+    kwargs['batch_size'] = kwargs['batch_size'] if 'batch_size' in kwargs else 128
+    kwargs['start_step'] = kwargs['start_step'] if 'start_step' in kwargs else 0
+    kwargs['max_step'] = kwargs['max_step'] if 'max_step' in kwargs else 64000
+    kwargs['fast_auto_augment'] = kwargs['fast_auto_augment'] if 'fast_auto_augment' in kwargs else False
+    kwargs['augment_path'] = kwargs['augment_path'] if 'augment_path' in kwargs else None
+
+    # to named tuple
+    args = dict_to_namedtuple(kwargs)
+    return args, kwargs
+
+
+def select_model(args):
+    if args.network in models.__dict__:
+        backbone = models.__dict__[args.network]()
+        model = BaseNet(backbone, args)
+    else:
+        Net = getattr(importlib.import_module('networks.{}'.format(args.network)), 'Net')
+        model = Net(args)
+
+    print(model)
+    return model
+
+
+def select_optimizer(args, model):
+    if args.optimizer == 'sgd':
+        optimizer = torch.optim.SGD(model.parameters(), lr=args.learning_rate, momentum=0.9, weight_decay=0.0001)
+    elif args.optimizer == 'rms':
+        #optimizer = torch.optim.RMSprop(model.parameters(), lr=args.learning_rate, momentum=0.9, weight_decay=1e-5)
+        optimizer = torch.optim.RMSprop(model.parameters(), lr=args.learning_rate)
+    elif args.optimizer == 'adam':
+        optimizer = torch.optim.Adam(model.parameters(), lr=args.learning_rate)
+    else:
+        raise Exception('Unknown Optimizer')
+    return optimizer
+
+
+def select_scheduler(args, optimizer):
+    if not args.scheduler or args.scheduler == 'None':
+        return None
+    elif args.scheduler =='clr':
+        return torch.optim.lr_scheduler.CyclicLR(optimizer, 0.01, 0.015, mode='triangular2', step_size_up=250000, cycle_momentum=False)
+    elif args.scheduler =='exp':
+        return torch.optim.lr_scheduler.ExponentialLR(optimizer, gamma=0.9999283, last_epoch=-1)
+    else:
+        raise Exception('Unknown Scheduler')
+
+
+class CustomDataset(Dataset): 
+  def __init__(self, path, transform = None):
+      self.path = path
+      self.transform = transform
+      self.img = np.load(path)
+      self.len = self.img.shape[0] 
+
+  def __len__(self):
+      return self.len
+
+  def __getitem__(self, idx): 
+       if self.transforms is not None:
+            img = self.transforms(img)
+        return img
+
+def get_dataset(args, transform, split='train'):
+    assert split in ['train', 'val', 'test', 'trainval']
+
+    if args.dataset == 'cifar10':
+        train = split in ['train', 'val', 'trainval']
+        dataset = torchvision.datasets.CIFAR10(DATASET_PATH,
+                                               train=train,
+                                               transform=transform,
+                                               download=True)
+
+        if split in ['train', 'val']:
+            split_path = os.path.join(DATASET_PATH,
+                    'cifar-10-batches-py', 'train_val_index.cp')
+
+            if not os.path.exists(split_path):
+                [train_index], [val_index] = split_dataset(args, dataset, k=1)
+                split_index = {'train':train_index, 'val':val_index}
+                cp.dump(split_index, open(split_path, 'wb'))
+
+            split_index = cp.load(open(split_path, 'rb'))
+            dataset = Subset(dataset, split_index[split])
+
+    elif args.dataset == 'imagenet':
+        dataset = torchvision.datasets.ImageNet(DATASET_PATH,
+                                                split=split,
+                                                transform=transform,
+                                                download=(split is 'val'))
+
+    elif args.dataset == 'BraTS':
+      if split in ['train']:
+        dataset = CustomDataset(TRAIN_DATASET_PATH, transform=transform)
+      else:
+        dataset = CustomDataset(VAL_DATASET_PATH, transform=transform)
+      
+
+    else:
+        raise Exception('Unknown dataset')
+
+    return dataset
+
+
+def get_dataloader(args, dataset, shuffle=False, pin_memory=True):
+    data_loader = torch.utils.data.DataLoader(dataset,
+                                              batch_size=args.batch_size,
+                                              shuffle=shuffle,
+                                              num_workers=args.num_workers,
+                                              pin_memory=pin_memory)
+    return data_loader
+
+
+def get_inf_dataloader(args, dataset):
+    global current_epoch
+    data_loader = iter(get_dataloader(args, dataset, shuffle=True))
+
+    while True:
+        try:
+            batch = next(data_loader)
+
+        except StopIteration:
+            current_epoch += 1
+            data_loader = iter(get_dataloader(args, dataset, shuffle=True))
+            batch = next(data_loader)
+
+        yield batch
+
+
+def get_train_transform(args, model, log_dir=None):
+    if args.fast_auto_augment:
+        assert args.dataset == 'BraTS' # TODO: FastAutoAugment for Imagenet
+
+        from fast_auto_augment import fast_auto_augment
+        if args.augment_path:
+            transform = cp.load(open(args.augment_path, 'rb'))
+            os.system('cp {} {}'.format(
+                args.augment_path, os.path.join(log_dir, 'augmentation.cp')))
+        else:
+            transform = fast_auto_augment(args, model, K=4, B=1, num_process=4)
+            if log_dir:
+                cp.dump(transform, open(os.path.join(log_dir, 'augmentation.cp'), 'wb'))
+
+    elif args.dataset == 'cifar10':
+        transform = transforms.Compose([
+            transforms.Pad(4),
+            transforms.RandomCrop(32),
+            transforms.RandomHorizontalFlip(),
+            transforms.ToTensor()
+        ])
+
+    elif args.dataset == 'imagenet':
+        resize_h, resize_w = model.img_size[0], int(model.img_size[1]*1.875)
+        transform = transforms.Compose([
+            transforms.Resize([resize_h, resize_w]),
+            transforms.RandomCrop(model.img_size),
+            transforms.RandomHorizontalFlip(),
+            transforms.ToTensor()
+        ])
+
+   elif args.dataset == 'BraTS':
+        resize_h, resize_w = 256, 256
+        transform = transforms.Compose([
+            transforms.Resize([resize_h, resize_w]),
+            transforms.RandomCrop(model.img_size),
+            transforms.RandomHorizontalFlip(),
+            transforms.ToTensor()
+        ])
+    else:
+        raise Exception('Unknown Dataset')
+
+    print(transform)
+
+    return transform
+
+
+def get_valid_transform(args, model):
+    if args.dataset == 'cifar10':
+        val_transform = transforms.Compose([
+            transforms.Resize(32),
+            transforms.ToTensor()
+        ])
+
+    elif args.dataset == 'imagenet':
+        resize_h, resize_w = model.img_size[0], int(model.img_size[1]*1.875)
+        val_transform = transforms.Compose([
+            transforms.Resize([resize_h, resize_w]),
+            transforms.ToTensor()
+        ])
+   elif args.dataset == 'BraTS':
+        resize_h, resize_w = 256, 256
+        val_transform = transforms.Compose([
+            transforms.Resize([resize_h, resize_w]),
+            transforms.ToTensor()
+        ])
+    else:
+        raise Exception('Unknown Dataset')
+
+    return val_transform
+
+
+def train_step(args, model, optimizer, scheduler, criterion, batch, step, writer, device=None):
+    model.train()
+    images, target = batch
+
+    if device:
+        images = images.to(device)
+        target = target.to(device)
+
+    elif args.use_cuda:
+        images = images.cuda(non_blocking=True)
+        target = target.cuda(non_blocking=True)
+
+    # compute output
+    start_t = time.time()
+    output, first = model(images)
+    forward_t = time.time() - start_t
+    loss = criterion(output, target)
+
+    # measure accuracy and record loss
+    acc1, acc5 = accuracy(output, target, topk=(1, 5))
+    acc1 /= images.size(0)
+    acc5 /= images.size(0)
+
+    # compute gradient and do SGD step
+    optimizer.zero_grad()
+    start_t = time.time()
+    loss.backward()
+    backward_t = time.time() - start_t
+    optimizer.step()
+    if scheduler: scheduler.step()
+
+    if writer and step % args.print_step == 0:
+        n_imgs = min(images.size(0), 10)
+        for j in range(n_imgs):
+            writer.add_image('train/input_image',
+                    concat_image_features(images[j], first[j]), global_step=step)
+
+    return acc1, acc5, loss, forward_t, backward_t
+
+
+def validate(args, model, criterion, valid_loader, step, writer, device=None):
+    # switch to evaluate mode
+    model.eval()
+
+    acc1, acc5 = 0, 0
+    samples = 0
+    infer_t = 0
+
+    with torch.no_grad():
+        for i, (images, target) in enumerate(valid_loader):
+
+            start_t = time.time()
+            if device:
+                images = images.to(device)
+                target = target.to(device)
+
+            elif args.use_cuda is not None:
+                images = images.cuda(non_blocking=True)
+                target = target.cuda(non_blocking=True)
+
+            # compute output
+            output, first = model(images)
+            loss = criterion(output, target)
+            infer_t += time.time() - start_t
+
+            # measure accuracy and record loss
+            _acc1, _acc5 = accuracy(output, target, topk=(1, 5))
+            acc1 += _acc1
+            acc5 += _acc5
+            samples += images.size(0)
+
+    acc1 /= samples
+    acc5 /= samples
+
+    if writer:
+        n_imgs = min(images.size(0), 10)
+        for j in range(n_imgs):
+            writer.add_image('valid/input_image',
+                    concat_image_features(images[j], first[j]), global_step=step)
+
+    return acc1, acc5, loss, infer_t
+
+
+def accuracy(output, target, topk=(1,)):
+    """Computes the accuracy over the k top predictions for the specified values of k"""
+    with torch.no_grad():
+        maxk = max(topk)
+        batch_size = target.size(0)
+
+        _, pred = output.topk(maxk, 1, True, True)
+        pred = pred.t()
+        correct = pred.eq(target.view(1, -1).expand_as(pred))
+
+        res = []
+        for k in topk:
+            correct_k = correct[:k].view(-1).float().sum(0, keepdim=True)
+            res.append(correct_k)
+        return res
+