brainscore-vision 2.1__py3-none-any.whl

brainscore_vision/models/voneresnet_50_non_stochastic/vonenet/README.md

@@ -0,0 +1,105 @@
+
+# VOneNet: CNNs with a Primary Visual Cortex Front-End
+
+A family of biologically-inspired Convolutional Neural Networks (CNNs). VOneNets have the following features:
+- Fixed-weight neural network model of the primate primary visual cortex (V1) as the front-end.
+- Robust to image perturbations
+- Brain-mapped
+- Flexible: can be adapted to different back-end architectures
+
+[read more...](#longer-motivation)
+
+## Available Models
+*(Click on model names to download the weights of ImageNet-trained models. Alternatively, you can use the function get_model in the vonenet package to download the weights.)*
+
+| Name     | Description                                                              |
+| -------- | ------------------------------------------------------------------------ |
+| [VOneResNet50](https://vonenet-models.s3.us-east-2.amazonaws.com/voneresnet50_e70.pth.tar) | Our best performing VOneNet with a ResNet50 back-end |
+| [VOneCORnet-S](https://vonenet-models.s3.us-east-2.amazonaws.com/vonecornets_e70.pth.tar) | VOneNet with a recurrent neural network back-end based on the CORnet-S |
+| [VOneAlexNet](https://vonenet-models.s3.us-east-2.amazonaws.com/vonealexnet_e70.pth.tar) | VOneNet with a back-end based on AlexNet         |
+
+
+## Quick Start
+
+VOneNets was trained with images normalized with mean=[0.5,0.5,0.5] and std=[0.5,0.5,0.5]
+
+More information coming soon...
+
+
+## Longer Motivation
+
+Current state-of-the-art object recognition models are largely based on convolutional neural network (CNN) architectures, which are loosely inspired by the primate visual system. However, these CNNs can be fooled by imperceptibly small, explicitly crafted perturbations, and struggle to recognize objects in corrupted images that are easily recognized by humans. Recently, we observed that CNN models with a neural hidden layer that better matches primate primary visual cortex (V1) are also more robust to adversarial attacks. Inspired by this observation, we developed VOneNets, a new class of hybrid CNN vision models. Each VOneNet contains a fixed weight neural network front-end that simulates primate V1, called the VOneBlock, followed by a neural network back-end adapted from current CNN vision models. The VOneBlock is based on a classical neuroscientific model of V1: the linear-nonlinear-Poisson model, consisting of a biologically-constrained Gabor filter bank, simple and complex cell nonlinearities, and a V1 neuronal stochasticity generator. After training, VOneNets retain high ImageNet performance, but each is substantially more robust, outperforming the base CNNs and state-of-the-art methods by 18% and 3%, respectively, on a conglomerate benchmark of perturbations comprised of white box adversarial attacks and common image corruptions. Additionally, all components of the VOneBlock work in synergy to improve robustness. 
+Read more: [Dapello\*, Marques\*, et al. (biorxiv, 2020)](https://doi.org/10.1101/2020.06.16.154542)
+
+
+
+## Requirements
+
+- Python 3.6+
+- PyTorch 0.4.1+
+- numpy
+- pandas
+- tqdm
+- scipy
+
+
+## Citation
+
+Dapello, J., Marques, T., Schrimpf, M., Geiger, F., Cox, D.D., DiCarlo, J.J. (2020) Simulating a Primary Visual Cortex at the Front of CNNs Improves Robustness to Image Perturbations. *biorxiv.* doi.org/10.1101/2020.06.16.154542
+
+
+## License
+
+GNU GPL 3+
+
+
+## FAQ
+
+Soon...
+
+## Setup and Run
+
+1. You need to clone it in your local repository
+  $ git clone https://github.com/dicarlolab/vonenet.git
+   
+2. And when you setup its codes, you must need 'val' directory. so here is link.
+  this link is from Korean's blog I refered as below https://seongkyun.github.io/others/2019/03/06/imagenet_dn/
+  
+       ** Download link**
+  https://academictorrents.com/collection/imagenet-2012 
+  
+  Once you download that large tar files, you must unzip that files
+  -- all instructions below are refered above link, I only translate it
+  
+  # Unzip training dataset
+  $ mkdir train && mb ILSVRC2012_img_train.tar train/ && cd train
+  $ tar -xvf ILSVRC2012_img_train.tar
+  $ rm -f ILSVRC2012_img_train.tar    (If you want to remove zipped file(tar))
+  $ find . -name "*.tar" | while read NAME ; do mkdir -p "${NAME%.tar}"; tar -xvf "${NAME}" -C "${NAME%.tar}"; rm -f "${NAME}"; done
+  $ cd ..
+  
+  # Unzip validation dataset
+  $ mkdir val && mv ILSVRC2012_img_val.tar val/ && cd val && tar -xvf ILSVRC2012_img_val.tar
+  $ wget -qO- https://raw.githubusercontent.com/soumith/imagenetloader.torch/master/valprep.sh | bash
+  
+  when it's finished, you can see train directory, val directory
+  that 'val' directory is needed when setting up 
+  
+  ## Caution!!!!
+  after all execution above, must remove directory or file not having name n0000
+  -> there will be fault in training
+  -> ex) 'ILSVRC2012_img_train' in train directory, 'ILSVRC2012_img_val.tar' in val directory
+  
+3. if you've done getting data, then we can setting up
+  go to local repository which into you cloned and open terminal (you must check your versions of python, pytorch, cudatoolkit if okay then,)
+  $ python3 setup.py install
+  $ python3 run.py --in_path {directory including above dataset, 'val' directory must be in!} 
+  
+  # If you see any GPU related problem especially 'GPU is not available' although you already got
+  $ python3 run.py --in_path {directory including above dataset, 'val' directory must be in!} --ngpus 0
+  
+  ngpus is 1 as default. if you don't care running on CPU you do so  
+  
+  
+  
+

brainscore_vision/models/voneresnet_50_non_stochastic/vonenet/run.py

@@ -0,0 +1,136 @@
+
+import os, argparse, time, subprocess, io, shlex
+import pandas as pd
+import tqdm
+
+parser = argparse.ArgumentParser(description='ImageNet Validation')
+
+parser.add_argument('--in_path', required=True,
+                    help='path to ImageNet folder that contains val folder')
+parser.add_argument('--batch_size', default=128, type=int,
+                    help='size of batch for validation')
+parser.add_argument('--workers', default=20,
+                    help='number of data loading workers')
+parser.add_argument('--ngpus', default=1, type=int,
+                    help='number of GPUs to use; 0 if you want to run on CPU')
+parser.add_argument('--model_arch', choices=['alexnet', 'resnet50', 'resnet50_at', 'cornets'], default='resnet50',
+                    help='back-end model architecture to load')
+
+FLAGS, FIRE_FLAGS = parser.parse_known_args()
+
+
+def set_gpus(n=2):
+    """
+    Finds all GPUs on the system and restricts to n of them that have the most
+    free memory.
+    """
+    if n > 0:
+        gpus = subprocess.run(shlex.split(
+            'nvidia-smi --query-gpu=index,memory.free,memory.total --format=csv,nounits'), check=True,
+            stdout=subprocess.PIPE).stdout
+        gpus = pd.read_csv(io.BytesIO(gpus), sep=', ', engine='python')
+        gpus = gpus[gpus['memory.total [MiB]'] > 10000]  # only above 10 GB
+        if os.environ.get('CUDA_VISIBLE_DEVICES') is not None:
+            visible = [int(i)
+                       for i in os.environ['CUDA_VISIBLE_DEVICES'].split(',')]
+            gpus = gpus[gpus['index'].isin(visible)]
+        gpus = gpus.sort_values(by='memory.free [MiB]', ascending=False)
+        os.environ['CUDA_DEVICE_ORDER'] = 'PCI_BUS_ID'  # making sure GPUs are numbered the same way as in nvidia_smi
+        os.environ['CUDA_VISIBLE_DEVICES'] = ','.join(
+            [str(i) for i in gpus['index'].iloc[:n]])
+    else:
+        os.environ['CUDA_VISIBLE_DEVICES'] = '-1'
+
+
+set_gpus(FLAGS.ngpus)
+
+import torch
+import torch.nn as nn
+import torchvision
+from vonenet import get_model
+
+device = torch.device("cuda" if FLAGS.ngpus > 0 else "cpu")
+
+
+def val():
+    model = get_model(model_arch=FLAGS.model_arch, pretrained=True)
+
+    if FLAGS.ngpus == 0:
+        print('Running on CPU')
+    if FLAGS.ngpus > 0 and torch.cuda.device_count() > 1:
+        print('Running on multiple GPUs')
+        model = model.to(device)
+    elif FLAGS.ngpus > 0 and torch.cuda.device_count() is 1:
+        print('Running on single GPU')
+        model = model.to(device)
+    else:
+        print('No GPU detected!')
+        model = model.module
+
+    validator = ImageNetVal(model)
+    record = validator()
+
+    print(record['top1'])
+    print(record['top5'])
+    return
+
+
+class ImageNetVal(object):
+
+    def __init__(self, model):
+        self.name = 'val'
+        self.model = model
+        self.data_loader = self.data()
+        self.loss = nn.CrossEntropyLoss(size_average=False)
+        self.loss = self.loss.to(device)
+
+    def data(self):
+        dataset = torchvision.datasets.ImageFolder(
+            os.path.join(FLAGS.in_path, 'val'),
+            torchvision.transforms.Compose([
+                torchvision.transforms.Resize(256),
+                torchvision.transforms.CenterCrop(224),
+                torchvision.transforms.ToTensor(),
+                torchvision.transforms.Normalize(mean=[0.5, 0.5, 0.5],
+                                                 std=[0.5, 0.5, 0.5]),
+            ]))
+        data_loader = torch.utils.data.DataLoader(dataset,
+                                                  batch_size=FLAGS.batch_size,
+                                                  shuffle=False,
+                                                  num_workers=FLAGS.workers,
+                                                  pin_memory=True)
+
+        return data_loader
+
+    def __call__(self):
+        self.model.eval()
+        start = time.time()
+        record = {'loss': 0, 'top1': 0, 'top5': 0}
+        with torch.no_grad():
+            for (inp, target) in tqdm.tqdm(self.data_loader, desc=self.name):
+                target = target.to(device)
+                output = self.model(inp)
+
+                record['loss'] += self.loss(output, target).item()
+                p1, p5 = accuracy(output, target, topk=(1, 5))
+                record['top1'] += p1
+                record['top5'] += p5
+
+        for key in record:
+            record[key] /= len(self.data_loader.dataset.samples)
+        record['dur'] = (time.time() - start) / len(self.data_loader)
+
+        return record
+
+
+def accuracy(output, target, topk=(1,)):
+    with torch.no_grad():
+        _, pred = output.topk(max(topk), dim=1, largest=True, sorted=True)
+        pred = pred.t()
+        correct = pred.eq(target.view(1, -1).expand_as(pred))
+        res = [correct[:k].sum().item() for k in topk]
+        return res
+
+
+if __name__ == '__main__':
+    val()

brainscore_vision/models/voneresnet_50_non_stochastic/vonenet/setup.py

@@ -0,0 +1,41 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+
+from setuptools import setup
+
+with open('README.md') as readme_file:
+    readme = readme_file.read()
+
+requirements = [
+    "torch>=0.4.0+",
+    "torchvision",
+    "numpy",
+    "pandas",
+    "scipy",
+    "tqdm",
+    "fire",
+    "requests",
+]
+
+setup(
+    name='vonenet',
+    version='0.1.0',
+    description="CNNs with a Primary Visual Cortex Front-End ",
+    long_description=readme,
+    author="Tiago Marques, Joel Dapello",
+    author_email='tmarques@mit.edu, dapello@mit.edu',
+    url='https://github.com/dicarlolab/vonenet',
+    packages=['vonenet'],
+    include_package_data=True,
+    install_requires=requirements,
+    license="GNU GPL v3",
+    zip_safe=False,
+    keywords='VOneNet, Robustness, Primary Visual Cortex',
+    classifiers=[
+        'Development Status :: 5 - Production/Stable',
+        'Intended Audience :: Developers',
+        'License :: OSI Approved :: GNU GPL v3',
+        'Natural Language :: English',
+        'Programming Language :: Python :: 3.6'
+    ],
+)

brainscore_vision/models/voneresnet_50_non_stochastic/vonenet/train.py

@@ -0,0 +1,383 @@
+
+import os, argparse, time, subprocess, io, shlex, pickle, pprint
+import pandas as pd
+import numpy as np
+import tqdm
+import fire
+
+parser = argparse.ArgumentParser(description='ImageNet Training')
+## General parameters
+parser.add_argument('--in_path', required=True,
+                    help='path to ImageNet folder that contains train and val folders')
+parser.add_argument('-o', '--output_path', default=None,
+                    help='path for storing ')
+parser.add_argument('-restore_epoch', '--restore_epoch', default=0, type=int,
+                    help='epoch number for restoring model training ')
+parser.add_argument('-restore_path', '--restore_path', default=None, type=str,
+                    help='path of folder containing specific epoch file for restoring model training')
+
+## Training parameters
+parser.add_argument('--ngpus', default=0, type=int,
+                    help='number of GPUs to use; 0 if you want to run on CPU')
+parser.add_argument('-j', '--workers', default=20, type=int,
+                    help='number of data loading workers')
+parser.add_argument('--epochs', default=70, type=int,
+                    help='number of total epochs to run')
+parser.add_argument('--batch_size', default=256, type=int,
+                    help='mini-batch size')
+parser.add_argument('--optimizer', choices=['stepLR', 'plateauLR'], default='stepLR',
+                    help='Optimizer')
+parser.add_argument('--lr', '--learning_rate', default=.1, type=float,
+                    help='initial learning rate')
+parser.add_argument('--step_size', default=20, type=int,
+                    help='after how many epochs learning rate should be decreased by step_factor')
+parser.add_argument('--step_factor', default=0.1, type=float,
+                    help='factor by which to decrease the learning rate')
+parser.add_argument('--momentum', default=.9, type=float, help='momentum')
+parser.add_argument('--weight_decay', default=1e-4, type=float,
+                    help='weight decay ')
+
+## Model parameters
+parser.add_argument('--torch_seed', default=0, type=int,
+                    help='seed for weights initializations and torch RNG')
+parser.add_argument('--model_arch', choices=['alexnet', 'resnet50', 'resnet50_at', 'cornets'], default='resnet50',
+                    help='back-end model architecture to load')
+parser.add_argument('--normalization', choices=['vonenet', 'imagenet'], default='vonenet',
+                    help='image normalization to apply to models')
+parser.add_argument('--visual_degrees', default=8, type=float,
+                    help='Field-of-View of the model in visual degrees')
+
+## VOneBlock parameters
+# Gabor filter bank
+parser.add_argument('--stride', default=4, type=int,
+                    help='stride for the first convolution (Gabor Filter Bank)')
+parser.add_argument('--ksize', default=25, type=int,
+                    help='kernel size for the first convolution (Gabor Filter Bank)')
+parser.add_argument('--simple_channels', default=256, type=int,
+                    help='number of simple channels in V1 block')
+parser.add_argument('--complex_channels', default=256, type=int,
+                    help='number of complex channels in V1 block')
+parser.add_argument('--gabor_seed', default=0, type=int,
+                    help='seed for gabor initialization')
+parser.add_argument('--sf_corr', default=0.75, type=float,
+                    help='')
+parser.add_argument('--sf_max', default=6, type=float,
+                    help='')
+parser.add_argument('--sf_min', default=0, type=float,
+                    help='')
+parser.add_argument('--rand_param', choices=[True, False], default=False, type=bool,
+                    help='random gabor params')
+parser.add_argument('--k_exc', default=25, type=float,
+                    help='')
+
+# Noise layer
+parser.add_argument('--noise_mode', choices=['gaussian', 'neuronal', None],
+                    default=None,
+                    help='noise distribution')
+parser.add_argument('--noise_scale', default=1, type=float,
+                    help='noise scale factor')
+parser.add_argument('--noise_level', default=1, type=float,
+                    help='noise level')
+
+
+FLAGS, FIRE_FLAGS = parser.parse_known_args()
+
+
+def set_gpus(n=2):
+    """
+    Finds all GPUs on the system and restricts to n of them that have the most
+    free memory.
+    """
+    if n > 0:
+        gpus = subprocess.run(shlex.split(
+            'nvidia-smi --query-gpu=index,memory.free,memory.total --format=csv,nounits'), check=True,
+            stdout=subprocess.PIPE).stdout
+        gpus = pd.read_csv(io.BytesIO(gpus), sep=', ', engine='python')
+        gpus = gpus[gpus['memory.total [MiB]'] > 10000]  # only above 10 GB
+        if os.environ.get('CUDA_VISIBLE_DEVICES') is not None:
+            visible = [int(i)
+                       for i in os.environ['CUDA_VISIBLE_DEVICES'].split(',')]
+            gpus = gpus[gpus['index'].isin(visible)]
+        gpus = gpus.sort_values(by='memory.free [MiB]', ascending=False)
+        os.environ['CUDA_DEVICE_ORDER'] = 'PCI_BUS_ID'  # making sure GPUs are numbered the same way as in nvidia_smi
+        os.environ['CUDA_VISIBLE_DEVICES'] = ','.join(
+            [str(i) for i in gpus['index'].iloc[:n]])
+    else:
+        os.environ['CUDA_VISIBLE_DEVICES'] = '-1'
+
+
+if FLAGS.ngpus > 0:
+    set_gpus(FLAGS.ngpus)
+
+import torch
+import torch.nn as nn
+import torch.utils.model_zoo
+import torchvision
+from vonenet import get_model
+
+torch.manual_seed(FLAGS.torch_seed)
+
+torch.backends.cudnn.benchmark = True
+
+if FLAGS.ngpus > 0:
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+else:
+    device = 'cpu'
+
+if FLAGS.normalization == 'vonenet':
+    print('VOneNet normalization')
+    norm_mean = [0.5, 0.5, 0.5]
+    norm_std = [0.5, 0.5, 0.5]
+elif FLAGS.normalization == 'imagenet':
+    print('Imagenet standard normalization')
+    norm_mean = [0.485, 0.456, 0.406]
+    norm_std = [0.229, 0.224, 0.225]
+
+
+def load_model():
+    map_location = None if FLAGS.ngpus > 0 else 'cpu'
+    print('Getting VOneNet')
+    model = get_model(map_location=map_location, model_arch=FLAGS.model_arch, pretrained=False,
+                      visual_degrees=FLAGS.visual_degrees, stride=FLAGS.stride, ksize=FLAGS.ksize,
+                      sf_corr=FLAGS.sf_corr, sf_max=FLAGS.sf_max, sf_min=FLAGS.sf_min, rand_param=FLAGS.rand_param,
+                      gabor_seed=FLAGS.gabor_seed, simple_channels=FLAGS.simple_channels,
+                      complex_channels=FLAGS.simple_channels, noise_mode=FLAGS.noise_mode,
+                      noise_scale=FLAGS.noise_scale, noise_level=FLAGS.noise_level, k_exc=FLAGS.k_exc)
+
+    if FLAGS.ngpus > 0 and torch.cuda.device_count() > 1:
+        print('We have multiple GPUs detected')
+        model = model.to(device)
+    elif FLAGS.ngpus > 0 and torch.cuda.device_count() is 1:
+        print('We run on GPU')
+        model = model.to(device)
+    else:
+        print('No GPU detected!')
+        model = model.module
+
+    return model
+
+
+def train(save_train_epochs=.2,  # how often save output during training
+          save_val_epochs=.5,  # how often save output during validation
+          save_model_epochs=1,  # how often save model weights
+          save_model_secs=720 * 10  # how often save model (in sec)
+          ):
+
+    model = load_model()
+
+    trainer = ImageNetTrain(model)
+    validator = ImageNetVal(model)
+
+    start_epoch = 0
+    records = []
+
+    if FLAGS.restore_epoch > 0:
+        print('Restoring from previous...')
+        ckpt_data = torch.load(os.path.join(FLAGS.restore_path, f'epoch_{FLAGS.restore_epoch:02d}.pth.tar'))
+        start_epoch = ckpt_data['epoch']
+        print('Loaded epoch: '+str(start_epoch))
+        model.load_state_dict(ckpt_data['state_dict'])
+        trainer.optimizer.load_state_dict(ckpt_data['optimizer'])
+        results_old = pickle.load(open(os.path.join(FLAGS.restore_path, 'results.pkl'), 'rb'))
+        for result in results_old:
+            records.append(result)
+
+    results = {'meta': {'step_in_epoch': 0,
+                        'epoch': start_epoch,
+                        'wall_time': time.time()}
+               }
+
+    # records = []
+    recent_time = time.time()
+
+    nsteps = len(trainer.data_loader)
+
+    if save_train_epochs is not None:
+        save_train_steps = (np.arange(0, FLAGS.epochs + 1,
+                                      save_train_epochs) * nsteps).astype(int)
+    if save_val_epochs is not None:
+        save_val_steps = (np.arange(0, FLAGS.epochs + 1,
+                                    save_val_epochs) * nsteps).astype(int)
+    if save_model_epochs is not None:
+        save_model_steps = (np.arange(0, FLAGS.epochs + 1,
+                                      save_model_epochs) * nsteps).astype(int)
+
+    for epoch in tqdm.trange(start_epoch, FLAGS.epochs + 1, initial=0, desc='epoch'):
+        print(epoch)
+        data_load_start = np.nan
+
+        data_loader_iter = trainer.data_loader
+
+        for step, data in enumerate(tqdm.tqdm(data_loader_iter, desc=trainer.name)):
+            data_load_time = time.time() - data_load_start
+            global_step = epoch * nsteps + step
+
+            if save_val_steps is not None:
+                if global_step in save_val_steps:
+                    results[validator.name] = validator()
+                    if FLAGS.optimizer == 'plateauLR' and step == 0:
+                        trainer.lr.step(results[validator.name]['loss'])
+                    trainer.model.train()
+                    print('LR: ', trainer.optimizer.param_groups[0]["lr"])
+
+            if FLAGS.output_path is not None:
+                if not (os.path.isdir(FLAGS.output_path)):
+                    os.mkdir(FLAGS.output_path)
+
+                records.append(results)
+                if len(results) > 1:
+                    pickle.dump(records, open(os.path.join(FLAGS.output_path, 'results.pkl'), 'wb'))
+
+                ckpt_data = {}
+                ckpt_data['flags'] = FLAGS.__dict__.copy()
+                ckpt_data['epoch'] = epoch
+                ckpt_data['state_dict'] = model.state_dict()
+                ckpt_data['optimizer'] = trainer.optimizer.state_dict()
+
+                if save_model_secs is not None:
+                    if time.time() - recent_time > save_model_secs:
+                        torch.save(ckpt_data, os.path.join(FLAGS.output_path,
+                                                           'latest_checkpoint.pth.tar'))
+                        recent_time = time.time()
+
+                if save_model_steps is not None:
+                    if global_step in save_model_steps:
+                        torch.save(ckpt_data, os.path.join(FLAGS.output_path,
+                                                           f'epoch_{epoch:02d}.pth.tar'))
+
+            else:
+                if len(results) > 1:
+                    pprint.pprint(results)
+
+            if epoch < FLAGS.epochs:
+                frac_epoch = (global_step + 1) / nsteps
+                record = trainer(frac_epoch, *data)
+                record['data_load_dur'] = data_load_time
+                results = {'meta': {'step_in_epoch': step + 1,
+                                    'epoch': frac_epoch,
+                                    'wall_time': time.time()}
+                           }
+                if save_train_steps is not None:
+                    if step in save_train_steps:
+                        results[trainer.name] = record
+
+            data_load_start = time.time()
+
+
+class ImageNetTrain(object):
+
+    def __init__(self, model):
+        self.name = 'train'
+        self.model = model
+        self.data_loader = self.data()
+        self.optimizer = torch.optim.SGD(self.model.parameters(), FLAGS.lr, momentum=FLAGS.momentum,
+                                         weight_decay=FLAGS.weight_decay)
+        if FLAGS.optimizer == 'stepLR':
+            self.lr = torch.optim.lr_scheduler.StepLR(self.optimizer, gamma=FLAGS.step_factor,
+                                                      step_size=FLAGS.step_size)
+        elif FLAGS.optimizer == 'plateauLR':
+            self.lr = torch.optim.lr_scheduler.ReduceLROnPlateau(self.optimizer, factor=FLAGS.step_factor,
+                                                                 patience=FLAGS.step_size-1, threshold=0.01)
+        self.loss = nn.CrossEntropyLoss()
+        if FLAGS.ngpus > 0:
+            self.loss = self.loss.cuda()
+
+    def data(self):
+        dataset = torchvision.datasets.ImageFolder(
+            os.path.join(FLAGS.in_path, 'train'),
+            torchvision.transforms.Compose([
+                torchvision.transforms.RandomResizedCrop(224),
+                torchvision.transforms.RandomHorizontalFlip(),
+                torchvision.transforms.ToTensor(),
+                torchvision.transforms.Normalize(mean=norm_mean, std=norm_std)
+            ]))
+        data_loader = torch.utils.data.DataLoader(dataset,
+                                                  batch_size=FLAGS.batch_size,
+                                                  shuffle=True,
+                                                  num_workers=FLAGS.workers,
+                                                  pin_memory=True)
+
+        return data_loader
+
+    def __call__(self, frac_epoch, inp, target):
+        start = time.time()
+        if FLAGS.optimizer == 'stepLR':
+            self.lr.step(epoch=frac_epoch)
+        target = target.to(device)
+
+        output = self.model(inp)
+
+        record = {}
+        loss = self.loss(output, target)
+        record['loss'] = loss.item()
+        record['top1'], record['top5'] = accuracy(output, target, topk=(1, 5))
+        record['top1'] /= len(output)
+        record['top5'] /= len(output)
+        # record['learning_rate'] = self.lr.get_lr()[0]
+        record['learning_rate'] = self.optimizer.param_groups[0]["lr"]
+        self.optimizer.zero_grad()
+        loss.backward()
+        self.optimizer.step()
+
+        record['dur'] = time.time() - start
+        return record
+
+
+class ImageNetVal(object):
+
+    def __init__(self, model):
+        self.name = 'val'
+        self.model = model
+        self.data_loader = self.data()
+        self.loss = nn.CrossEntropyLoss(size_average=False)
+        self.loss = self.loss.to(device)
+
+    def data(self):
+        dataset = torchvision.datasets.ImageFolder(
+            os.path.join(FLAGS.in_path, 'val'),
+            torchvision.transforms.Compose([
+                torchvision.transforms.Resize(256),
+                torchvision.transforms.CenterCrop(224),
+                torchvision.transforms.ToTensor(),
+                torchvision.transforms.Normalize(mean=norm_mean, std=norm_std),
+            ]))
+        data_loader = torch.utils.data.DataLoader(dataset,
+                                                  batch_size=FLAGS.batch_size,
+                                                  shuffle=False,
+                                                  num_workers=FLAGS.workers,
+                                                  pin_memory=True)
+
+        return data_loader
+
+    def __call__(self):
+        self.model.eval()
+        start = time.time()
+        record = {'loss': 0, 'top1': 0, 'top5': 0}
+        with torch.no_grad():
+            for (inp, target) in tqdm.tqdm(self.data_loader, desc=self.name):
+                target = target.to(device)
+                output = self.model(inp)
+
+                record['loss'] += self.loss(output, target).item()
+                p1, p5 = accuracy(output, target, topk=(1, 5))
+                record['top1'] += p1
+                record['top5'] += p5
+
+        for key in record:
+            record[key] /= len(self.data_loader.dataset.samples)
+        record['dur'] = (time.time() - start) / len(self.data_loader)
+
+        return record
+
+
+def accuracy(output, target, topk=(1,)):
+    with torch.no_grad():
+        _, pred = output.topk(max(topk), dim=1, largest=True, sorted=True)
+        pred = pred.t()
+        correct = pred.eq(target.view(1, -1).expand_as(pred))
+        res = [correct[:k].sum().item() for k in topk]
+        return res
+
+
+if __name__ == '__main__':
+    fire.Fire(command=FIRE_FLAGS)