flair2seg & filter normal

+<?xml version="1.0" encoding="UTF-8"?>
+<module type="PYTHON_MODULE" version="4">
+  <component name="NewModuleRootManager">
+    <content url="file://$MODULE_DIR$" />
+    <orderEntry type="inheritedJdk" />
+    <orderEntry type="sourceFolder" forTests="false" />
+  </component>
+  <component name="TestRunnerService">
+    <option name="PROJECT_TEST_RUNNER" value="Unittests" />
+  </component>
+</module>
\ No newline at end of file

+<?xml version="1.0" encoding="UTF-8"?>
+<project version="4">
+  <component name="Encoding" addBOMForNewFiles="with NO BOM" />
+</project>
\ No newline at end of file

+<?xml version="1.0" encoding="UTF-8"?>
+<project version="4">
+  <component name="JavaScriptSettings">
+    <option name="languageLevel" value="ES6" />
+  </component>
+  <component name="ProjectRootManager" version="2" project-jdk-name="Python 3.7 (code) (3)" project-jdk-type="Python SDK" />
+</project>
\ No newline at end of file

+<?xml version="1.0" encoding="UTF-8"?>
+<project version="4">
+  <component name="ProjectModuleManager">
+    <modules>
+      <module fileurl="file://$PROJECT_DIR$/.idea/code.iml" filepath="$PROJECT_DIR$/.idea/code.iml" />
+    </modules>
+  </component>
+</project>
\ No newline at end of file

+<?xml version="1.0" encoding="UTF-8"?>
+<project version="4">
+  <component name="ChangeListManager">
+    <list default="true" id="38bbd81d-cf79-4a3f-8979-7c3ceb27bc32" name="Default Changelist" comment="" />
+    <option name="EXCLUDED_CONVERTED_TO_IGNORED" value="true" />
+    <option name="SHOW_DIALOG" value="false" />
+    <option name="HIGHLIGHT_CONFLICTS" value="true" />
+    <option name="HIGHLIGHT_NON_ACTIVE_CHANGELIST" value="false" />
+    <option name="LAST_RESOLUTION" value="IGNORE" />
+  </component>
+  <component name="ProjectFrameBounds">
+    <option name="x" value="-9" />
+    <option name="width" value="1825" />
+    <option name="height" value="1039" />
+  </component>
+  <component name="ProjectView">
+    <navigator proportions="" version="1">
+      <foldersAlwaysOnTop value="true" />
+    </navigator>
+    <panes>
+      <pane id="Scope" />
+      <pane id="ProjectPane" />
+    </panes>
+  </component>
+  <component name="PropertiesComponent">
+    <property name="WebServerToolWindowFactoryState" value="false" />
+    <property name="nodejs_interpreter_path.stuck_in_default_project" value="undefined stuck path" />
+    <property name="nodejs_npm_path_reset_for_default_project" value="true" />
+    <property name="settings.editor.selected.configurable" value="com.jetbrains.python.configuration.PyActiveSdkModuleConfigurable" />
+  </component>
+  <component name="RunDashboard">
+    <option name="ruleStates">
+      <list>
+        <RuleState>
+          <option name="name" value="ConfigurationTypeDashboardGroupingRule" />
+        </RuleState>
+        <RuleState>
+          <option name="name" value="StatusDashboardGroupingRule" />
+        </RuleState>
+      </list>
+    </option>
+  </component>
+  <component name="SvnConfiguration">
+    <configuration />
+  </component>
+  <component name="TaskManager">
+    <task active="true" id="Default" summary="Default task">
+      <changelist id="38bbd81d-cf79-4a3f-8979-7c3ceb27bc32" name="Default Changelist" comment="" />
+      <created>1584871834688</created>
+      <option name="number" value="Default" />
+      <option name="presentableId" value="Default" />
+      <updated>1584871834688</updated>
+      <workItem from="1584871840396" duration="72000" />
+    </task>
+    <servers />
+  </component>
+  <component name="TimeTrackingManager">
+    <option name="totallyTimeSpent" value="72000" />
+  </component>
+  <component name="ToolWindowManager">
+    <frame x="-7" y="0" width="1460" height="831" extended-state="0" />
+    <layout>
+      <window_info id="Favorites" side_tool="true" />
+      <window_info active="true" content_ui="combo" id="Project" order="0" visible="true" weight="0.25" />
+      <window_info id="Structure" order="1" side_tool="true" weight="0.25" />
+      <window_info anchor="bottom" id="Docker" show_stripe_button="false" />
+      <window_info anchor="bottom" id="Database Changes" />
+      <window_info anchor="bottom" id="Version Control" />
+      <window_info anchor="bottom" id="Python Console" />
+      <window_info anchor="bottom" id="Terminal" />
+      <window_info anchor="bottom" id="Event Log" side_tool="true" />
+      <window_info anchor="bottom" id="Message" order="0" />
+      <window_info anchor="bottom" id="Find" order="1" />
+      <window_info anchor="bottom" id="Run" order="2" />
+      <window_info anchor="bottom" id="Debug" order="3" weight="0.4" />
+      <window_info anchor="bottom" id="Cvs" order="4" weight="0.25" />
+      <window_info anchor="bottom" id="Inspection" order="5" weight="0.4" />
+      <window_info anchor="bottom" id="TODO" order="6" />
+      <window_info anchor="right" id="SciView" />
+      <window_info anchor="right" id="Database" />
+      <window_info anchor="right" id="Commander" internal_type="SLIDING" order="0" type="SLIDING" weight="0.4" />
+      <window_info anchor="right" id="Ant Build" order="1" weight="0.25" />
+      <window_info anchor="right" content_ui="combo" id="Hierarchy" order="2" weight="0.25" />
+    </layout>
+  </component>
+  <component name="TypeScriptGeneratedFilesManager">
+    <option name="version" value="1" />
+  </component>
+</project>
\ No newline at end of file

+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import pandas as pd"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/html": [
+       "<div>\n",
+       "<style scoped>\n",
+       "    .dataframe tbody tr th:only-of-type {\n",
+       "        vertical-align: middle;\n",
+       "    }\n",
+       "\n",
+       "    .dataframe tbody tr th {\n",
+       "        vertical-align: top;\n",
+       "    }\n",
+       "\n",
+       "    .dataframe thead th {\n",
+       "        text-align: right;\n",
+       "    }\n",
+       "</style>\n",
+       "<table border=\"1\" class=\"dataframe\">\n",
+       "  <thead>\n",
+       "    <tr style=\"text-align: right;\">\n",
+       "      <th></th>\n",
+       "      <th>id</th>\n",
+       "      <th>Label</th>\n",
+       "      <th>Subject</th>\n",
+       "      <th>Date</th>\n",
+       "      <th>Gender</th>\n",
+       "      <th>Age</th>\n",
+       "      <th>mmse</th>\n",
+       "      <th>ageAtEntry</th>\n",
+       "      <th>cdr</th>\n",
+       "      <th>commun</th>\n",
+       "      <th>...</th>\n",
+       "      <th>memory</th>\n",
+       "      <th>orient</th>\n",
+       "      <th>perscare</th>\n",
+       "      <th>apoe</th>\n",
+       "      <th>sumbox</th>\n",
+       "      <th>acsparnt</th>\n",
+       "      <th>height</th>\n",
+       "      <th>weight</th>\n",
+       "      <th>primStudy</th>\n",
+       "      <th>acsStudy</th>\n",
+       "    </tr>\n",
+       "  </thead>\n",
+       "  <tbody>\n",
+       "    <tr>\n",
+       "      <th>0</th>\n",
+       "      <td>/@WEBAPP/images/r.gif</td>\n",
+       "      <td>OAS30001_ClinicalData_d3025</td>\n",
+       "      <td>OAS30001</td>\n",
+       "      <td>NaN</td>\n",
+       "      <td>female</td>\n",
+       "      <td>NaN</td>\n",
+       "      <td>30.0</td>\n",
+       "      <td>65.149895</td>\n",
+       "      <td>0.0</td>\n",
+       "      <td>0.0</td>\n",
+       "      <td>...</td>\n",
+       "      <td>0.0</td>\n",
+       "      <td>0.0</td>\n",
+       "      <td>0.0</td>\n",
+       "      <td>23.0</td>\n",
+       "      <td>0.0</td>\n",
+       "      <td>NaN</td>\n",
+       "      <td>64.0</td>\n",
+       "      <td>180.0</td>\n",
+       "      <td>NaN</td>\n",
+       "      <td>NaN</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>1</th>\n",
+       "      <td>/@WEBAPP/images/r.gif</td>\n",
+       "      <td>OAS30001_ClinicalData_d3977</td>\n",
+       "      <td>OAS30001</td>\n",
+       "      <td>NaN</td>\n",
+       "      <td>female</td>\n",
+       "      <td>NaN</td>\n",
+       "      <td>29.0</td>\n",
+       "      <td>65.149895</td>\n",
+       "      <td>0.0</td>\n",
+       "      <td>0.0</td>\n",
+       "      <td>...</td>\n",
+       "      <td>0.0</td>\n",
+       "      <td>0.0</td>\n",
+       "      <td>0.0</td>\n",
+       "      <td>23.0</td>\n",
+       "      <td>0.0</td>\n",
+       "      <td>NaN</td>\n",
+       "      <td>NaN</td>\n",
+       "      <td>NaN</td>\n",
+       "      <td>NaN</td>\n",
+       "      <td>NaN</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>2</th>\n",
+       "      <td>/@WEBAPP/images/r.gif</td>\n",
+       "      <td>OAS30001_ClinicalData_d3332</td>\n",
+       "      <td>OAS30001</td>\n",
+       "      <td>NaN</td>\n",
+       "      <td>female</td>\n",
+       "      <td>NaN</td>\n",
+       "      <td>30.0</td>\n",
+       "      <td>65.149895</td>\n",
+       "      <td>0.0</td>\n",
+       "      <td>0.0</td>\n",
+       "      <td>...</td>\n",
+       "      <td>0.0</td>\n",
+       "      <td>0.0</td>\n",
+       "      <td>0.0</td>\n",
+       "      <td>23.0</td>\n",
+       "      <td>0.0</td>\n",
+       "      <td>NaN</td>\n",
+       "      <td>63.0</td>\n",
+       "      <td>185.0</td>\n",
+       "      <td>NaN</td>\n",
+       "      <td>NaN</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>3</th>\n",
+       "      <td>/@WEBAPP/images/r.gif</td>\n",
+       "      <td>OAS30001_ClinicalData_d0000</td>\n",
+       "      <td>OAS30001</td>\n",
+       "      <td>NaN</td>\n",
+       "      <td>female</td>\n",
+       "      <td>NaN</td>\n",
+       "      <td>28.0</td>\n",
+       "      <td>65.149895</td>\n",
+       "      <td>0.0</td>\n",
+       "      <td>0.0</td>\n",
+       "      <td>...</td>\n",
+       "      <td>0.0</td>\n",
+       "      <td>0.0</td>\n",
+       "      <td>0.0</td>\n",
+       "      <td>23.0</td>\n",
+       "      <td>0.0</td>\n",
+       "      <td>NaN</td>\n",
+       "      <td>NaN</td>\n",
+       "      <td>NaN</td>\n",
+       "      <td>NaN</td>\n",
+       "      <td>NaN</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>4</th>\n",
+       "      <td>/@WEBAPP/images/r.gif</td>\n",
+       "      <td>OAS30001_ClinicalData_d1456</td>\n",
+       "      <td>OAS30001</td>\n",
+       "      <td>NaN</td>\n",
+       "      <td>female</td>\n",
+       "      <td>NaN</td>\n",
+       "      <td>30.0</td>\n",
+       "      <td>65.149895</td>\n",
+       "      <td>0.0</td>\n",
+       "      <td>0.0</td>\n",
+       "      <td>...</td>\n",
+       "      <td>0.0</td>\n",
+       "      <td>0.0</td>\n",
+       "      <td>0.0</td>\n",
+       "      <td>23.0</td>\n",
+       "      <td>0.0</td>\n",
+       "      <td>NaN</td>\n",
+       "      <td>63.0</td>\n",
+       "      <td>173.0</td>\n",
+       "      <td>NaN</td>\n",
+       "      <td>NaN</td>\n",
+       "    </tr>\n",
+       "  </tbody>\n",
+       "</table>\n",
+       "<p>5 rows × 27 columns</p>\n",
+       "</div>"
+      ],
+      "text/plain": [
+       "                      id                        Label   Subject  Date  Gender  \\\n",
+       "0  /@WEBAPP/images/r.gif  OAS30001_ClinicalData_d3025  OAS30001   NaN  female   \n",
+       "1  /@WEBAPP/images/r.gif  OAS30001_ClinicalData_d3977  OAS30001   NaN  female   \n",
+       "2  /@WEBAPP/images/r.gif  OAS30001_ClinicalData_d3332  OAS30001   NaN  female   \n",
+       "3  /@WEBAPP/images/r.gif  OAS30001_ClinicalData_d0000  OAS30001   NaN  female   \n",
+       "4  /@WEBAPP/images/r.gif  OAS30001_ClinicalData_d1456  OAS30001   NaN  female   \n",
+       "\n",
+       "   Age  mmse  ageAtEntry  cdr  commun  ... memory orient perscare  apoe  \\\n",
+       "0  NaN  30.0   65.149895  0.0     0.0  ...    0.0    0.0      0.0  23.0   \n",
+       "1  NaN  29.0   65.149895  0.0     0.0  ...    0.0    0.0      0.0  23.0   \n",
+       "2  NaN  30.0   65.149895  0.0     0.0  ...    0.0    0.0      0.0  23.0   \n",
+       "3  NaN  28.0   65.149895  0.0     0.0  ...    0.0    0.0      0.0  23.0   \n",
+       "4  NaN  30.0   65.149895  0.0     0.0  ...    0.0    0.0      0.0  23.0   \n",
+       "\n",
+       "  sumbox  acsparnt  height  weight  primStudy  acsStudy  \n",
+       "0    0.0       NaN    64.0   180.0        NaN       NaN  \n",
+       "1    0.0       NaN     NaN     NaN        NaN       NaN  \n",
+       "2    0.0       NaN    63.0   185.0        NaN       NaN  \n",
+       "3    0.0       NaN     NaN     NaN        NaN       NaN  \n",
+       "4    0.0       NaN    63.0   173.0        NaN       NaN  \n",
+       "\n",
+       "[5 rows x 27 columns]"
+      ]
+     },
+     "execution_count": 3,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "all_data = pd.read_csv(\"..\\data\\ADRC clinical data_all.csv\")\n",
+    "\n",
+    "all_data.head()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 18,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Subject\n",
+      "OAS30001    0.0\n",
+      "OAS30002    0.0\n",
+      "OAS30003    0.0\n",
+      "OAS30004    0.0\n",
+      "OAS30005    0.0\n",
+      "           ... \n",
+      "OAS31168    0.0\n",
+      "OAS31169    3.0\n",
+      "OAS31170    2.0\n",
+      "OAS31171    2.0\n",
+      "OAS31172    0.0\n",
+      "Name: cdr, Length: 1098, dtype: float64\n"
+     ]
+    }
+   ],
+   "source": [
+    "ad = all_data.groupby(['Subject'])['cdr'].max()\n",
+    "print(ad)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 21,
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "'OAS30001'"
+      ]
+     },
+     "execution_count": 21,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "ad.index[0]"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 22,
+   "metadata": {},
+   "outputs": [
+    {
+     "ename": "SyntaxError",
+     "evalue": "unexpected EOF while parsing (<ipython-input-22-b8a078b72aca>, line 5)",
+     "output_type": "error",
+     "traceback": [
+      "\u001b[1;36m  File \u001b[1;32m\"<ipython-input-22-b8a078b72aca>\"\u001b[1;36m, line \u001b[1;32m5\u001b[0m\n\u001b[1;33m    #print(filtered)\u001b[0m\n\u001b[1;37m                    ^\u001b[0m\n\u001b[1;31mSyntaxError\u001b[0m\u001b[1;31m:\u001b[0m unexpected EOF while parsing\n"
+     ]
+    }
+   ],
+   "source": [
+    "filtered = []\n",
+    "for i, val in enumerate(ad):\n",
+    "    if ad[i] == 0:\n",
+    "        filtered.append(ad.index[i])\n",
+    "#print(filtered)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 23,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "df_filtered = pd.DataFrame(filtered)\n",
+    "df_filtered.to_csv('..\\data\\ADRC clinical data_normal.csv')"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "ML",
+   "language": "python",
+   "name": "ml"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.7.4"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}

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

+"""
+Reference :
+- https://github.com/hysts/pytorch_image_classification/blob/master/augmentations/mixup.py
+- https://github.com/tensorflow/tpu/blob/master/models/official/efficientnet/imagenet_input.py#L120
+"""
+
+import numpy as np
+import torch
+
+from FastAutoAugment.metrics import CrossEntropyLabelSmooth
+
+
+def mixup(data, targets, alpha):
+    indices = torch.randperm(data.size(0))
+    shuffled_data = data[indices]
+    shuffled_targets = targets[indices]
+
+    lam = np.random.beta(alpha, alpha)
+    lam = max(lam, 1. - lam)
+    assert 0.0 <= lam <= 1.0, lam
+    data = data * lam + shuffled_data * (1 - lam)
+
+    return data, targets, shuffled_targets, lam
+
+
+class CrossEntropyMixUpLabelSmooth(torch.nn.Module):
+    def __init__(self, num_classes, epsilon, reduction='mean'):
+        super(CrossEntropyMixUpLabelSmooth, self).__init__()
+        self.ce = CrossEntropyLabelSmooth(num_classes, epsilon, reduction=reduction)
+
+    def forward(self, input, target1, target2, lam):  # pylint: disable=redefined-builtin
+        return lam * self.ce(input, target1) + (1 - lam) * self.ce(input, target2)

+# code in this file is adpated from rpmcruz/autoaugment
+# https://github.com/rpmcruz/autoaugment/blob/master/transformations.py
+import random
+
+import PIL, PIL.ImageOps, PIL.ImageEnhance, PIL.ImageDraw
+import numpy as np
+import torch
+from torchvision.transforms.transforms import Compose
+
+random_mirror = True
+
+
+def ShearX(img, v):  # [-0.3, 0.3]
+    assert -0.3 <= v <= 0.3
+    if random_mirror and random.random() > 0.5:
+        v = -v
+    return img.transform(img.size, PIL.Image.AFFINE, (1, v, 0, 0, 1, 0))
+
+
+def ShearY(img, v):  # [-0.3, 0.3]
+    assert -0.3 <= v <= 0.3
+    if random_mirror and random.random() > 0.5:
+        v = -v
+    return img.transform(img.size, PIL.Image.AFFINE, (1, 0, 0, v, 1, 0))
+
+
+def TranslateX(img, v):  # [-150, 150] => percentage: [-0.45, 0.45]
+    assert -0.45 <= v <= 0.45
+    if random_mirror and random.random() > 0.5:
+        v = -v
+    v = v * img.size[0]
+    return img.transform(img.size, PIL.Image.AFFINE, (1, 0, v, 0, 1, 0))
+
+
+def TranslateY(img, v):  # [-150, 150] => percentage: [-0.45, 0.45]
+    assert -0.45 <= v <= 0.45
+    if random_mirror and random.random() > 0.5:
+        v = -v
+    v = v * img.size[1]
+    return img.transform(img.size, PIL.Image.AFFINE, (1, 0, 0, 0, 1, v))
+
+
+def TranslateXAbs(img, v):  # [-150, 150] => percentage: [-0.45, 0.45]
+    assert 0 <= v <= 10
+    if random.random() > 0.5:
+        v = -v
+    return img.transform(img.size, PIL.Image.AFFINE, (1, 0, v, 0, 1, 0))
+
+
+def TranslateYAbs(img, v):  # [-150, 150] => percentage: [-0.45, 0.45]
+    assert 0 <= v <= 10
+    if random.random() > 0.5:
+        v = -v
+    return img.transform(img.size, PIL.Image.AFFINE, (1, 0, 0, 0, 1, v))
+
+
+def Rotate(img, v):  # [-30, 30]
+    assert -30 <= v <= 30
+    if random_mirror and random.random() > 0.5:
+        v = -v
+    return img.rotate(v)
+
+
+def AutoContrast(img, _):
+    return PIL.ImageOps.autocontrast(img)
+
+
+def Invert(img, _):
+    return PIL.ImageOps.invert(img)
+
+
+def Equalize(img, _):
+    return PIL.ImageOps.equalize(img)
+
+
+def Flip(img, _):  # not from the paper
+    return PIL.ImageOps.mirror(img)
+
+
+def Solarize(img, v):  # [0, 256]
+    assert 0 <= v <= 256
+    return PIL.ImageOps.solarize(img, v)
+
+
+def Posterize(img, v):  # [4, 8]
+    assert 4 <= v <= 8
+    v = int(v)
+    return PIL.ImageOps.posterize(img, v)
+
+
+def Posterize2(img, v):  # [0, 4]
+    assert 0 <= v <= 4
+    v = int(v)
+    return PIL.ImageOps.posterize(img, v)
+
+
+def Contrast(img, v):  # [0.1,1.9]
+    assert 0.1 <= v <= 1.9
+    return PIL.ImageEnhance.Contrast(img).enhance(v)
+
+
+def Color(img, v):  # [0.1,1.9]
+    assert 0.1 <= v <= 1.9
+    return PIL.ImageEnhance.Color(img).enhance(v)
+
+
+def Brightness(img, v):  # [0.1,1.9]
+    assert 0.1 <= v <= 1.9
+    return PIL.ImageEnhance.Brightness(img).enhance(v)
+
+
+def Sharpness(img, v):  # [0.1,1.9]
+    assert 0.1 <= v <= 1.9
+    return PIL.ImageEnhance.Sharpness(img).enhance(v)
+
+
+def Cutout(img, v):  # [0, 60] => percentage: [0, 0.2]
+    assert 0.0 <= v <= 0.2
+    if v <= 0.:
+        return img
+
+    v = v * img.size[0]
+    return CutoutAbs(img, v)
+
+
+def CutoutAbs(img, v):  # [0, 60] => percentage: [0, 0.2]
+    # assert 0 <= v <= 20
+    if v < 0:
+        return img
+    w, h = img.size
+    x0 = np.random.uniform(w)
+    y0 = np.random.uniform(h)
+
+    x0 = int(max(0, x0 - v / 2.))
+    y0 = int(max(0, y0 - v / 2.))
+    x1 = min(w, x0 + v)
+    y1 = min(h, y0 + v)
+
+    xy = (x0, y0, x1, y1)
+    color = (125, 123, 114)
+    # color = (0, 0, 0)
+    img = img.copy()
+    PIL.ImageDraw.Draw(img).rectangle(xy, color)
+    return img
+
+
+def SamplePairing(imgs):  # [0, 0.4]
+    def f(img1, v):
+        i = np.random.choice(len(imgs))
+        img2 = PIL.Image.fromarray(imgs[i])
+        return PIL.Image.blend(img1, img2, v)
+
+    return f
+
+
+def augment_list(for_autoaug=True):  # 16 oeprations and their ranges
+    l = [
+        (ShearX, -0.3, 0.3),  # 0
+        (ShearY, -0.3, 0.3),  # 1
+        (TranslateX, -0.45, 0.45),  # 2
+        (TranslateY, -0.45, 0.45),  # 3
+        (Rotate, -30, 30),  # 4
+        (AutoContrast, 0, 1),  # 5
+        (Invert, 0, 1),  # 6
+        (Equalize, 0, 1),  # 7
+        (Solarize, 0, 256),  # 8
+        (Posterize, 4, 8),  # 9
+        (Contrast, 0.1, 1.9),  # 10
+        (Color, 0.1, 1.9),  # 11
+        (Brightness, 0.1, 1.9),  # 12
+        (Sharpness, 0.1, 1.9),  # 13
+        (Cutout, 0, 0.2),  # 14
+        # (SamplePairing(imgs), 0, 0.4),  # 15
+    ]
+    if for_autoaug:
+        l += [
+            (CutoutAbs, 0, 20),  # compatible with auto-augment
+            (Posterize2, 0, 4),  # 9
+            (TranslateXAbs, 0, 10),  # 9
+            (TranslateYAbs, 0, 10),  # 9
+        ]
+    return l
+
+
+augment_dict = {fn.__name__: (fn, v1, v2) for fn, v1, v2 in augment_list()}
+
+
+def get_augment(name):
+    return augment_dict[name]
+
+
+def apply_augment(img, name, level):
+    augment_fn, low, high = get_augment(name)
+    return augment_fn(img.copy(), level * (high - low) + low)
+
+
+class Lighting(object):
+    """Lighting noise(AlexNet - style PCA - based noise)"""
+
+    def __init__(self, alphastd, eigval, eigvec):
+        self.alphastd = alphastd
+        self.eigval = torch.Tensor(eigval)
+        self.eigvec = torch.Tensor(eigvec)
+
+    def __call__(self, img):
+        if self.alphastd == 0:
+            return img
+
+        alpha = img.new().resize_(3).normal_(0, self.alphastd)
+        rgb = self.eigvec.type_as(img).clone() \
+            .mul(alpha.view(1, 3).expand(3, 3)) \
+            .mul(self.eigval.view(1, 3).expand(3, 3)) \
+            .sum(1).squeeze()
+
+        return img.add(rgb.view(3, 1, 1).expand_as(img))

+import copy
+import logging
+import warnings
+
+formatter = logging.Formatter('[%(asctime)s] [%(name)s] [%(levelname)s] %(message)s')
+warnings.filterwarnings("ignore", "(Possibly )?corrupt EXIF data", UserWarning)
+warnings.filterwarnings("ignore", "DeprecationWarning: 'saved_variables' is deprecated", UserWarning)
+
+
+def get_logger(name, level=logging.DEBUG):
+    logger = logging.getLogger(name)
+    logger.handlers.clear()
+    logger.setLevel(level)
+    ch = logging.StreamHandler()
+    ch.setLevel(level)
+    ch.setFormatter(formatter)
+    logger.addHandler(ch)
+    return logger
+
+
+def add_filehandler(logger, filepath, level=logging.DEBUG):
+    fh = logging.FileHandler(filepath)
+    fh.setLevel(level)
+    fh.setFormatter(formatter)
+    logger.addHandler(fh)
+
+
+class EMA:
+    def __init__(self, mu):
+        self.mu = mu
+        self.shadow = {}
+
+    def state_dict(self):
+        return copy.deepcopy(self.shadow)
+
+    def __len__(self):
+        return len(self.shadow)
+
+    def __call__(self, module, step=None):
+        if step is None:
+            mu = self.mu
+        else:
+            # see : https://www.tensorflow.org/versions/r1.15/api_docs/python/tf/train/ExponentialMovingAverage?hl=PL
+            mu = min(self.mu, (1. + step) / (10 + step))
+
+        for name, x in module.state_dict().items():
+            if name in self.shadow:
+                new_average = (1.0 - mu) * x + mu * self.shadow[name]
+                self.shadow[name] = new_average.clone()
+            else:
+                self.shadow[name] = x.clone()

+import logging
+
+import numpy as np
+import os
+
+import math
+import random
+import torch
+import torchvision
+from PIL import Image
+
+from torch.utils.data import SubsetRandomSampler, Sampler, Subset, ConcatDataset
+import torch.distributed as dist
+from torchvision.transforms import transforms
+from sklearn.model_selection import StratifiedShuffleSplit
+from theconf import Config as C
+
+from FastAutoAugment.archive import arsaug_policy, autoaug_policy, autoaug_paper_cifar10, fa_reduced_cifar10, fa_reduced_svhn, fa_resnet50_rimagenet
+from FastAutoAugment.augmentations import *
+from FastAutoAugment.common import get_logger
+from FastAutoAugment.imagenet import ImageNet
+from FastAutoAugment.networks.efficientnet_pytorch.model import EfficientNet
+
+logger = get_logger('Fast AutoAugment')
+logger.setLevel(logging.INFO)
+_IMAGENET_PCA = {
+    'eigval': [0.2175, 0.0188, 0.0045],
+    'eigvec': [
+        [-0.5675,  0.7192,  0.4009],
+        [-0.5808, -0.0045, -0.8140],
+        [-0.5836, -0.6948,  0.4203],
+    ]
+}
+_CIFAR_MEAN, _CIFAR_STD = (0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010)
+
+
+def get_dataloaders(dataset, batch, dataroot, split=0.15, split_idx=0, multinode=False, target_lb=-1):
+    if 'cifar' in dataset or 'svhn' in dataset:
+        transform_train = transforms.Compose([
+            transforms.RandomCrop(32, padding=4),
+            transforms.RandomHorizontalFlip(),
+            transforms.ToTensor(),
+            transforms.Normalize(_CIFAR_MEAN, _CIFAR_STD),
+        ])
+        transform_test = transforms.Compose([
+            transforms.ToTensor(),
+            transforms.Normalize(_CIFAR_MEAN, _CIFAR_STD),
+        ])
+    elif 'imagenet' in dataset:
+        input_size = 224
+        sized_size = 256
+
+        if 'efficientnet' in C.get()['model']['type']:
+            input_size = EfficientNet.get_image_size(C.get()['model']['type'])
+            sized_size = input_size + 32    # TODO
+            # sized_size = int(round(input_size / 224. * 256))
+            # sized_size = input_size
+            logger.info('size changed to %d/%d.' % (input_size, sized_size))
+
+        transform_train = transforms.Compose([
+            EfficientNetRandomCrop(input_size),
+            transforms.Resize((input_size, input_size), interpolation=Image.BICUBIC),
+            # transforms.RandomResizedCrop(input_size, scale=(0.1, 1.0), interpolation=Image.BICUBIC),
+            transforms.RandomHorizontalFlip(),
+            transforms.ColorJitter(
+                brightness=0.4,
+                contrast=0.4,
+                saturation=0.4,
+            ),
+            transforms.ToTensor(),
+            Lighting(0.1, _IMAGENET_PCA['eigval'], _IMAGENET_PCA['eigvec']),
+            transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
+        ])
+
+        transform_test = transforms.Compose([
+            EfficientNetCenterCrop(input_size),
+            transforms.Resize((input_size, input_size), interpolation=Image.BICUBIC),
+            transforms.ToTensor(),
+            transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
+        ])
+
+    else:
+        raise ValueError('dataset=%s' % dataset)
+
+    total_aug = augs = None
+    if isinstance(C.get()['aug'], list):
+        logger.debug('augmentation provided.')
+        transform_train.transforms.insert(0, Augmentation(C.get()['aug']))
+    else:
+        logger.debug('augmentation: %s' % C.get()['aug'])
+        if C.get()['aug'] == 'fa_reduced_cifar10':
+            transform_train.transforms.insert(0, Augmentation(fa_reduced_cifar10()))
+
+        elif C.get()['aug'] == 'fa_reduced_imagenet':
+            transform_train.transforms.insert(0, Augmentation(fa_resnet50_rimagenet()))
+
+        elif C.get()['aug'] == 'fa_reduced_svhn':
+            transform_train.transforms.insert(0, Augmentation(fa_reduced_svhn()))
+
+        elif C.get()['aug'] == 'arsaug':
+            transform_train.transforms.insert(0, Augmentation(arsaug_policy()))
+        elif C.get()['aug'] == 'autoaug_cifar10':
+            transform_train.transforms.insert(0, Augmentation(autoaug_paper_cifar10()))
+        elif C.get()['aug'] == 'autoaug_extend':
+            transform_train.transforms.insert(0, Augmentation(autoaug_policy()))
+        elif C.get()['aug'] in ['default']:
+            pass
+        else:
+            raise ValueError('not found augmentations. %s' % C.get()['aug'])
+
+    if C.get()['cutout'] > 0:
+        transform_train.transforms.append(CutoutDefault(C.get()['cutout']))
+
+    if dataset == 'cifar10':
+        total_trainset = torchvision.datasets.CIFAR10(root=dataroot, train=True, download=True, transform=transform_train)
+        testset = torchvision.datasets.CIFAR10(root=dataroot, train=False, download=True, transform=transform_test)
+    elif dataset == 'reduced_cifar10':
+        total_trainset = torchvision.datasets.CIFAR10(root=dataroot, train=True, download=True, transform=transform_train)
+        sss = StratifiedShuffleSplit(n_splits=1, test_size=46000, random_state=0)   # 4000 trainset
+        sss = sss.split(list(range(len(total_trainset))), total_trainset.targets)
+        train_idx, valid_idx = next(sss)
+        targets = [total_trainset.targets[idx] for idx in train_idx]
+        total_trainset = Subset(total_trainset, train_idx)
+        total_trainset.targets = targets
+
+        testset = torchvision.datasets.CIFAR10(root=dataroot, train=False, download=True, transform=transform_test)
+    elif dataset == 'cifar100':
+        total_trainset = torchvision.datasets.CIFAR100(root=dataroot, train=True, download=True, transform=transform_train)
+        testset = torchvision.datasets.CIFAR100(root=dataroot, train=False, download=True, transform=transform_test)
+    elif dataset == 'svhn':
+        trainset = torchvision.datasets.SVHN(root=dataroot, split='train', download=True, transform=transform_train)
+        extraset = torchvision.datasets.SVHN(root=dataroot, split='extra', download=True, transform=transform_train)
+        total_trainset = ConcatDataset([trainset, extraset])
+        testset = torchvision.datasets.SVHN(root=dataroot, split='test', download=True, transform=transform_test)
+    elif dataset == 'reduced_svhn':
+        total_trainset = torchvision.datasets.SVHN(root=dataroot, split='train', download=True, transform=transform_train)
+        sss = StratifiedShuffleSplit(n_splits=1, test_size=73257-1000, random_state=0)  # 1000 trainset
+        sss = sss.split(list(range(len(total_trainset))), total_trainset.targets)
+        train_idx, valid_idx = next(sss)
+        targets = [total_trainset.targets[idx] for idx in train_idx]
+        total_trainset = Subset(total_trainset, train_idx)
+        total_trainset.targets = targets
+
+        testset = torchvision.datasets.SVHN(root=dataroot, split='test', download=True, transform=transform_test)
+    elif dataset == 'imagenet':
+        total_trainset = ImageNet(root=os.path.join(dataroot, 'imagenet-pytorch'), transform=transform_train)
+        testset = ImageNet(root=os.path.join(dataroot, 'imagenet-pytorch'), split='val', transform=transform_test)
+
+        # compatibility
+        total_trainset.targets = [lb for _, lb in total_trainset.samples]
+    elif dataset == 'reduced_imagenet':
+        # randomly chosen indices
+#         idx120 = sorted(random.sample(list(range(1000)), k=120))
+        idx120 = [16, 23, 52, 57, 76, 93, 95, 96, 99, 121, 122, 128, 148, 172, 181, 189, 202, 210, 232, 238, 257, 258, 259, 277, 283, 289, 295, 304, 307, 318, 322, 331, 337, 338, 345, 350, 361, 375, 376, 381, 388, 399, 401, 408, 424, 431, 432, 440, 447, 462, 464, 472, 483, 497, 506, 512, 530, 541, 553, 554, 557, 564, 570, 584, 612, 614, 619, 626, 631, 632, 650, 657, 658, 660, 674, 675, 680, 682, 691, 695, 699, 711, 734, 736, 741, 754, 757, 764, 769, 770, 780, 781, 787, 797, 799, 811, 822, 829, 830, 835, 837, 842, 843, 845, 873, 883, 897, 900, 902, 905, 913, 920, 925, 937, 938, 940, 941, 944, 949, 959]
+        total_trainset = ImageNet(root=os.path.join(dataroot, 'imagenet-pytorch'), transform=transform_train)
+        testset = ImageNet(root=os.path.join(dataroot, 'imagenet-pytorch'), split='val', transform=transform_test)
+
+        # compatibility
+        total_trainset.targets = [lb for _, lb in total_trainset.samples]
+
+        sss = StratifiedShuffleSplit(n_splits=1, test_size=len(total_trainset) - 50000, random_state=0)  # 4000 trainset
+        sss = sss.split(list(range(len(total_trainset))), total_trainset.targets)
+        train_idx, valid_idx = next(sss)
+
+        # filter out
+        train_idx = list(filter(lambda x: total_trainset.labels[x] in idx120, train_idx))
+        valid_idx = list(filter(lambda x: total_trainset.labels[x] in idx120, valid_idx))
+        test_idx = list(filter(lambda x: testset.samples[x][1] in idx120, range(len(testset))))
+
+        targets = [idx120.index(total_trainset.targets[idx]) for idx in train_idx]
+        for idx in range(len(total_trainset.samples)):
+            if total_trainset.samples[idx][1] not in idx120:
+                continue
+            total_trainset.samples[idx] = (total_trainset.samples[idx][0], idx120.index(total_trainset.samples[idx][1]))
+        total_trainset = Subset(total_trainset, train_idx)
+        total_trainset.targets = targets
+
+        for idx in range(len(testset.samples)):
+            if testset.samples[idx][1] not in idx120:
+                continue
+            testset.samples[idx] = (testset.samples[idx][0], idx120.index(testset.samples[idx][1]))
+        testset = Subset(testset, test_idx)
+        print('reduced_imagenet train=', len(total_trainset))
+    else:
+        raise ValueError('invalid dataset name=%s' % dataset)
+
+    if total_aug is not None and augs is not None:
+        total_trainset.set_preaug(augs, total_aug)
+        print('set_preaug-')
+
+    train_sampler = None
+    if split > 0.0:
+        sss = StratifiedShuffleSplit(n_splits=5, test_size=split, random_state=0)
+        sss = sss.split(list(range(len(total_trainset))), total_trainset.targets)
+        for _ in range(split_idx + 1):
+            train_idx, valid_idx = next(sss)
+
+        if target_lb >= 0:
+            train_idx = [i for i in train_idx if total_trainset.targets[i] == target_lb]
+            valid_idx = [i for i in valid_idx if total_trainset.targets[i] == target_lb]
+
+        train_sampler = SubsetRandomSampler(train_idx)
+        valid_sampler = SubsetSampler(valid_idx)
+
+        if multinode:
+            train_sampler = torch.utils.data.distributed.DistributedSampler(Subset(total_trainset, train_idx), num_replicas=dist.get_world_size(), rank=dist.get_rank())
+    else:
+        valid_sampler = SubsetSampler([])
+
+        if multinode:
+            train_sampler = torch.utils.data.distributed.DistributedSampler(total_trainset, num_replicas=dist.get_world_size(), rank=dist.get_rank())
+            logger.info(f'----- dataset with DistributedSampler  {dist.get_rank()}/{dist.get_world_size()}')
+
+    trainloader = torch.utils.data.DataLoader(
+        total_trainset, batch_size=batch, shuffle=True if train_sampler is None else False, num_workers=8, pin_memory=True,
+        sampler=train_sampler, drop_last=True)
+    validloader = torch.utils.data.DataLoader(
+        total_trainset, batch_size=batch, shuffle=False, num_workers=4, pin_memory=True,
+        sampler=valid_sampler, drop_last=False)
+
+    testloader = torch.utils.data.DataLoader(
+        testset, batch_size=batch, shuffle=False, num_workers=8, pin_memory=True,
+        drop_last=False
+    )
+    return train_sampler, trainloader, validloader, testloader
+
+
+class CutoutDefault(object):
+    """
+    Reference : https://github.com/quark0/darts/blob/master/cnn/utils.py
+    """
+    def __init__(self, length):
+        self.length = length
+
+    def __call__(self, img):
+        h, w = img.size(1), img.size(2)
+        mask = np.ones((h, w), np.float32)
+        y = np.random.randint(h)
+        x = np.random.randint(w)
+
+        y1 = np.clip(y - self.length // 2, 0, h)
+        y2 = np.clip(y + self.length // 2, 0, h)
+        x1 = np.clip(x - self.length // 2, 0, w)
+        x2 = np.clip(x + self.length // 2, 0, w)
+
+        mask[y1: y2, x1: x2] = 0.
+        mask = torch.from_numpy(mask)
+        mask = mask.expand_as(img)
+        img *= mask
+        return img
+
+
+class Augmentation(object):
+    def __init__(self, policies):
+        self.policies = policies
+
+    def __call__(self, img):
+        for _ in range(1):
+            policy = random.choice(self.policies)
+            for name, pr, level in policy:
+                if random.random() > pr:
+                    continue
+                img = apply_augment(img, name, level)
+        return img
+
+
+class EfficientNetRandomCrop:
+    def __init__(self, imgsize, min_covered=0.1, aspect_ratio_range=(3./4, 4./3), area_range=(0.08, 1.0), max_attempts=10):
+        assert 0.0 < min_covered
+        assert 0 < aspect_ratio_range[0] <= aspect_ratio_range[1]
+        assert 0 < area_range[0] <= area_range[1]
+        assert 1 <= max_attempts
+
+        self.min_covered = min_covered
+        self.aspect_ratio_range = aspect_ratio_range
+        self.area_range = area_range
+        self.max_attempts = max_attempts
+        self._fallback = EfficientNetCenterCrop(imgsize)
+
+    def __call__(self, img):
+        # https://github.com/tensorflow/tensorflow/blob/9274bcebb31322370139467039034f8ff852b004/tensorflow/core/kernels/sample_distorted_bounding_box_op.cc#L111
+        original_width, original_height = img.size
+        min_area = self.area_range[0] * (original_width * original_height)
+        max_area = self.area_range[1] * (original_width * original_height)
+
+        for _ in range(self.max_attempts):
+            aspect_ratio = random.uniform(*self.aspect_ratio_range)
+            height = int(round(math.sqrt(min_area / aspect_ratio)))
+            max_height = int(round(math.sqrt(max_area / aspect_ratio)))
+
+            if max_height * aspect_ratio > original_width:
+                max_height = (original_width + 0.5 - 1e-7) / aspect_ratio
+                max_height = int(max_height)
+                if max_height * aspect_ratio > original_width:
+                    max_height -= 1
+
+            if max_height > original_height:
+                max_height = original_height
+
+            if height >= max_height:
+                height = max_height
+
+            height = int(round(random.uniform(height, max_height)))
+            width = int(round(height * aspect_ratio))
+            area = width * height
+
+            if area < min_area or area > max_area:
+                continue
+            if width > original_width or height > original_height:
+                continue
+            if area < self.min_covered * (original_width * original_height):
+                continue
+            if width == original_width and height == original_height:
+                return self._fallback(img)      # https://github.com/tensorflow/tpu/blob/master/models/official/efficientnet/preprocessing.py#L102
+
+            x = random.randint(0, original_width - width)
+            y = random.randint(0, original_height - height)
+            return img.crop((x, y, x + width, y + height))
+
+        return self._fallback(img)
+
+
+class EfficientNetCenterCrop:
+    def __init__(self, imgsize):
+        self.imgsize = imgsize
+
+    def __call__(self, img):
+        """Crop the given PIL Image and resize it to desired size.
+
+        Args:
+            img (PIL Image): Image to be cropped. (0,0) denotes the top left corner of the image.
+            output_size (sequence or int): (height, width) of the crop box. If int,
+                it is used for both directions
+        Returns:
+            PIL Image: Cropped image.
+        """
+        image_width, image_height = img.size
+        image_short = min(image_width, image_height)
+
+        crop_size = float(self.imgsize) / (self.imgsize + 32) * image_short
+
+        crop_height, crop_width = crop_size, crop_size
+        crop_top = int(round((image_height - crop_height) / 2.))
+        crop_left = int(round((image_width - crop_width) / 2.))
+        return img.crop((crop_left, crop_top, crop_left + crop_width, crop_top + crop_height))
+
+
+class SubsetSampler(Sampler):
+    r"""Samples elements from a given list of indices, without replacement.
+
+    Arguments:
+        indices (sequence): a sequence of indices
+    """
+
+    def __init__(self, indices):
+        self.indices = indices
+
+    def __iter__(self):
+        return (i for i in self.indices)
+
+    def __len__(self):
+        return len(self.indices)

+from __future__ import print_function
+import os
+import shutil
+import torch
+
+ARCHIVE_DICT = {
+    'train': {
+        'url': 'http://www.image-net.org/challenges/LSVRC/2012/nnoupb/ILSVRC2012_img_train.tar',
+        'md5': '1d675b47d978889d74fa0da5fadfb00e',
+    },
+    'val': {
+        'url': 'http://www.image-net.org/challenges/LSVRC/2012/nnoupb/ILSVRC2012_img_val.tar',
+        'md5': '29b22e2961454d5413ddabcf34fc5622',
+    },
+    'devkit': {
+        'url': 'http://www.image-net.org/challenges/LSVRC/2012/nnoupb/ILSVRC2012_devkit_t12.tar.gz',
+        'md5': 'fa75699e90414af021442c21a62c3abf',
+    }
+}
+
+
+import torchvision
+from torchvision.datasets.utils import check_integrity, download_url
+
+
+# copy ILSVRC/ImageSets/CLS-LOC/train_cls.txt to ./root/
+# to skip os walk (it's too slow) using ILSVRC/ImageSets/CLS-LOC/train_cls.txt file
+class ImageNet(torchvision.datasets.ImageFolder):
+    """`ImageNet <http://image-net.org/>`_ 2012 Classification Dataset.
+
+    Args:
+        root (string): Root directory of the ImageNet Dataset.
+        split (string, optional): The dataset split, supports ``train``, or ``val``.
+        download (bool, optional): If true, downloads the dataset from the internet and
+            puts it in root directory. If dataset is already downloaded, it is not
+            downloaded again.
+        transform (callable, optional): A function/transform that  takes in an PIL image
+            and returns a transformed version. E.g, ``transforms.RandomCrop``
+        target_transform (callable, optional): A function/transform that takes in the
+            target and transforms it.
+        loader (callable, optional): A function to load an image given its path.
+
+     Attributes:
+        classes (list): List of the class names.
+        class_to_idx (dict): Dict with items (class_name, class_index).
+        wnids (list): List of the WordNet IDs.
+        wnid_to_idx (dict): Dict with items (wordnet_id, class_index).
+        imgs (list): List of (image path, class_index) tuples
+        targets (list): The class_index value for each image in the dataset
+    """
+
+    def __init__(self, root, split='train', download=False, **kwargs):
+        root = self.root = os.path.expanduser(root)
+        self.split = self._verify_split(split)
+
+        if download:
+            self.download()
+        wnid_to_classes = self._load_meta_file()[0]
+
+        # to skip os walk (it's too slow) using ILSVRC/ImageSets/CLS-LOC/train_cls.txt file
+        listfile = os.path.join(root, 'train_cls.txt')
+        if split == 'train' and os.path.exists(listfile):
+            torchvision.datasets.VisionDataset.__init__(self, root, **kwargs)
+            with open(listfile, 'r') as f:
+                datalist = [
+                    line.strip().split(' ')[0]
+                    for line in f.readlines()
+                    if line.strip()
+                ]
+
+            classes = list(set([line.split('/')[0] for line in datalist]))
+            classes.sort()
+            class_to_idx = {classes[i]: i for i in range(len(classes))}
+
+            samples = [
+                (os.path.join(self.split_folder, line + '.JPEG'), class_to_idx[line.split('/')[0]])
+                for line in datalist
+            ]
+
+            self.loader = torchvision.datasets.folder.default_loader
+            self.extensions = torchvision.datasets.folder.IMG_EXTENSIONS
+
+            self.classes = classes
+            self.class_to_idx = class_to_idx
+            self.samples = samples
+            self.targets = [s[1] for s in samples]
+
+            self.imgs = self.samples
+        else:
+            super(ImageNet, self).__init__(self.split_folder, **kwargs)
+
+        self.root = root
+
+        idcs = [idx for _, idx in self.imgs]
+        self.wnids = self.classes
+        self.wnid_to_idx = {wnid: idx for idx, wnid in zip(idcs, self.wnids)}
+        self.classes = [wnid_to_classes[wnid] for wnid in self.wnids]
+        self.class_to_idx = {cls: idx
+                             for clss, idx in zip(self.classes, idcs)
+                             for cls in clss}
+
+    def download(self):
+        if not check_integrity(self.meta_file):
+            tmpdir = os.path.join(self.root, 'tmp')
+
+            archive_dict = ARCHIVE_DICT['devkit']
+            download_and_extract_tar(archive_dict['url'], self.root,
+                                     extract_root=tmpdir,
+                                     md5=archive_dict['md5'])
+            devkit_folder = _splitexts(os.path.basename(archive_dict['url']))[0]
+            meta = parse_devkit(os.path.join(tmpdir, devkit_folder))
+            self._save_meta_file(*meta)
+
+            shutil.rmtree(tmpdir)
+
+        if not os.path.isdir(self.split_folder):
+            archive_dict = ARCHIVE_DICT[self.split]
+            download_and_extract_tar(archive_dict['url'], self.root,
+                                     extract_root=self.split_folder,
+                                     md5=archive_dict['md5'])
+
+            if self.split == 'train':
+                prepare_train_folder(self.split_folder)
+            elif self.split == 'val':
+                val_wnids = self._load_meta_file()[1]
+                prepare_val_folder(self.split_folder, val_wnids)
+        else:
+            msg = ("You set download=True, but a folder '{}' already exist in "
+                   "the root directory. If you want to re-download or re-extract the "
+                   "archive, delete the folder.")
+            print(msg.format(self.split))
+
+    @property
+    def meta_file(self):
+        return os.path.join(self.root, 'meta.bin')
+
+    def _load_meta_file(self):
+        if check_integrity(self.meta_file):
+            return torch.load(self.meta_file)
+        raise RuntimeError("Meta file not found or corrupted.",
+                           "You can use download=True to create it.")
+
+    def _save_meta_file(self, wnid_to_class, val_wnids):
+        torch.save((wnid_to_class, val_wnids), self.meta_file)
+
+    def _verify_split(self, split):
+        if split not in self.valid_splits:
+            msg = "Unknown split {} .".format(split)
+            msg += "Valid splits are {{}}.".format(", ".join(self.valid_splits))
+            raise ValueError(msg)
+        return split
+
+    @property
+    def valid_splits(self):
+        return 'train', 'val'
+
+    @property
+    def split_folder(self):
+        return os.path.join(self.root, self.split)
+
+    def extra_repr(self):
+        return "Split: {split}".format(**self.__dict__)
+
+
+def extract_tar(src, dest=None, gzip=None, delete=False):
+    import tarfile
+
+    if dest is None:
+        dest = os.path.dirname(src)
+    if gzip is None:
+        gzip = src.lower().endswith('.gz')
+
+    mode = 'r:gz' if gzip else 'r'
+    with tarfile.open(src, mode) as tarfh:
+        tarfh.extractall(path=dest)
+
+    if delete:
+        os.remove(src)
+
+
+def download_and_extract_tar(url, download_root, extract_root=None, filename=None,
+                             md5=None, **kwargs):
+    download_root = os.path.expanduser(download_root)
+    if extract_root is None:
+        extract_root = download_root
+    if filename is None:
+        filename = os.path.basename(url)
+
+    if not check_integrity(os.path.join(download_root, filename), md5):
+        download_url(url, download_root, filename=filename, md5=md5)
+
+    extract_tar(os.path.join(download_root, filename), extract_root, **kwargs)
+
+
+def parse_devkit(root):
+    idx_to_wnid, wnid_to_classes = parse_meta(root)
+    val_idcs = parse_val_groundtruth(root)
+    val_wnids = [idx_to_wnid[idx] for idx in val_idcs]
+    return wnid_to_classes, val_wnids
+
+
+def parse_meta(devkit_root, path='data', filename='meta.mat'):
+    import scipy.io as sio
+
+    metafile = os.path.join(devkit_root, path, filename)
+    meta = sio.loadmat(metafile, squeeze_me=True)['synsets']
+    nums_children = list(zip(*meta))[4]
+    meta = [meta[idx] for idx, num_children in enumerate(nums_children)
+            if num_children == 0]
+    idcs, wnids, classes = list(zip(*meta))[:3]
+    classes = [tuple(clss.split(', ')) for clss in classes]
+    idx_to_wnid = {idx: wnid for idx, wnid in zip(idcs, wnids)}
+    wnid_to_classes = {wnid: clss for wnid, clss in zip(wnids, classes)}
+    return idx_to_wnid, wnid_to_classes
+
+
+def parse_val_groundtruth(devkit_root, path='data',
+                          filename='ILSVRC2012_validation_ground_truth.txt'):
+    with open(os.path.join(devkit_root, path, filename), 'r') as txtfh:
+        val_idcs = txtfh.readlines()
+    return [int(val_idx) for val_idx in val_idcs]
+
+
+def prepare_train_folder(folder):
+    for archive in [os.path.join(folder, archive) for archive in os.listdir(folder)]:
+        extract_tar(archive, os.path.splitext(archive)[0], delete=True)
+
+
+def prepare_val_folder(folder, wnids):
+    img_files = sorted([os.path.join(folder, file) for file in os.listdir(folder)])
+
+    for wnid in set(wnids):
+        os.mkdir(os.path.join(folder, wnid))
+
+    for wnid, img_file in zip(wnids, img_files):
+        shutil.move(img_file, os.path.join(folder, wnid, os.path.basename(img_file)))
+
+
+def _splitexts(root):
+    exts = []
+    ext = '.'
+    while ext:
+        root, ext = os.path.splitext(root)
+        exts.append(ext)
+    return root, ''.join(reversed(exts))

+import torch
+
+from theconf import Config as C
+
+
+def adjust_learning_rate_resnet(optimizer):
+    """
+    Sets the learning rate to the initial LR decayed by 10 on every predefined epochs
+    Ref: AutoAugment
+    """
+
+    if C.get()['epoch'] == 90:
+        return torch.optim.lr_scheduler.MultiStepLR(optimizer, [30, 60, 80])
+    elif C.get()['epoch'] == 270:   # autoaugment
+        return torch.optim.lr_scheduler.MultiStepLR(optimizer, [90, 180, 240])
+    else:
+        raise ValueError('invalid epoch=%d for resnet scheduler' % C.get()['epoch'])

+import copy
+
+import torch
+import numpy as np
+from collections import defaultdict
+
+from torch import nn
+
+
+def accuracy(output, target, topk=(1,)):
+    """Computes the precision@k for the specified values of k"""
+    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)
+        res.append(correct_k.mul_(1. / batch_size))
+    return res
+
+
+class CrossEntropyLabelSmooth(torch.nn.Module):
+    def __init__(self, num_classes, epsilon, reduction='mean'):
+        super(CrossEntropyLabelSmooth, self).__init__()
+        self.num_classes = num_classes
+        self.epsilon = epsilon
+        self.reduction = reduction
+        self.logsoftmax = torch.nn.LogSoftmax(dim=1)
+
+    def forward(self, input, target):  # pylint: disable=redefined-builtin
+        log_probs = self.logsoftmax(input)
+        targets = torch.zeros_like(log_probs).scatter_(1, target.unsqueeze(1), 1)
+        if self.epsilon > 0.0:
+            targets = (1 - self.epsilon) * targets + self.epsilon / self.num_classes
+        targets = targets.detach()
+        loss = (-targets * log_probs)
+
+        if self.reduction in ['avg', 'mean']:
+            loss = torch.mean(torch.sum(loss, dim=1))
+        elif self.reduction == 'sum':
+            loss = loss.sum()
+        return loss
+
+
+class Accumulator:
+    def __init__(self):
+        self.metrics = defaultdict(lambda: 0.)
+
+    def add(self, key, value):
+        self.metrics[key] += value
+
+    def add_dict(self, dict):
+        for key, value in dict.items():
+            self.add(key, value)
+
+    def __getitem__(self, item):
+        return self.metrics[item]
+
+    def __setitem__(self, key, value):
+        self.metrics[key] = value
+
+    def get_dict(self):
+        return copy.deepcopy(dict(self.metrics))
+
+    def items(self):
+        return self.metrics.items()
+
+    def __str__(self):
+        return str(dict(self.metrics))
+
+    def __truediv__(self, other):
+        newone = Accumulator()
+        for key, value in self.items():
+            if isinstance(other, str):
+                if other != key:
+                    newone[key] = value / self[other]
+                else:
+                    newone[key] = value
+            else:
+                newone[key] = value / other
+        return newone
+
+
+class SummaryWriterDummy:
+    def __init__(self, log_dir):
+        pass
+
+    def add_scalar(self, *args, **kwargs):
+        pass

+import torch
+
+from torch import nn
+from torch.nn import DataParallel
+from torch.nn.parallel import DistributedDataParallel
+import torch.backends.cudnn as cudnn
+# from torchvision import models
+import numpy as np
+
+from FastAutoAugment.networks.resnet import ResNet
+from FastAutoAugment.networks.pyramidnet import PyramidNet
+from FastAutoAugment.networks.shakeshake.shake_resnet import ShakeResNet
+from FastAutoAugment.networks.wideresnet import WideResNet
+from FastAutoAugment.networks.shakeshake.shake_resnext import ShakeResNeXt
+from FastAutoAugment.networks.efficientnet_pytorch import EfficientNet, RoutingFn
+from FastAutoAugment.tf_port.tpu_bn import TpuBatchNormalization
+
+
+def get_model(conf, num_class=10, local_rank=-1):
+    name = conf['type']
+
+    if name == 'resnet50':
+        model = ResNet(dataset='imagenet', depth=50, num_classes=num_class, bottleneck=True)
+    elif name == 'resnet200':
+        model = ResNet(dataset='imagenet', depth=200, num_classes=num_class, bottleneck=True)
+    elif name == 'wresnet40_2':
+        model = WideResNet(40, 2, dropout_rate=0.0, num_classes=num_class)
+    elif name == 'wresnet28_10':
+        model = WideResNet(28, 10, dropout_rate=0.0, num_classes=num_class)
+
+    elif name == 'shakeshake26_2x32d':
+        model = ShakeResNet(26, 32, num_class)
+    elif name == 'shakeshake26_2x64d':
+        model = ShakeResNet(26, 64, num_class)
+    elif name == 'shakeshake26_2x96d':
+        model = ShakeResNet(26, 96, num_class)
+    elif name == 'shakeshake26_2x112d':
+        model = ShakeResNet(26, 112, num_class)
+
+    elif name == 'shakeshake26_2x96d_next':
+        model = ShakeResNeXt(26, 96, 4, num_class)
+
+    elif name == 'pyramid':
+        model = PyramidNet('cifar10', depth=conf['depth'], alpha=conf['alpha'], num_classes=num_class, bottleneck=conf['bottleneck'])
+
+    elif 'efficientnet' in name:
+        model = EfficientNet.from_name(name, condconv_num_expert=conf['condconv_num_expert'], norm_layer=None)  # TpuBatchNormalization
+        if local_rank >= 0:
+            model = nn.SyncBatchNorm.convert_sync_batchnorm(model)
+        def kernel_initializer(module):
+            def get_fan_in_out(module):
+                num_input_fmaps = module.weight.size(1)
+                num_output_fmaps = module.weight.size(0)
+                receptive_field_size = 1
+                if module.weight.dim() > 2:
+                    receptive_field_size = module.weight[0][0].numel()
+                fan_in = num_input_fmaps * receptive_field_size
+                fan_out = num_output_fmaps * receptive_field_size
+                return fan_in, fan_out
+
+            if isinstance(module, torch.nn.Conv2d):
+                # https://github.com/tensorflow/tpu/blob/master/models/official/efficientnet/efficientnet_model.py#L58
+                fan_in, fan_out = get_fan_in_out(module)
+                torch.nn.init.normal_(module.weight, mean=0.0, std=np.sqrt(2.0 / fan_out))
+                if module.bias is not None:
+                    torch.nn.init.constant_(module.bias, val=0.)
+            elif isinstance(module, RoutingFn):
+                torch.nn.init.xavier_uniform_(module.weight)
+                torch.nn.init.constant_(module.bias, val=0.)
+            elif isinstance(module, torch.nn.Linear):
+                # https://github.com/tensorflow/tpu/blob/master/models/official/efficientnet/efficientnet_model.py#L82
+                fan_in, fan_out = get_fan_in_out(module)
+                delta = 1.0 / np.sqrt(fan_out)
+                torch.nn.init.uniform_(module.weight, a=-delta, b=delta)
+                if module.bias is not None:
+                    torch.nn.init.constant_(module.bias, val=0.)
+        model.apply(kernel_initializer)
+    else:
+        raise NameError('no model named, %s' % name)
+
+    if local_rank >= 0:
+        device = torch.device('cuda', local_rank)
+        model = model.to(device)
+        model = DistributedDataParallel(model, device_ids=[local_rank], output_device=local_rank)
+    else:
+        model = model.cuda()
+#         model = DataParallel(model)
+
+    cudnn.benchmark = True
+    return model
+
+
+def num_class(dataset):
+    return {
+        'cifar10': 10,
+        'reduced_cifar10': 10,
+        'cifar10.1': 10,
+        'cifar100': 100,
+        'svhn': 10,
+        'reduced_svhn': 10,
+        'imagenet': 1000,
+        'reduced_imagenet': 120,
+    }[dataset]

+__version__ = "0.5.1"
+from .model import EfficientNet, RoutingFn
+from .utils import (
+    GlobalParams,
+    BlockArgs,
+    BlockDecoder,
+    efficientnet,
+    get_model_params,
+)
\ No newline at end of file

+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch._six import container_abcs
+
+from itertools import repeat
+from functools import partial
+from typing import Union, List, Tuple, Optional, Callable
+import numpy as np
+import math
+
+
+def _ntuple(n):
+    def parse(x):
+        if isinstance(x, container_abcs.Iterable):
+            return x
+        return tuple(repeat(x, n))
+    return parse
+
+
+_single = _ntuple(1)
+_pair = _ntuple(2)
+_triple = _ntuple(3)
+_quadruple = _ntuple(4)
+
+
+def _is_static_pad(kernel_size, stride=1, dilation=1, **_):
+    return stride == 1 and (dilation * (kernel_size - 1)) % 2 == 0
+
+
+def _get_padding(kernel_size, stride=1, dilation=1, **_):
+    padding = ((stride - 1) + dilation * (kernel_size - 1)) // 2
+    return padding
+
+
+def _calc_same_pad(i: int, k: int, s: int, d: int):
+    return max((math.ceil(i / s) - 1) * s + (k - 1) * d + 1 - i, 0)
+
+
+def conv2d_same(
+        x, weight: torch.Tensor, bias: Optional[torch.Tensor] = None, stride: Tuple[int, int] = (1, 1),
+        padding: Tuple[int, int] = (0, 0), dilation: Tuple[int, int] = (1, 1), groups: int = 1):
+    ih, iw = x.size()[-2:]
+    kh, kw = weight.size()[-2:]
+    pad_h = _calc_same_pad(ih, kh, stride[0], dilation[0])
+    pad_w = _calc_same_pad(iw, kw, stride[1], dilation[1])
+    if pad_h > 0 or pad_w > 0:
+        x = F.pad(x, [pad_w // 2, pad_w - pad_w // 2, pad_h // 2, pad_h - pad_h // 2])
+    return F.conv2d(x, weight, bias, stride, (0, 0), dilation, groups)
+
+
+def get_padding_value(padding, kernel_size, **kwargs):
+    dynamic = False
+    if isinstance(padding, str):
+        # for any string padding, the padding will be calculated for you, one of three ways
+        padding = padding.lower()
+        if padding == 'same':
+            # TF compatible 'SAME' padding, has a performance and GPU memory allocation impact
+            if _is_static_pad(kernel_size, **kwargs):
+                # static case, no extra overhead
+                padding = _get_padding(kernel_size, **kwargs)
+            else:
+                # dynamic padding
+                padding = 0
+                dynamic = True
+        elif padding == 'valid':
+            # 'VALID' padding, same as padding=0
+            padding = 0
+        else:
+            # Default to PyTorch style 'same'-ish symmetric padding
+            padding = _get_padding(kernel_size, **kwargs)
+    return padding, dynamic
+
+
+def get_condconv_initializer(initializer, num_experts, expert_shape):
+    def condconv_initializer(weight):
+        """CondConv initializer function."""
+        num_params = np.prod(expert_shape)
+        if (len(weight.shape) != 2 or weight.shape[0] != num_experts or weight.shape[1] != num_params):
+            raise (ValueError('CondConv variables must have shape [num_experts, num_params]'))
+        for i in range(num_experts):
+            initializer(weight[i].view(expert_shape))
+    return condconv_initializer
+
+
+class CondConv2d(nn.Module):
+    """ Conditional Convolution
+    Inspired by: https://github.com/tensorflow/tpu/blob/master/models/official/efficientnet/condconv/condconv_layers.py
+    Grouped convolution hackery for parallel execution of the per-sample kernel filters inspired by this discussion:
+    https://github.com/pytorch/pytorch/issues/17983
+    """
+    __constants__ = ['bias', 'in_channels', 'out_channels', 'dynamic_padding']
+
+    def __init__(self, in_channels, out_channels, kernel_size=3,
+                 stride=1, padding='', dilation=1, groups=1, bias=False, num_experts=4):
+        super(CondConv2d, self).__init__()
+        assert num_experts > 1
+
+        if isinstance(stride, container_abcs.Iterable) and len(stride) == 1:
+            stride = stride[0]
+        # print('CondConv', num_experts)
+
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+        self.kernel_size = _pair(kernel_size)
+        self.stride = _pair(stride)
+        padding_val, is_padding_dynamic = get_padding_value(padding, kernel_size, stride=stride, dilation=dilation)
+        self.dynamic_padding = is_padding_dynamic  # if in forward to work with torchscript
+        self.padding = _pair(padding_val)
+        self.dilation = _pair(dilation)
+        self.groups = groups
+        self.num_experts = num_experts
+
+        self.weight_shape = (self.out_channels, self.in_channels // self.groups) + self.kernel_size
+        weight_num_param = 1
+        for wd in self.weight_shape:
+            weight_num_param *= wd
+        self.weight = torch.nn.Parameter(torch.Tensor(self.num_experts, weight_num_param))
+
+        if bias:
+            self.bias_shape = (self.out_channels,)
+            self.bias = torch.nn.Parameter(torch.Tensor(self.num_experts, self.out_channels))
+        else:
+            self.register_parameter('bias', None)
+
+        self.reset_parameters()
+
+    def reset_parameters(self):
+        num_input_fmaps = self.weight.size(1)
+        num_output_fmaps = self.weight.size(0)
+        receptive_field_size = 1
+        if self.weight.dim() > 2:
+            receptive_field_size = self.weight[0][0].numel()
+        fan_in = num_input_fmaps * receptive_field_size
+        fan_out = num_output_fmaps * receptive_field_size
+
+        init_weight = get_condconv_initializer(partial(nn.init.normal_, mean=0.0, std=np.sqrt(2.0 / fan_out)), self.num_experts, self.weight_shape)
+        init_weight(self.weight)
+        if self.bias is not None:
+            # fan_in = np.prod(self.weight_shape[1:])
+            # bound = 1 / math.sqrt(fan_in)
+            init_bias = get_condconv_initializer(partial(nn.init.constant_, val=0), self.num_experts, self.bias_shape)
+            init_bias(self.bias)
+
+    def forward(self, x, routing_weights):
+        x_orig = x
+        B, C, H, W = x.shape
+        weight = torch.matmul(routing_weights, self.weight)     # (Expert x out x in x 3x3) --> (B x out x in x 3x3)
+        new_weight_shape = (B * self.out_channels, self.in_channels // self.groups) + self.kernel_size
+        weight = weight.view(new_weight_shape)                  # (B*out x in x 3 x 3)
+        bias = None
+        if self.bias is not None:
+            bias = torch.matmul(routing_weights, self.bias)
+            bias = bias.view(B * self.out_channels)
+        # move batch elements with channels so each batch element can be efficiently convolved with separate kernel
+        x = x.view(1, B * C, H, W)
+        if self.dynamic_padding:
+            out = conv2d_same(
+                x, weight, bias, stride=self.stride, padding=self.padding,
+                dilation=self.dilation, groups=self.groups * B)
+        else:
+            out = F.conv2d(
+                x, weight, bias, stride=self.stride, padding=self.padding,
+                dilation=self.dilation, groups=self.groups * B)
+
+        # out : (1 x B*out x ...)
+        out = out.permute([1, 0, 2, 3]).view(B, self.out_channels, out.shape[-2], out.shape[-1])
+
+        # out2 = self.forward_legacy(x_orig, routing_weights)
+        # lt = torch.lt(torch.abs(torch.add(out, -out2)), 1e-8)
+        # assert torch.all(lt), torch.abs(torch.add(out, -out2))[lt]
+        # print('checked')
+        return out
+
+    def forward_legacy(self, x, routing_weights):
+        # Literal port (from TF definition)
+        B, C, H, W = x.shape
+        weight = torch.matmul(routing_weights, self.weight)  # (Expert x out x in x 3x3) --> (B x out x in x 3x3)
+        x = torch.split(x, 1, 0)
+        weight = torch.split(weight, 1, 0)
+        if self.bias is not None:
+            bias = torch.matmul(routing_weights, self.bias)
+            bias = torch.split(bias, 1, 0)
+        else:
+            bias = [None] * B
+        out = []
+        if self.dynamic_padding:
+            conv_fn = conv2d_same
+        else:
+            conv_fn = F.conv2d
+        for xi, wi, bi in zip(x, weight, bias):
+            wi = wi.view(*self.weight_shape)
+            if bi is not None:
+                bi = bi.view(*self.bias_shape)
+            out.append(conv_fn(
+                xi, wi, bi, stride=self.stride, padding=self.padding,
+                dilation=self.dilation, groups=self.groups))
+        out = torch.cat(out, 0)
+        return out

+import torch
+from torch import nn
+from torch.nn import functional as F
+
+from functools import partial
+from .utils import (
+    round_filters,
+    round_repeats,
+    drop_connect,
+    get_same_padding_conv2d,
+    get_model_params,
+    efficientnet_params,
+    load_pretrained_weights,
+    MemoryEfficientSwish,
+)
+
+
+class RoutingFn(nn.Linear):
+    pass
+
+
+class MBConvBlock(nn.Module):
+    """
+    Mobile Inverted Residual Bottleneck Block
+
+    Args:
+        block_args (namedtuple): BlockArgs, see above
+        global_params (namedtuple): GlobalParam, see above
+
+    Attributes:
+        has_se (bool): Whether the block contains a Squeeze and Excitation layer.
+    """
+
+    def __init__(self, block_args, global_params, norm_layer=None):
+        super().__init__()
+        self._block_args = block_args
+        self._bn_mom = 1 - global_params.batch_norm_momentum
+        self._bn_eps = global_params.batch_norm_epsilon
+        self.has_se = (self._block_args.se_ratio is not None) and (0 < self._block_args.se_ratio <= 1)
+        self.id_skip = block_args.id_skip  # skip connection and drop connect
+        if norm_layer is None:
+            norm_layer = nn.BatchNorm2d
+
+        self.condconv_num_expert = block_args.condconv_num_expert
+        if self._is_condconv():
+            self.routing_fn = RoutingFn(self._block_args.input_filters, self.condconv_num_expert)
+
+        # Get static or dynamic convolution depending on image size
+        Conv2d = get_same_padding_conv2d(image_size=global_params.image_size, condconv_num_expert=block_args.condconv_num_expert)
+        Conv2dse = get_same_padding_conv2d(image_size=global_params.image_size)
+
+        # Expansion phase
+        inp = self._block_args.input_filters  # number of input channels
+        oup = self._block_args.input_filters * self._block_args.expand_ratio  # number of output channels
+        if self._block_args.expand_ratio != 1:
+            self._expand_conv = Conv2d(in_channels=inp, out_channels=oup, kernel_size=1, bias=False)
+            self._bn0 = norm_layer(num_features=oup, momentum=self._bn_mom, eps=self._bn_eps)
+
+        # Depthwise convolution phase
+        k = self._block_args.kernel_size
+        s = self._block_args.stride
+        self._depthwise_conv = Conv2d(
+            in_channels=oup, out_channels=oup, groups=oup,  # groups makes it depthwise
+            kernel_size=k, stride=s, bias=False)
+        self._bn1 = norm_layer(num_features=oup, momentum=self._bn_mom, eps=self._bn_eps)
+
+        # Squeeze and Excitation layer, if desired
+        if self.has_se:
+            num_squeezed_channels = max(1, int(self._block_args.input_filters * self._block_args.se_ratio))
+            self._se_reduce = Conv2dse(in_channels=oup, out_channels=num_squeezed_channels, kernel_size=1)
+            self._se_expand = Conv2dse(in_channels=num_squeezed_channels, out_channels=oup, kernel_size=1)
+
+        # Output phase
+        final_oup = self._block_args.output_filters
+        self._project_conv = Conv2d(in_channels=oup, out_channels=final_oup, kernel_size=1, bias=False)
+        self._bn2 = norm_layer(num_features=final_oup, momentum=self._bn_mom, eps=self._bn_eps)
+        self._swish = MemoryEfficientSwish()
+
+    def _is_condconv(self):
+        return self.condconv_num_expert > 1
+
+    def forward(self, inputs, drop_connect_rate=None):
+        """
+        :param inputs: input tensor
+        :param drop_connect_rate: drop connect rate (float, between 0 and 1)
+        :return: output of block
+        """
+
+        if self._is_condconv():
+            feat = F.adaptive_avg_pool2d(inputs, 1).flatten(1)
+            routing_w = torch.sigmoid(self.routing_fn(feat))
+
+            if self._block_args.expand_ratio != 1:
+                _expand_conv = partial(self._expand_conv, routing_weights=routing_w)
+            _depthwise_conv = partial(self._depthwise_conv, routing_weights=routing_w)
+            _project_conv = partial(self._project_conv, routing_weights=routing_w)
+        else:
+            if self._block_args.expand_ratio != 1:
+                _expand_conv = self._expand_conv
+            _depthwise_conv, _project_conv = self._depthwise_conv, self._project_conv
+
+        # Expansion and Depthwise Convolution
+        x = inputs
+        if self._block_args.expand_ratio != 1:
+            x = self._swish(self._bn0(_expand_conv(inputs)))
+        x = self._swish(self._bn1(_depthwise_conv(x)))
+
+        # Squeeze and Excitation
+        if self.has_se:
+            x_squeezed = F.adaptive_avg_pool2d(x, 1)
+            x_squeezed = self._se_expand(self._swish(self._se_reduce(x_squeezed)))
+            x = torch.sigmoid(x_squeezed) * x
+
+        x = self._bn2(_project_conv(x))
+
+        # Skip connection and drop connect
+        input_filters, output_filters = self._block_args.input_filters, self._block_args.output_filters
+        if self.id_skip and self._block_args.stride == 1 and input_filters == output_filters:
+            if drop_connect_rate:
+                x = drop_connect(x, drop_p=drop_connect_rate, training=self.training)
+            x = x + inputs  # skip connection
+        return x
+
+    def set_swish(self):
+        """Sets swish function as memory efficient (for training) or standard (for export)"""
+        self._swish = MemoryEfficientSwish()
+
+
+class EfficientNet(nn.Module):
+    """
+    An EfficientNet model. Most easily loaded with the .from_name or .from_pretrained methods
+
+    Args:
+        blocks_args (list): A list of BlockArgs to construct blocks
+        global_params (namedtuple): A set of GlobalParams shared between blocks
+
+    Example:
+        model = EfficientNet.from_pretrained('efficientnet-b0')
+
+    """
+
+    def __init__(self, blocks_args=None, global_params=None, norm_layer=None):
+        super().__init__()
+        assert isinstance(blocks_args, list), 'blocks_args should be a list'
+        assert len(blocks_args) > 0, 'block args must be greater than 0'
+        self._global_params = global_params
+        self._blocks_args = blocks_args
+        if norm_layer is None:
+            norm_layer = nn.BatchNorm2d
+
+        # Get static or dynamic convolution depending on image size
+        Conv2d = get_same_padding_conv2d(image_size=global_params.image_size)
+
+        # Batch norm parameters
+        bn_mom = 1 - self._global_params.batch_norm_momentum
+        bn_eps = self._global_params.batch_norm_epsilon
+
+        # Stem
+        in_channels = 3  # rgb
+        out_channels = round_filters(32, self._global_params)  # number of output channels
+        self._conv_stem = Conv2d(in_channels, out_channels, kernel_size=3, stride=2, bias=False)
+        self._bn0 = norm_layer(num_features=out_channels, momentum=bn_mom, eps=bn_eps)
+
+        # Build blocks
+        self._blocks = nn.ModuleList([])
+        for idx, block_args in enumerate(self._blocks_args):
+            # Update block input and output filters based on depth multiplier.
+            block_args = block_args._replace(
+                input_filters=round_filters(block_args.input_filters, self._global_params),
+                output_filters=round_filters(block_args.output_filters, self._global_params),
+                num_repeat=round_repeats(block_args.num_repeat, self._global_params)
+            )
+
+            # The first block needs to take care of stride and filter size increase.
+            self._blocks.append(MBConvBlock(block_args, self._global_params, norm_layer=norm_layer))
+            if block_args.num_repeat > 1:
+                block_args = block_args._replace(input_filters=block_args.output_filters, stride=1)
+            for _ in range(block_args.num_repeat - 1):
+                self._blocks.append(MBConvBlock(block_args, self._global_params, norm_layer=norm_layer))
+
+        # Head
+        in_channels = block_args.output_filters  # output of final block
+        out_channels = round_filters(1280, self._global_params)
+        self._conv_head = Conv2d(in_channels, out_channels, kernel_size=1, bias=False)
+        self._bn1 = norm_layer(num_features=out_channels, momentum=bn_mom, eps=bn_eps)
+
+        # Final linear layer
+        self._avg_pooling = nn.AdaptiveAvgPool2d(1)
+        self._dropout = nn.Dropout(self._global_params.dropout_rate)
+        self._fc = nn.Linear(out_channels, self._global_params.num_classes)
+        self._swish = MemoryEfficientSwish()
+
+    def set_swish(self):
+        """Sets swish function as memory efficient (for training) or standard (for export)"""
+        self._swish = MemoryEfficientSwish()
+        for block in self._blocks:
+            block.set_swish()
+
+    def extract_features(self, inputs):
+        """ Returns output of the final convolution layer """
+
+        # Stem
+        x = self._swish(self._bn0(self._conv_stem(inputs)))
+
+        # Blocks
+        for idx, block in enumerate(self._blocks):
+            drop_connect_rate = self._global_params.drop_connect_rate
+            if drop_connect_rate:
+                drop_connect_rate *= float(idx) / len(self._blocks)
+            x = block(x, drop_connect_rate=drop_connect_rate)
+
+        # Head
+        x = self._swish(self._bn1(self._conv_head(x)))
+
+        return x
+
+    def forward(self, inputs):
+        """ Calls extract_features to extract features, applies final linear layer, and returns logits. """
+        bs = inputs.size(0)
+        # Convolution layers
+        x = self.extract_features(inputs)
+
+        # Pooling and final linear layer
+        x = self._avg_pooling(x)
+        x = x.view(bs, -1)
+        x = self._dropout(x)
+        x = self._fc(x)
+        return x
+
+    @classmethod
+    def from_name(cls, model_name, override_params=None, norm_layer=None, condconv_num_expert=1):
+        cls._check_model_name_is_valid(model_name)
+        blocks_args, global_params = get_model_params(model_name, override_params, condconv_num_expert=condconv_num_expert)
+        return cls(blocks_args, global_params, norm_layer=norm_layer)
+
+    @classmethod
+    def from_pretrained(cls, model_name, num_classes=1000):
+        model = cls.from_name(model_name, override_params={'num_classes': num_classes})
+        load_pretrained_weights(model, model_name, load_fc=(num_classes == 1000))
+
+        return model
+
+    @classmethod
+    def get_image_size(cls, model_name):
+        cls._check_model_name_is_valid(model_name)
+        _, _, res, _ = efficientnet_params(model_name)
+        return res
+
+    @classmethod
+    def _check_model_name_is_valid(cls, model_name, also_need_pretrained_weights=False):
+        """ Validates model name. None that pretrained weights are only available for
+        the first four models (efficientnet-b{i} for i in 0,1,2,3) at the moment. """
+        num_models = 4 if also_need_pretrained_weights else 8
+        valid_models = ['efficientnet-b'+str(i) for i in range(num_models)]
+        if model_name not in valid_models:
+            raise ValueError(f'model_name={model_name} should be one of: ' + ', '.join(valid_models))

+"""
+This file contains helper functions for building the model and for loading model parameters.
+These helper functions are built to mirror those in the official TensorFlow implementation.
+"""
+
+import re
+import math
+import collections
+from functools import partial
+import torch
+from torch import nn
+from torch.nn import functional as F
+from torch.utils import model_zoo
+
+########################################################################
+############### HELPERS FUNCTIONS FOR MODEL ARCHITECTURE ###############
+########################################################################
+
+
+# Parameters for the entire model (stem, all blocks, and head)
+from FastAutoAugment.networks.efficientnet_pytorch.condconv import CondConv2d
+
+GlobalParams = collections.namedtuple('GlobalParams', [
+    'batch_norm_momentum', 'batch_norm_epsilon', 'dropout_rate',
+    'num_classes', 'width_coefficient', 'depth_coefficient',
+    'depth_divisor', 'min_depth', 'drop_connect_rate', 'image_size'])
+
+# Parameters for an individual model block
+BlockArgs = collections.namedtuple('BlockArgs', [
+    'kernel_size', 'num_repeat', 'input_filters', 'output_filters',
+    'expand_ratio', 'id_skip', 'stride', 'se_ratio', 'condconv_num_expert'])
+
+# Change namedtuple defaults
+GlobalParams.__new__.__defaults__ = (None,) * len(GlobalParams._fields)
+BlockArgs.__new__.__defaults__ = (None,) * len(BlockArgs._fields)
+
+
+class SwishImplementation(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, i):
+        result = i * torch.sigmoid(i)
+        ctx.save_for_backward(i)
+        return result
+
+    @staticmethod
+    def backward(ctx, grad_output):
+        i = ctx.saved_tensors[0]
+        sigmoid_i = torch.sigmoid(i)
+        return grad_output * (sigmoid_i * (1 + i * (1 - sigmoid_i)))
+
+
+class MemoryEfficientSwish(nn.Module):
+    def forward(self, x):
+        return SwishImplementation.apply(x)
+
+
+def round_filters(filters, global_params):
+    """ Calculate and round number of filters based on depth multiplier. """
+    multiplier = global_params.width_coefficient
+    if not multiplier:
+        return filters
+    divisor = global_params.depth_divisor
+    min_depth = global_params.min_depth
+    filters *= multiplier
+    min_depth = min_depth or divisor
+    new_filters = max(min_depth, int(filters + divisor / 2) // divisor * divisor)
+    if new_filters < 0.9 * filters:  # prevent rounding by more than 10%
+        new_filters += divisor
+    return int(new_filters)
+
+
+def round_repeats(repeats, global_params):
+    """ Round number of filters based on depth multiplier. """
+    multiplier = global_params.depth_coefficient
+    if not multiplier:
+        return repeats
+    return int(math.ceil(multiplier * repeats))
+
+
+def drop_connect(inputs, drop_p, training):
+    """ Drop connect. """
+    if not training:
+        return inputs * (1. - drop_p)
+    batch_size = inputs.shape[0]
+    random_tensor = torch.rand([batch_size, 1, 1, 1], dtype=inputs.dtype, device=inputs.device)
+    binary_tensor = random_tensor > drop_p
+    output = inputs * binary_tensor.float()
+    # output = inputs / (1. - drop_p) * binary_tensor.float()
+    return output
+
+    # if not training: return inputs
+    # batch_size = inputs.shape[0]
+    # keep_prob = 1 - drop_p
+    # random_tensor = keep_prob
+    # random_tensor += torch.rand([batch_size, 1, 1, 1], dtype=inputs.dtype, device=inputs.device)
+    # binary_tensor = torch.floor(random_tensor)
+    # output = inputs / keep_prob * binary_tensor
+    # return output
+
+
+def get_same_padding_conv2d(image_size=None, condconv_num_expert=1):
+    """ Chooses static padding if you have specified an image size, and dynamic padding otherwise.
+        Static padding is necessary for ONNX exporting of models. """
+    if condconv_num_expert > 1:
+        return partial(CondConv2d, num_experts=condconv_num_expert)
+    elif image_size is None:
+        return Conv2dDynamicSamePadding
+    else:
+        return partial(Conv2dStaticSamePadding, image_size=image_size)
+
+
+class Conv2dDynamicSamePadding(nn.Conv2d):
+    """ 2D Convolutions like TensorFlow, for a dynamic image size """
+
+    def __init__(self, in_channels, out_channels, kernel_size, stride=1, dilation=1, groups=1, bias=True):
+        super().__init__(in_channels, out_channels, kernel_size, stride, 0, dilation, groups, bias)
+        self.stride = self.stride if len(self.stride) == 2 else [self.stride[0]] * 2
+
+    def forward(self, x):
+        ih, iw = x.size()[-2:]
+        kh, kw = self.weight.size()[-2:]
+        sh, sw = self.stride
+        oh, ow = math.ceil(ih / sh), math.ceil(iw / sw)
+        pad_h = max((oh - 1) * self.stride[0] + (kh - 1) * self.dilation[0] + 1 - ih, 0)
+        pad_w = max((ow - 1) * self.stride[1] + (kw - 1) * self.dilation[1] + 1 - iw, 0)
+        if pad_h > 0 or pad_w > 0:
+            x = F.pad(x, [pad_w // 2, pad_w - pad_w // 2, pad_h // 2, pad_h - pad_h // 2])
+        return F.conv2d(x, self.weight, self.bias, self.stride, self.padding, self.dilation, self.groups)
+
+
+class Conv2dStaticSamePadding(nn.Conv2d):
+    """ 2D Convolutions like TensorFlow, for a fixed image size"""
+
+    def __init__(self, in_channels, out_channels, kernel_size, image_size=None, **kwargs):
+        super().__init__(in_channels, out_channels, kernel_size, **kwargs)
+        self.stride = self.stride if len(self.stride) == 2 else [self.stride[0]] * 2
+
+        # Calculate padding based on image size and save it
+        assert image_size is not None
+        ih, iw = image_size if type(image_size) == list else [image_size, image_size]
+        kh, kw = self.weight.size()[-2:]
+        sh, sw = self.stride
+        oh, ow = math.ceil(ih / sh), math.ceil(iw / sw)
+        pad_h = max((oh - 1) * self.stride[0] + (kh - 1) * self.dilation[0] + 1 - ih, 0)
+        pad_w = max((ow - 1) * self.stride[1] + (kw - 1) * self.dilation[1] + 1 - iw, 0)
+        if pad_h > 0 or pad_w > 0:
+            self.static_padding = nn.ZeroPad2d((pad_w // 2, pad_w - pad_w // 2, pad_h // 2, pad_h - pad_h // 2))
+        else:
+            self.static_padding = Identity()
+
+    def forward(self, x):
+        x = self.static_padding(x)
+        x = F.conv2d(x, self.weight, self.bias, self.stride, self.padding, self.dilation, self.groups)
+        return x
+
+
+class Identity(nn.Module):
+    def __init__(self, ):
+        super(Identity, self).__init__()
+
+    def forward(self, input):
+        return input
+
+
+########################################################################
+############## HELPERS FUNCTIONS FOR LOADING MODEL PARAMS ##############
+########################################################################
+
+
+def efficientnet_params(model_name):
+    """ Map EfficientNet model name to parameter coefficients. """
+    params_dict = {
+        # Coefficients:   width,depth,res,dropout
+        'efficientnet-b0': (1.0, 1.0, 224, 0.2),
+        'efficientnet-b1': (1.0, 1.1, 240, 0.2),
+        'efficientnet-b2': (1.1, 1.2, 260, 0.3),
+        'efficientnet-b3': (1.2, 1.4, 300, 0.3),
+        'efficientnet-b4': (1.4, 1.8, 380, 0.4),
+        'efficientnet-b5': (1.6, 2.2, 456, 0.4),
+        'efficientnet-b6': (1.8, 2.6, 528, 0.5),
+        'efficientnet-b7': (2.0, 3.1, 600, 0.5),
+    }
+    return params_dict[model_name]
+
+
+class BlockDecoder(object):
+    """ Block Decoder for readability, straight from the official TensorFlow repository """
+
+    @staticmethod
+    def _decode_block_string(block_string):
+        """ Gets a block through a string notation of arguments. """
+        assert isinstance(block_string, str)
+
+        ops = block_string.split('_')
+        options = {}
+        for op in ops:
+            splits = re.split(r'(\d.*)', op)
+            if len(splits) >= 2:
+                key, value = splits[:2]
+                options[key] = value
+
+        # Check stride
+        assert (('s' in options and len(options['s']) == 1) or
+                (len(options['s']) == 2 and options['s'][0] == options['s'][1]))
+
+        return BlockArgs(
+            kernel_size=int(options['k']),
+            num_repeat=int(options['r']),
+            input_filters=int(options['i']),
+            output_filters=int(options['o']),
+            expand_ratio=int(options['e']),
+            id_skip=('noskip' not in block_string),
+            se_ratio=float(options['se']) if 'se' in options else None,
+            stride=[int(options['s'][0])],
+            condconv_num_expert=0
+        )
+
+    @staticmethod
+    def _encode_block_string(block):
+        """Encodes a block to a string."""
+        args = [
+            'r%d' % block.num_repeat,
+            'k%d' % block.kernel_size,
+            's%d%d' % (block.strides[0], block.strides[1]),
+            'e%s' % block.expand_ratio,
+            'i%d' % block.input_filters,
+            'o%d' % block.output_filters
+        ]
+        if 0 < block.se_ratio <= 1:
+            args.append('se%s' % block.se_ratio)
+        if block.id_skip is False:
+            args.append('noskip')
+        return '_'.join(args)
+
+    @staticmethod
+    def decode(string_list):
+        """
+        Decodes a list of string notations to specify blocks inside the network.
+
+        :param string_list: a list of strings, each string is a notation of block
+        :return: a list of BlockArgs namedtuples of block args
+        """
+        assert isinstance(string_list, list)
+        blocks_args = []
+        for block_string in string_list:
+            blocks_args.append(BlockDecoder._decode_block_string(block_string))
+        return blocks_args
+
+    @staticmethod
+    def encode(blocks_args):
+        """
+        Encodes a list of BlockArgs to a list of strings.
+
+        :param blocks_args: a list of BlockArgs namedtuples of block args
+        :return: a list of strings, each string is a notation of block
+        """
+        block_strings = []
+        for block in blocks_args:
+            block_strings.append(BlockDecoder._encode_block_string(block))
+        return block_strings
+
+
+def efficientnet(width_coefficient=None, depth_coefficient=None, dropout_rate=0.2,
+                 drop_connect_rate=0.2, image_size=None, num_classes=1000, condconv_num_expert=1):
+    """ Creates a efficientnet model. """
+
+    blocks_args = [
+        'r1_k3_s11_e1_i32_o16_se0.25', 'r2_k3_s22_e6_i16_o24_se0.25',
+        'r2_k5_s22_e6_i24_o40_se0.25', 'r3_k3_s22_e6_i40_o80_se0.25',
+        'r3_k5_s11_e6_i80_o112_se0.25', 'r4_k5_s22_e6_i112_o192_se0.25',
+        'r1_k3_s11_e6_i192_o320_se0.25',
+    ]
+    blocks_args = BlockDecoder.decode(blocks_args)
+
+    blocks_args_new = blocks_args[:-3]
+    for blocks_arg in blocks_args[-3:]:
+        blocks_arg = blocks_arg._replace(condconv_num_expert=condconv_num_expert)
+        blocks_args_new.append(blocks_arg)
+    blocks_args = blocks_args_new
+
+    global_params = GlobalParams(
+        batch_norm_momentum=0.99,
+        batch_norm_epsilon=1e-3,
+        dropout_rate=dropout_rate,
+        drop_connect_rate=drop_connect_rate,
+        # data_format='channels_last',  # removed, this is always true in PyTorch
+        num_classes=num_classes,
+        width_coefficient=width_coefficient,
+        depth_coefficient=depth_coefficient,
+        depth_divisor=8,
+        min_depth=None,
+        image_size=image_size,
+    )
+
+    return blocks_args, global_params
+
+
+def get_model_params(model_name, override_params, condconv_num_expert=1):
+    """ Get the block args and global params for a given model """
+    if model_name.startswith('efficientnet'):
+        w, d, s, p = efficientnet_params(model_name)
+        # note: all models have drop connect rate = 0.2
+        blocks_args, global_params = efficientnet(
+            width_coefficient=w, depth_coefficient=d, dropout_rate=p, image_size=s, condconv_num_expert=condconv_num_expert)
+    else:
+        raise NotImplementedError('model name is not pre-defined: %s' % model_name)
+    if override_params:
+        # ValueError will be raised here if override_params has fields not included in global_params.
+        global_params = global_params._replace(**override_params)
+    return blocks_args, global_params
+
+
+url_map = {
+    'efficientnet-b0': 'http://storage.googleapis.com/public-models/efficientnet/efficientnet-b0-355c32eb.pth',
+    'efficientnet-b1': 'http://storage.googleapis.com/public-models/efficientnet/efficientnet-b1-f1951068.pth',
+    'efficientnet-b2': 'http://storage.googleapis.com/public-models/efficientnet/efficientnet-b2-8bb594d6.pth',
+    'efficientnet-b3': 'http://storage.googleapis.com/public-models/efficientnet/efficientnet-b3-5fb5a3c3.pth',
+    'efficientnet-b4': 'http://storage.googleapis.com/public-models/efficientnet/efficientnet-b4-6ed6700e.pth',
+    'efficientnet-b5': 'http://storage.googleapis.com/public-models/efficientnet/efficientnet-b5-b6417697.pth',
+    'efficientnet-b6': 'http://storage.googleapis.com/public-models/efficientnet/efficientnet-b6-c76e70fd.pth',
+    'efficientnet-b7': 'http://storage.googleapis.com/public-models/efficientnet/efficientnet-b7-dcc49843.pth',
+}
+
+
+def load_pretrained_weights(model, model_name, load_fc=True):
+    """ Loads pretrained weights, and downloads if loading for the first time. """
+    state_dict = model_zoo.load_url(url_map[model_name])
+    if load_fc:
+        model.load_state_dict(state_dict)
+    else:
+        state_dict.pop('_fc.weight')
+        state_dict.pop('_fc.bias')
+        res = model.load_state_dict(state_dict, strict=False)
+        assert set(res.missing_keys) == set(['_fc.weight', '_fc.bias']), 'issue loading pretrained weights'
+    print('Loaded pretrained weights for {}'.format(model_name))

+import torch
+import torch.nn as nn
+import math
+
+from FastAutoAugment.networks.shakedrop import ShakeDrop
+
+
+def conv3x3(in_planes, out_planes, stride=1):
+    """
+    3x3 convolution with padding
+    """
+    return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride, padding=1, bias=False)
+
+
+class BasicBlock(nn.Module):
+    outchannel_ratio = 1
+
+    def __init__(self, inplanes, planes, stride=1, downsample=None, p_shakedrop=1.0):
+        super(BasicBlock, self).__init__()
+        self.bn1 = nn.BatchNorm2d(inplanes)
+        self.conv1 = conv3x3(inplanes, planes, stride)
+        self.bn2 = nn.BatchNorm2d(planes)
+        self.conv2 = conv3x3(planes, planes)
+        self.bn3 = nn.BatchNorm2d(planes)
+        self.relu = nn.ReLU(inplace=True)
+        self.downsample = downsample
+        self.stride = stride
+        self.shake_drop = ShakeDrop(p_shakedrop)
+
+    def forward(self, x):
+
+        out = self.bn1(x)
+        out = self.conv1(out)
+        out = self.bn2(out)
+        out = self.relu(out)
+        out = self.conv2(out)
+        out = self.bn3(out)
+
+        out = self.shake_drop(out)
+
+        if self.downsample is not None:
+            shortcut = self.downsample(x)
+            featuremap_size = shortcut.size()[2:4]
+        else:
+            shortcut = x
+            featuremap_size = out.size()[2:4]
+
+        batch_size = out.size()[0]
+        residual_channel = out.size()[1]
+        shortcut_channel = shortcut.size()[1]
+
+        if residual_channel != shortcut_channel:
+            padding = torch.autograd.Variable(
+                torch.cuda.FloatTensor(batch_size, residual_channel - shortcut_channel, featuremap_size[0],
+                                       featuremap_size[1]).fill_(0))
+            out += torch.cat((shortcut, padding), 1)
+        else:
+            out += shortcut
+
+        return out
+
+
+class Bottleneck(nn.Module):
+    outchannel_ratio = 4
+
+    def __init__(self, inplanes, planes, stride=1, downsample=None, p_shakedrop=1.0):
+        super(Bottleneck, self).__init__()
+        self.bn1 = nn.BatchNorm2d(inplanes)
+        self.conv1 = nn.Conv2d(inplanes, planes, kernel_size=1, bias=False)
+        self.bn2 = nn.BatchNorm2d(planes)
+        self.conv2 = nn.Conv2d(planes, (planes * 1), kernel_size=3, stride=stride,
+                               padding=1, bias=False)
+        self.bn3 = nn.BatchNorm2d((planes * 1))
+        self.conv3 = nn.Conv2d((planes * 1), planes * Bottleneck.outchannel_ratio, kernel_size=1, bias=False)
+        self.bn4 = nn.BatchNorm2d(planes * Bottleneck.outchannel_ratio)
+        self.relu = nn.ReLU(inplace=True)
+        self.downsample = downsample
+        self.stride = stride
+        self.shake_drop = ShakeDrop(p_shakedrop)
+
+    def forward(self, x):
+
+        out = self.bn1(x)
+        out = self.conv1(out)
+
+        out = self.bn2(out)
+        out = self.relu(out)
+        out = self.conv2(out)
+
+        out = self.bn3(out)
+        out = self.relu(out)
+        out = self.conv3(out)
+
+        out = self.bn4(out)
+
+        out = self.shake_drop(out)
+
+        if self.downsample is not None:
+            shortcut = self.downsample(x)
+            featuremap_size = shortcut.size()[2:4]
+        else:
+            shortcut = x
+            featuremap_size = out.size()[2:4]
+
+        batch_size = out.size()[0]
+        residual_channel = out.size()[1]
+        shortcut_channel = shortcut.size()[1]
+
+        if residual_channel != shortcut_channel:
+            padding = torch.autograd.Variable(
+                torch.cuda.FloatTensor(batch_size, residual_channel - shortcut_channel, featuremap_size[0],
+                                       featuremap_size[1]).fill_(0))
+            out += torch.cat((shortcut, padding), 1)
+        else:
+            out += shortcut
+
+        return out
+
+
+class PyramidNet(nn.Module):
+
+    def __init__(self, dataset, depth, alpha, num_classes, bottleneck=True):
+        super(PyramidNet, self).__init__()
+        self.dataset = dataset
+        if self.dataset.startswith('cifar'):
+            self.inplanes = 16
+            if bottleneck:
+                n = int((depth - 2) / 9)
+                block = Bottleneck
+            else:
+                n = int((depth - 2) / 6)
+                block = BasicBlock
+
+            self.addrate = alpha / (3 * n * 1.0)
+            self.ps_shakedrop = [1. - (1.0 - (0.5 / (3 * n)) * (i + 1)) for i in range(3 * n)]
+
+            self.input_featuremap_dim = self.inplanes
+            self.conv1 = nn.Conv2d(3, self.input_featuremap_dim, kernel_size=3, stride=1, padding=1, bias=False)
+            self.bn1 = nn.BatchNorm2d(self.input_featuremap_dim)
+
+            self.featuremap_dim = self.input_featuremap_dim
+            self.layer1 = self.pyramidal_make_layer(block, n)
+            self.layer2 = self.pyramidal_make_layer(block, n, stride=2)
+            self.layer3 = self.pyramidal_make_layer(block, n, stride=2)
+
+            self.final_featuremap_dim = self.input_featuremap_dim
+            self.bn_final = nn.BatchNorm2d(self.final_featuremap_dim)
+            self.relu_final = nn.ReLU(inplace=True)
+            self.avgpool = nn.AvgPool2d(8)
+            self.fc = nn.Linear(self.final_featuremap_dim, num_classes)
+
+        elif dataset == 'imagenet':
+            blocks = {18: BasicBlock, 34: BasicBlock, 50: Bottleneck, 101: Bottleneck, 152: Bottleneck, 200: Bottleneck}
+            layers = {18: [2, 2, 2, 2], 34: [3, 4, 6, 3], 50: [3, 4, 6, 3], 101: [3, 4, 23, 3], 152: [3, 8, 36, 3],
+                      200: [3, 24, 36, 3]}
+
+            if layers.get(depth) is None:
+                if bottleneck == True:
+                    blocks[depth] = Bottleneck
+                    temp_cfg = int((depth - 2) / 12)
+                else:
+                    blocks[depth] = BasicBlock
+                    temp_cfg = int((depth - 2) / 8)
+
+                layers[depth] = [temp_cfg, temp_cfg, temp_cfg, temp_cfg]
+                print('=> the layer configuration for each stage is set to', layers[depth])
+
+            self.inplanes = 64
+            self.addrate = alpha / (sum(layers[depth]) * 1.0)
+
+            self.input_featuremap_dim = self.inplanes
+            self.conv1 = nn.Conv2d(3, self.input_featuremap_dim, kernel_size=7, stride=2, padding=3, bias=False)
+            self.bn1 = nn.BatchNorm2d(self.input_featuremap_dim)
+            self.relu = nn.ReLU(inplace=True)
+            self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
+
+            self.featuremap_dim = self.input_featuremap_dim
+            self.layer1 = self.pyramidal_make_layer(blocks[depth], layers[depth][0])
+            self.layer2 = self.pyramidal_make_layer(blocks[depth], layers[depth][1], stride=2)
+            self.layer3 = self.pyramidal_make_layer(blocks[depth], layers[depth][2], stride=2)
+            self.layer4 = self.pyramidal_make_layer(blocks[depth], layers[depth][3], stride=2)
+
+            self.final_featuremap_dim = self.input_featuremap_dim
+            self.bn_final = nn.BatchNorm2d(self.final_featuremap_dim)
+            self.relu_final = nn.ReLU(inplace=True)
+            self.avgpool = nn.AvgPool2d(7)
+            self.fc = nn.Linear(self.final_featuremap_dim, num_classes)
+
+        for m in self.modules():
+            if isinstance(m, nn.Conv2d):
+                n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
+                m.weight.data.normal_(0, math.sqrt(2. / n))
+            elif isinstance(m, nn.BatchNorm2d):
+                m.weight.data.fill_(1)
+                m.bias.data.zero_()
+
+        assert len(self.ps_shakedrop) == 0, self.ps_shakedrop
+
+    def pyramidal_make_layer(self, block, block_depth, stride=1):
+        downsample = None
+        if stride != 1:  # or self.inplanes != int(round(featuremap_dim_1st)) * block.outchannel_ratio:
+            downsample = nn.AvgPool2d((2, 2), stride=(2, 2), ceil_mode=True)
+
+        layers = []
+        self.featuremap_dim = self.featuremap_dim + self.addrate
+        layers.append(block(self.input_featuremap_dim, int(round(self.featuremap_dim)), stride, downsample, p_shakedrop=self.ps_shakedrop.pop(0)))
+        for i in range(1, block_depth):
+            temp_featuremap_dim = self.featuremap_dim + self.addrate
+            layers.append(
+                block(int(round(self.featuremap_dim)) * block.outchannel_ratio, int(round(temp_featuremap_dim)), 1, p_shakedrop=self.ps_shakedrop.pop(0)))
+            self.featuremap_dim = temp_featuremap_dim
+        self.input_featuremap_dim = int(round(self.featuremap_dim)) * block.outchannel_ratio
+
+        return nn.Sequential(*layers)
+
+    def forward(self, x):
+        if self.dataset == 'cifar10' or self.dataset == 'cifar100':
+            x = self.conv1(x)
+            x = self.bn1(x)
+
+            x = self.layer1(x)
+            x = self.layer2(x)
+            x = self.layer3(x)
+
+            x = self.bn_final(x)
+            x = self.relu_final(x)
+            x = self.avgpool(x)
+            x = x.view(x.size(0), -1)
+            x = self.fc(x)
+
+        elif self.dataset == 'imagenet':
+            x = self.conv1(x)
+            x = self.bn1(x)
+            x = self.relu(x)
+            x = self.maxpool(x)
+
+            x = self.layer1(x)
+            x = self.layer2(x)
+            x = self.layer3(x)
+            x = self.layer4(x)
+
+            x = self.bn_final(x)
+            x = self.relu_final(x)
+            x = self.avgpool(x)
+            x = x.view(x.size(0), -1)
+            x = self.fc(x)
+
+        return x

+# Original code: https://github.com/pytorch/vision/blob/master/torchvision/models/resnet.py
+
+import torch.nn as nn
+import math
+
+
+def conv3x3(in_planes, out_planes, stride=1):
+    "3x3 convolution with padding"
+    return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride,
+                     padding=1, bias=False)
+
+
+class BasicBlock(nn.Module):
+    expansion = 1
+
+    def __init__(self, inplanes, planes, stride=1, downsample=None):
+        super(BasicBlock, self).__init__()
+        self.conv1 = conv3x3(inplanes, planes, stride)
+        self.bn1 = nn.BatchNorm2d(planes)
+        self.conv2 = conv3x3(planes, planes)
+        self.bn2 = nn.BatchNorm2d(planes)
+        self.relu = nn.ReLU(inplace=True)
+
+        self.downsample = downsample
+        self.stride = stride
+
+    def forward(self, x):
+        residual = x
+
+        out = self.conv1(x)
+        out = self.bn1(out)
+        out = self.relu(out)
+
+        out = self.conv2(out)
+        out = self.bn2(out)
+
+        if self.downsample is not None:
+            residual = self.downsample(x)
+
+        out += residual
+        out = self.relu(out)
+
+        return out
+
+
+class Bottleneck(nn.Module):
+    expansion = 4
+
+    def __init__(self, inplanes, planes, stride=1, downsample=None):
+        super(Bottleneck, self).__init__()
+
+        self.conv1 = nn.Conv2d(inplanes, planes, kernel_size=1, bias=False)
+        self.bn1 = nn.BatchNorm2d(planes)
+        self.conv2 = nn.Conv2d(planes, planes, kernel_size=3, stride=stride, padding=1, bias=False)
+        self.bn2 = nn.BatchNorm2d(planes)
+        self.conv3 = nn.Conv2d(planes, planes * Bottleneck.expansion, kernel_size=1, bias=False)
+        self.bn3 = nn.BatchNorm2d(planes * Bottleneck.expansion)
+        self.relu = nn.ReLU(inplace=True)
+
+        self.downsample = downsample
+        self.stride = stride
+
+    def forward(self, x):
+        residual = x
+
+        out = self.conv1(x)
+        out = self.bn1(out)
+        out = self.relu(out)
+
+        out = self.conv2(out)
+        out = self.bn2(out)
+        out = self.relu(out)
+
+        out = self.conv3(out)
+        out = self.bn3(out)
+        if self.downsample is not None:
+            residual = self.downsample(x)
+
+        out += residual
+        out = self.relu(out)
+
+        return out
+
+class ResNet(nn.Module):
+    def __init__(self, dataset, depth, num_classes, bottleneck=False):
+        super(ResNet, self).__init__()        
+        self.dataset = dataset
+        if self.dataset.startswith('cifar'):
+            self.inplanes = 16
+            print(bottleneck)
+            if bottleneck == True:
+                n = int((depth - 2) / 9)
+                block = Bottleneck
+            else:
+                n = int((depth - 2) / 6)
+                block = BasicBlock
+
+            self.conv1 = nn.Conv2d(3, self.inplanes, kernel_size=3, stride=1, padding=1, bias=False)
+            self.bn1 = nn.BatchNorm2d(self.inplanes)
+            self.relu = nn.ReLU(inplace=True)
+            self.layer1 = self._make_layer(block, 16, n)
+            self.layer2 = self._make_layer(block, 32, n, stride=2)
+            self.layer3 = self._make_layer(block, 64, n, stride=2) 
+            # self.avgpool = nn.AvgPool2d(8)
+            self.avgpool = nn.AdaptiveAvgPool2d((1, 1))
+            self.fc = nn.Linear(64 * block.expansion, num_classes)
+
+        elif dataset == 'imagenet':
+            blocks ={18: BasicBlock, 34: BasicBlock, 50: Bottleneck, 101: Bottleneck, 152: Bottleneck, 200: Bottleneck}
+            layers ={18: [2, 2, 2, 2], 34: [3, 4, 6, 3], 50: [3, 4, 6, 3], 101: [3, 4, 23, 3], 152: [3, 8, 36, 3], 200: [3, 24, 36, 3]}
+            assert layers[depth], 'invalid detph for ResNet (depth should be one of 18, 34, 50, 101, 152, and 200)'
+
+            self.inplanes = 64
+            self.conv1 = nn.Conv2d(3, self.inplanes, kernel_size=7, stride=2, padding=3, bias=False)
+            self.bn1 = nn.BatchNorm2d(64)
+            self.relu = nn.ReLU(inplace=True)
+            self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
+            self.layer1 = self._make_layer(blocks[depth], 64, layers[depth][0])
+            self.layer2 = self._make_layer(blocks[depth], 128, layers[depth][1], stride=2)
+            self.layer3 = self._make_layer(blocks[depth], 256, layers[depth][2], stride=2)
+            self.layer4 = self._make_layer(blocks[depth], 512, layers[depth][3], stride=2)
+            # self.avgpool = nn.AvgPool2d(7)
+            self.avgpool = nn.AdaptiveAvgPool2d((1, 1))
+            self.fc = nn.Linear(512 * blocks[depth].expansion, num_classes)
+
+        for m in self.modules():
+            if isinstance(m, nn.Conv2d):
+                n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
+                m.weight.data.normal_(0, math.sqrt(2. / n))
+            elif isinstance(m, nn.BatchNorm2d):
+                m.weight.data.fill_(1)
+                m.bias.data.zero_()
+
+    def _make_layer(self, block, planes, blocks, stride=1):
+        downsample = None
+        if stride != 1 or self.inplanes != planes * block.expansion:
+            downsample = nn.Sequential(
+                nn.Conv2d(self.inplanes, planes * block.expansion,
+                          kernel_size=1, stride=stride, bias=False),
+                nn.BatchNorm2d(planes * block.expansion),
+            )
+
+        layers = []
+        layers.append(block(self.inplanes, planes, stride, downsample))
+        self.inplanes = planes * block.expansion
+        for i in range(1, blocks):
+            layers.append(block(self.inplanes, planes))
+
+        return nn.Sequential(*layers)
+
+    def forward(self, x):
+        if self.dataset == 'cifar10' or self.dataset == 'cifar100':
+            x = self.conv1(x)
+            x = self.bn1(x)
+            x = self.relu(x)
+            
+            x = self.layer1(x)
+            x = self.layer2(x)
+            x = self.layer3(x)
+
+            x = self.avgpool(x)
+            x = x.view(x.size(0), -1)
+            x = self.fc(x)
+
+        elif self.dataset == 'imagenet':
+            x = self.conv1(x)
+            x = self.bn1(x)
+            x = self.relu(x)
+            x = self.maxpool(x)
+
+            x = self.layer1(x)
+            x = self.layer2(x)
+            x = self.layer3(x)
+            x = self.layer4(x)
+
+            x = self.avgpool(x)
+            x = x.view(x.size(0), -1)
+            x = self.fc(x)
+    
+        return x

+# -*- coding: utf-8 -*-
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.autograd import Variable
+
+
+class ShakeDropFunction(torch.autograd.Function):
+
+    @staticmethod
+    def forward(ctx, x, training=True, p_drop=0.5, alpha_range=[-1, 1]):
+        if training:
+            gate = torch.cuda.FloatTensor([0]).bernoulli_(1 - p_drop)
+            ctx.save_for_backward(gate)
+            if gate.item() == 0:
+                alpha = torch.cuda.FloatTensor(x.size(0)).uniform_(*alpha_range)
+                alpha = alpha.view(alpha.size(0), 1, 1, 1).expand_as(x)
+                return alpha * x
+            else:
+                return x
+        else:
+            return (1 - p_drop) * x
+
+    @staticmethod
+    def backward(ctx, grad_output):
+        gate = ctx.saved_tensors[0]
+        if gate.item() == 0:
+            beta = torch.cuda.FloatTensor(grad_output.size(0)).uniform_(0, 1)
+            beta = beta.view(beta.size(0), 1, 1, 1).expand_as(grad_output)
+            beta = Variable(beta)
+            return beta * grad_output, None, None, None
+        else:
+            return grad_output, None, None, None
+
+
+class ShakeDrop(nn.Module):
+
+    def __init__(self, p_drop=0.5, alpha_range=[-1, 1]):
+        super(ShakeDrop, self).__init__()
+        self.p_drop = p_drop
+        self.alpha_range = alpha_range
+
+    def forward(self, x):
+        return ShakeDropFunction.apply(x, self.training, self.p_drop, self.alpha_range)

+# -*- coding: utf-8 -*-
+
+import math
+
+import torch.nn as nn
+import torch.nn.functional as F
+
+from FastAutoAugment.networks.shakeshake.shakeshake import ShakeShake
+from FastAutoAugment.networks.shakeshake.shakeshake import Shortcut
+
+
+class ShakeBlock(nn.Module):
+
+    def __init__(self, in_ch, out_ch, stride=1):
+        super(ShakeBlock, self).__init__()
+        self.equal_io = in_ch == out_ch
+        self.shortcut = self.equal_io and None or Shortcut(in_ch, out_ch, stride=stride)
+
+        self.branch1 = self._make_branch(in_ch, out_ch, stride)
+        self.branch2 = self._make_branch(in_ch, out_ch, stride)
+
+    def forward(self, x):
+        h1 = self.branch1(x)
+        h2 = self.branch2(x)
+        h = ShakeShake.apply(h1, h2, self.training)
+        h0 = x if self.equal_io else self.shortcut(x)
+        return h + h0
+
+    def _make_branch(self, in_ch, out_ch, stride=1):
+        return nn.Sequential(
+            nn.ReLU(inplace=False),
+            nn.Conv2d(in_ch, out_ch, 3, padding=1, stride=stride, bias=False),
+            nn.BatchNorm2d(out_ch),
+            nn.ReLU(inplace=False),
+            nn.Conv2d(out_ch, out_ch, 3, padding=1, stride=1, bias=False),
+            nn.BatchNorm2d(out_ch))
+
+
+class ShakeResNet(nn.Module):
+
+    def __init__(self, depth, w_base, label):
+        super(ShakeResNet, self).__init__()
+        n_units = (depth - 2) / 6
+
+        in_chs = [16, w_base, w_base * 2, w_base * 4]
+        self.in_chs = in_chs
+
+        self.c_in = nn.Conv2d(3, in_chs[0], 3, padding=1)
+        self.layer1 = self._make_layer(n_units, in_chs[0], in_chs[1])
+        self.layer2 = self._make_layer(n_units, in_chs[1], in_chs[2], 2)
+        self.layer3 = self._make_layer(n_units, in_chs[2], in_chs[3], 2)
+        self.fc_out = nn.Linear(in_chs[3], label)
+
+        # Initialize paramters
+        for m in self.modules():
+            if isinstance(m, nn.Conv2d):
+                n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
+                m.weight.data.normal_(0, math.sqrt(2. / n))
+            elif isinstance(m, nn.BatchNorm2d):
+                m.weight.data.fill_(1)
+                m.bias.data.zero_()
+            elif isinstance(m, nn.Linear):
+                m.bias.data.zero_()
+
+    def forward(self, x):
+        h = self.c_in(x)
+        h = self.layer1(h)
+        h = self.layer2(h)
+        h = self.layer3(h)
+        h = F.relu(h)
+        h = F.avg_pool2d(h, 8)
+        h = h.view(-1, self.in_chs[3])
+        h = self.fc_out(h)
+        return h
+
+    def _make_layer(self, n_units, in_ch, out_ch, stride=1):
+        layers = []
+        for i in range(int(n_units)):
+            layers.append(ShakeBlock(in_ch, out_ch, stride=stride))
+            in_ch, stride = out_ch, 1
+        return nn.Sequential(*layers)

+# -*- coding: utf-8 -*-
+
+import math
+
+import torch.nn as nn
+import torch.nn.functional as F
+
+from FastAutoAugment.networks.shakeshake.shakeshake import ShakeShake
+from FastAutoAugment.networks.shakeshake.shakeshake import Shortcut
+
+
+class ShakeBottleNeck(nn.Module):
+
+    def __init__(self, in_ch, mid_ch, out_ch, cardinary, stride=1):
+        super(ShakeBottleNeck, self).__init__()
+        self.equal_io = in_ch == out_ch
+        self.shortcut = None if self.equal_io else Shortcut(in_ch, out_ch, stride=stride)
+
+        self.branch1 = self._make_branch(in_ch, mid_ch, out_ch, cardinary, stride)
+        self.branch2 = self._make_branch(in_ch, mid_ch, out_ch, cardinary, stride)
+
+    def forward(self, x):
+        h1 = self.branch1(x)
+        h2 = self.branch2(x)
+        h = ShakeShake.apply(h1, h2, self.training)
+        h0 = x if self.equal_io else self.shortcut(x)
+        return h + h0
+
+    def _make_branch(self, in_ch, mid_ch, out_ch, cardinary, stride=1):
+        return nn.Sequential(
+            nn.Conv2d(in_ch, mid_ch, 1, padding=0, bias=False),
+            nn.BatchNorm2d(mid_ch),
+            nn.ReLU(inplace=False),
+            nn.Conv2d(mid_ch, mid_ch, 3, padding=1, stride=stride, groups=cardinary, bias=False),
+            nn.BatchNorm2d(mid_ch),
+            nn.ReLU(inplace=False),
+            nn.Conv2d(mid_ch, out_ch, 1, padding=0, bias=False),
+            nn.BatchNorm2d(out_ch))
+
+
+class ShakeResNeXt(nn.Module):
+
+    def __init__(self, depth, w_base, cardinary, label):
+        super(ShakeResNeXt, self).__init__()
+        n_units = (depth - 2) // 9
+        n_chs = [64, 128, 256, 1024]
+        self.n_chs = n_chs
+        self.in_ch = n_chs[0]
+
+        self.c_in = nn.Conv2d(3, n_chs[0], 3, padding=1)
+        self.layer1 = self._make_layer(n_units, n_chs[0], w_base, cardinary)
+        self.layer2 = self._make_layer(n_units, n_chs[1], w_base, cardinary, 2)
+        self.layer3 = self._make_layer(n_units, n_chs[2], w_base, cardinary, 2)
+        self.fc_out = nn.Linear(n_chs[3], label)
+
+        # Initialize paramters
+        for m in self.modules():
+            if isinstance(m, nn.Conv2d):
+                n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
+                m.weight.data.normal_(0, math.sqrt(2. / n))
+            elif isinstance(m, nn.BatchNorm2d):
+                m.weight.data.fill_(1)
+                m.bias.data.zero_()
+            elif isinstance(m, nn.Linear):
+                m.bias.data.zero_()
+
+    def forward(self, x):
+        h = self.c_in(x)
+        h = self.layer1(h)
+        h = self.layer2(h)
+        h = self.layer3(h)
+        h = F.relu(h)
+        h = F.avg_pool2d(h, 8)
+        h = h.view(-1, self.n_chs[3])
+        h = self.fc_out(h)
+        return h
+
+    def _make_layer(self, n_units, n_ch, w_base, cardinary, stride=1):
+        layers = []
+        mid_ch, out_ch = n_ch * (w_base // 64) * cardinary, n_ch * 4
+        for i in range(n_units):
+            layers.append(ShakeBottleNeck(self.in_ch, mid_ch, out_ch, cardinary, stride=stride))
+            self.in_ch, stride = out_ch, 1
+        return nn.Sequential(*layers)

+# -*- coding: utf-8 -*-
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.autograd import Variable
+
+
+class ShakeShake(torch.autograd.Function):
+
+    @staticmethod
+    def forward(ctx, x1, x2, training=True):
+        if training:
+            alpha = torch.cuda.FloatTensor(x1.size(0)).uniform_()
+            alpha = alpha.view(alpha.size(0), 1, 1, 1).expand_as(x1)
+        else:
+            alpha = 0.5
+        return alpha * x1 + (1 - alpha) * x2
+
+    @staticmethod
+    def backward(ctx, grad_output):
+        beta = torch.cuda.FloatTensor(grad_output.size(0)).uniform_()
+        beta = beta.view(beta.size(0), 1, 1, 1).expand_as(grad_output)
+        beta = Variable(beta)
+
+        return beta * grad_output, (1 - beta) * grad_output, None
+
+
+class Shortcut(nn.Module):
+
+    def __init__(self, in_ch, out_ch, stride):
+        super(Shortcut, self).__init__()
+        self.stride = stride
+        self.conv1 = nn.Conv2d(in_ch, out_ch // 2, 1, stride=1, padding=0, bias=False)
+        self.conv2 = nn.Conv2d(in_ch, out_ch // 2, 1, stride=1, padding=0, bias=False)
+        self.bn = nn.BatchNorm2d(out_ch)
+
+    def forward(self, x):
+        h = F.relu(x)
+
+        h1 = F.avg_pool2d(h, 1, self.stride)
+        h1 = self.conv1(h1)
+
+        h2 = F.avg_pool2d(F.pad(h, (-1, 1, -1, 1)), 1, self.stride)
+        h2 = self.conv2(h2)
+
+        h = torch.cat((h1, h2), 1)
+        return self.bn(h)

+import torch.nn as nn
+import torch.nn.init as init
+import torch.nn.functional as F
+import numpy as np
+
+
+def conv3x3(in_planes, out_planes, stride=1):
+    return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride, padding=1, bias=True)
+
+
+def conv_init(m):
+    classname = m.__class__.__name__
+    if classname.find('Conv') != -1:
+        init.xavier_uniform_(m.weight, gain=np.sqrt(2))
+        init.constant_(m.bias, 0)
+    elif classname.find('BatchNorm') != -1:
+        init.constant_(m.weight, 1)
+        init.constant_(m.bias, 0)
+
+
+class WideBasic(nn.Module):
+    def __init__(self, in_planes, planes, dropout_rate, stride=1):
+        super(WideBasic, self).__init__()
+        self.bn1 = nn.BatchNorm2d(in_planes, momentum=0.9)
+        self.conv1 = nn.Conv2d(in_planes, planes, kernel_size=3, padding=1, bias=True)
+        self.dropout = nn.Dropout(p=dropout_rate)
+        self.bn2 = nn.BatchNorm2d(planes, momentum=0.9)
+        self.conv2 = nn.Conv2d(planes, planes, kernel_size=3, stride=stride, padding=1, bias=True)
+
+        self.shortcut = nn.Sequential()
+        if stride != 1 or in_planes != planes:
+            self.shortcut = nn.Sequential(
+                nn.Conv2d(in_planes, planes, kernel_size=1, stride=stride, bias=True),
+            )
+
+    def forward(self, x):
+        out = self.dropout(self.conv1(F.relu(self.bn1(x))))
+        out = self.conv2(F.relu(self.bn2(out)))
+        out += self.shortcut(x)
+
+        return out
+
+
+class WideResNet(nn.Module):
+    def __init__(self, depth, widen_factor, dropout_rate, num_classes):
+        super(WideResNet, self).__init__()
+        self.in_planes = 16
+
+        assert ((depth - 4) % 6 == 0), 'Wide-resnet depth should be 6n+4'
+        n = int((depth - 4) / 6)
+        k = widen_factor
+
+        nStages = [16, 16*k, 32*k, 64*k]
+
+        self.conv1 = conv3x3(3, nStages[0])
+        self.layer1 = self._wide_layer(WideBasic, nStages[1], n, dropout_rate, stride=1)
+        self.layer2 = self._wide_layer(WideBasic, nStages[2], n, dropout_rate, stride=2)
+        self.layer3 = self._wide_layer(WideBasic, nStages[3], n, dropout_rate, stride=2)
+        self.bn1 = nn.BatchNorm2d(nStages[3], momentum=0.9)
+        self.linear = nn.Linear(nStages[3], num_classes)
+
+        # self.apply(conv_init)
+
+    def _wide_layer(self, block, planes, num_blocks, dropout_rate, stride):
+        strides = [stride] + [1]*(num_blocks-1)
+        layers = []
+
+        for stride in strides:
+            layers.append(block(self.in_planes, planes, dropout_rate, stride))
+            self.in_planes = planes
+
+        return nn.Sequential(*layers)
+
+    def forward(self, x):
+        out = self.conv1(x)
+        out = self.layer1(out)
+        out = self.layer2(out)
+        out = self.layer3(out)
+        out = F.relu(self.bn1(out))
+        # out = F.avg_pool2d(out, 8)
+        out = F.adaptive_avg_pool2d(out, (1, 1))
+        out = out.view(out.size(0), -1)
+        out = self.linear(out)
+
+        return out

+# Copyright 2019 Uber Technologies, Inc. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+import os
+import psutil
+import re
+import signal
+import subprocess
+import sys
+import threading
+import time
+
+
+GRACEFUL_TERMINATION_TIME_S = 5
+
+
+def terminate_executor_shell_and_children(pid):
+    print('terminate_executor_shell_and_children+', pid)
+    # If the shell already ends, no need to terminate its child.
+    try:
+        p = psutil.Process(pid)
+    except psutil.NoSuchProcess:
+        print('nosuchprocess')
+        return
+
+    # Terminate children gracefully.
+    for child in p.children():
+        try:
+            child.terminate()
+        except psutil.NoSuchProcess:
+            pass
+
+    # Wait for graceful termination.
+    time.sleep(GRACEFUL_TERMINATION_TIME_S)
+
+    # Send STOP to executor shell to stop progress.
+    p.send_signal(signal.SIGSTOP)
+
+    # Kill children recursively.
+    for child in p.children(recursive=True):
+        try:
+            child.kill()
+        except psutil.NoSuchProcess:
+            pass
+
+    # Kill shell itself.
+    p.kill()
+    print('terminate_executor_shell_and_children-', pid)
+
+
+def forward_stream(src_fd, dst_stream, prefix, index):
+    with os.fdopen(src_fd, 'r') as src:
+        line_buffer = ''
+        while True:
+            text = os.read(src.fileno(), 1000)
+            if not isinstance(text, str):
+                text = text.decode('utf-8')
+            if not text:
+                break
+
+            for line in re.split('([\r\n])', text):
+                line_buffer += line
+                if line == '\r' or line == '\n':
+                    if index is not None:
+                        localtime = time.asctime(time.localtime(time.time()))
+                        line_buffer = '{time}[{rank}]<{prefix}>:{line}'.format(
+                            time=localtime,
+                            rank=str(index),
+                            prefix=prefix,
+                            line=line_buffer
+                        )
+
+                    dst_stream.write(line_buffer)
+                    dst_stream.flush()
+                    line_buffer = ''
+
+
+def execute(command, env=None, stdout=None, stderr=None, index=None, event=None):
+    # Make a pipe for the subprocess stdout/stderr.
+    (stdout_r, stdout_w) = os.pipe()
+    (stderr_r, stderr_w) = os.pipe()
+
+    # Make a pipe for notifying the child that parent has died.
+    (r, w) = os.pipe()
+
+    middleman_pid = os.fork()
+    if middleman_pid == 0:
+        # Close unused file descriptors to enforce PIPE behavior.
+        os.close(w)
+        os.setsid()
+
+        executor_shell = subprocess.Popen(command, shell=True, env=env,
+                                          stdout=stdout_w, stderr=stderr_w)
+
+        sigterm_received = threading.Event()
+
+        def set_sigterm_received(signum, frame):
+            sigterm_received.set()
+
+        signal.signal(signal.SIGINT, set_sigterm_received)
+        signal.signal(signal.SIGTERM, set_sigterm_received)
+
+        def kill_executor_children_if_parent_dies():
+            # This read blocks until the pipe is closed on the other side
+            # due to the process termination.
+            os.read(r, 1)
+            terminate_executor_shell_and_children(executor_shell.pid)
+
+        bg = threading.Thread(target=kill_executor_children_if_parent_dies)
+        bg.daemon = True
+        bg.start()
+
+        def kill_executor_children_if_sigterm_received():
+            sigterm_received.wait()
+            terminate_executor_shell_and_children(executor_shell.pid)
+
+        bg = threading.Thread(target=kill_executor_children_if_sigterm_received)
+        bg.daemon = True
+        bg.start()
+
+        exit_code = executor_shell.wait()
+        os._exit(exit_code)
+
+    # Close unused file descriptors to enforce PIPE behavior.
+    os.close(r)
+    os.close(stdout_w)
+    os.close(stderr_w)
+
+    # Redirect command stdout & stderr to provided streams or sys.stdout/sys.stderr.
+    # This is useful for Jupyter Notebook that uses custom sys.stdout/sys.stderr or
+    # for redirecting to a file on disk.
+    if stdout is None:
+        stdout = sys.stdout
+    if stderr is None:
+        stderr = sys.stderr
+    stdout_fwd = threading.Thread(target=forward_stream, args=(stdout_r, stdout, 'stdout', index))
+    stderr_fwd = threading.Thread(target=forward_stream, args=(stderr_r, stderr, 'stderr', index))
+    stdout_fwd.start()
+    stderr_fwd.start()
+
+    def kill_middleman_if_master_thread_terminate():
+        event.wait()
+        try:
+            os.kill(middleman_pid, signal.SIGTERM)
+        except:
+            # The process has already been killed elsewhere
+            pass
+
+    # TODO: Currently this requires explicitly declaration of the event and signal handler to set
+    #  the event (gloo_run.py:_launch_jobs()). Need to figure out a generalized way to hide this behind
+    #  interfaces.
+    if event is not None:
+        bg_thread = threading.Thread(target=kill_middleman_if_master_thread_terminate)
+        bg_thread.daemon = True
+        bg_thread.start()
+
+    try:
+        res, status = os.waitpid(middleman_pid, 0)
+    except:
+        # interrupted, send middleman TERM signal which will terminate children
+        os.kill(middleman_pid, signal.SIGTERM)
+        while True:
+            try:
+                _, status = os.waitpid(middleman_pid, 0)
+                break
+            except:
+                # interrupted, wait for middleman to finish
+                pass
+
+    stdout_fwd.join()
+    stderr_fwd.join()
+    exit_code = status >> 8
+    return exit_code

+import copy
+import os
+import sys
+import time
+from collections import OrderedDict, defaultdict
+
+import torch
+
+import numpy as np
+from hyperopt import hp
+import ray
+import gorilla
+from ray.tune.trial import Trial
+from ray.tune.trial_runner import TrialRunner
+from ray.tune.suggest import HyperOptSearch
+from ray.tune import register_trainable, run_experiments
+from tqdm import tqdm
+
+from FastAutoAugment.archive import remove_deplicates, policy_decoder
+from FastAutoAugment.augmentations import augment_list
+from FastAutoAugment.common import get_logger, add_filehandler
+from FastAutoAugment.data import get_dataloaders
+from FastAutoAugment.metrics import Accumulator
+from FastAutoAugment.networks import get_model, num_class
+from FastAutoAugment.train import train_and_eval
+from theconf import Config as C, ConfigArgumentParser
+
+
+top1_valid_by_cv = defaultdict(lambda: list)
+
+
+def step_w_log(self):
+    original = gorilla.get_original_attribute(ray.tune.trial_runner.TrialRunner, 'step')
+
+    # log
+    cnts = OrderedDict()
+    for status in [Trial.RUNNING, Trial.TERMINATED, Trial.PENDING, Trial.PAUSED, Trial.ERROR]:
+        cnt = len(list(filter(lambda x: x.status == status, self._trials)))
+        cnts[status] = cnt
+    best_top1_acc = 0.
+    for trial in filter(lambda x: x.status == Trial.TERMINATED, self._trials):
+        if not trial.last_result:
+            continue
+        best_top1_acc = max(best_top1_acc, trial.last_result['top1_valid'])
+    print('iter', self._iteration, 'top1_acc=%.3f' % best_top1_acc, cnts, end='\r')
+    return original(self)
+
+
+patch = gorilla.Patch(ray.tune.trial_runner.TrialRunner, 'step', step_w_log, settings=gorilla.Settings(allow_hit=True))
+gorilla.apply(patch)
+
+
+logger = get_logger('Fast AutoAugment')
+
+
+def _get_path(dataset, model, tag):
+    return os.path.join(os.path.dirname(os.path.realpath(__file__)), 'models/%s_%s_%s.model' % (dataset, model, tag))     # TODO
+
+
+@ray.remote(num_gpus=4, max_calls=1)
+def train_model(config, dataroot, augment, cv_ratio_test, cv_fold, save_path=None, skip_exist=False):
+    C.get()
+    C.get().conf = config
+    C.get()['aug'] = augment
+
+    result = train_and_eval(None, dataroot, cv_ratio_test, cv_fold, save_path=save_path, only_eval=skip_exist)
+    return C.get()['model']['type'], cv_fold, result
+
+
+def eval_tta(config, augment, reporter):
+    C.get()
+    C.get().conf = config
+    cv_ratio_test, cv_fold, save_path = augment['cv_ratio_test'], augment['cv_fold'], augment['save_path']
+
+    # setup - provided augmentation rules
+    C.get()['aug'] = policy_decoder(augment, augment['num_policy'], augment['num_op'])
+
+    # eval
+    model = get_model(C.get()['model'], num_class(C.get()['dataset']))
+    ckpt = torch.load(save_path)
+    if 'model' in ckpt:
+        model.load_state_dict(ckpt['model'])
+    else:
+        model.load_state_dict(ckpt)
+    model.eval()
+
+    loaders = []
+    for _ in range(augment['num_policy']):  # TODO
+        _, tl, validloader, tl2 = get_dataloaders(C.get()['dataset'], C.get()['batch'], augment['dataroot'], cv_ratio_test, split_idx=cv_fold)
+        loaders.append(iter(validloader))
+        del tl, tl2
+
+    start_t = time.time()
+    metrics = Accumulator()
+    loss_fn = torch.nn.CrossEntropyLoss(reduction='none')
+    try:
+        while True:
+            losses = []
+            corrects = []
+            for loader in loaders:
+                data, label = next(loader)
+                data = data.cuda()
+                label = label.cuda()
+
+                pred = model(data)
+
+                loss = loss_fn(pred, label)
+                losses.append(loss.detach().cpu().numpy())
+
+                _, pred = pred.topk(1, 1, True, True)
+                pred = pred.t()
+                correct = pred.eq(label.view(1, -1).expand_as(pred)).detach().cpu().numpy()
+                corrects.append(correct)
+                del loss, correct, pred, data, label
+
+            losses = np.concatenate(losses)
+            losses_min = np.min(losses, axis=0).squeeze()
+
+            corrects = np.concatenate(corrects)
+            corrects_max = np.max(corrects, axis=0).squeeze()
+            metrics.add_dict({
+                'minus_loss': -1 * np.sum(losses_min),
+                'correct': np.sum(corrects_max),
+                'cnt': len(corrects_max)
+            })
+            del corrects, corrects_max
+    except StopIteration:
+        pass
+
+    del model
+    metrics = metrics / 'cnt'
+    gpu_secs = (time.time() - start_t) * torch.cuda.device_count()
+    reporter(minus_loss=metrics['minus_loss'], top1_valid=metrics['correct'], elapsed_time=gpu_secs, done=True)
+    return metrics['correct']
+
+
+if __name__ == '__main__':
+    import json
+    from pystopwatch2 import PyStopwatch
+    w = PyStopwatch()
+
+    parser = ConfigArgumentParser(conflict_handler='resolve')
+    parser.add_argument('--dataroot', type=str, default='/data/private/pretrainedmodels', help='torchvision data folder')
+    parser.add_argument('--until', type=int, default=5)
+    parser.add_argument('--num-op', type=int, default=2)
+    parser.add_argument('--num-policy', type=int, default=5)
+    parser.add_argument('--num-search', type=int, default=200)
+    parser.add_argument('--cv-ratio', type=float, default=0.4)
+    parser.add_argument('--decay', type=float, default=-1)
+    parser.add_argument('--redis', type=str, default='gpu-cloud-vnode30.dakao.io:23655')
+    parser.add_argument('--per-class', action='store_true')
+    parser.add_argument('--resume', action='store_true')
+    parser.add_argument('--smoke-test', action='store_true')
+    args = parser.parse_args()
+
+    if args.decay > 0:
+        logger.info('decay=%.4f' % args.decay)
+        C.get()['optimizer']['decay'] = args.decay
+
+    add_filehandler(logger, os.path.join('models', '%s_%s_cv%.1f.log' % (C.get()['dataset'], C.get()['model']['type'], args.cv_ratio)))
+    logger.info('configuration...')
+    logger.info(json.dumps(C.get().conf, sort_keys=True, indent=4))
+    logger.info('initialize ray...')
+    ray.init(redis_address=args.redis)
+
+    num_result_per_cv = 10
+    cv_num = 5
+    copied_c = copy.deepcopy(C.get().conf)
+
+    logger.info('search augmentation policies, dataset=%s model=%s' % (C.get()['dataset'], C.get()['model']['type']))
+    logger.info('----- Train without Augmentations cv=%d ratio(test)=%.1f -----' % (cv_num, args.cv_ratio))
+    w.start(tag='train_no_aug')
+    paths = [_get_path(C.get()['dataset'], C.get()['model']['type'], 'ratio%.1f_fold%d' % (args.cv_ratio, i)) for i in range(cv_num)]
+    print(paths)
+    reqs = [
+        train_model.remote(copy.deepcopy(copied_c), args.dataroot, C.get()['aug'], args.cv_ratio, i, save_path=paths[i], skip_exist=True)
+        for i in range(cv_num)]
+
+    tqdm_epoch = tqdm(range(C.get()['epoch']))
+    is_done = False
+    for epoch in tqdm_epoch:
+        while True:
+            epochs_per_cv = OrderedDict()
+            for cv_idx in range(cv_num):
+                try:
+                    latest_ckpt = torch.load(paths[cv_idx])
+                    if 'epoch' not in latest_ckpt:
+                        epochs_per_cv['cv%d' % (cv_idx + 1)] = C.get()['epoch']
+                        continue
+                    epochs_per_cv['cv%d' % (cv_idx+1)] = latest_ckpt['epoch']
+                except Exception as e:
+                    continue
+            tqdm_epoch.set_postfix(epochs_per_cv)
+            if len(epochs_per_cv) == cv_num and min(epochs_per_cv.values()) >= C.get()['epoch']:
+                is_done = True
+            if len(epochs_per_cv) == cv_num and min(epochs_per_cv.values()) >= epoch:
+                break
+            time.sleep(10)
+        if is_done:
+            break
+
+    logger.info('getting results...')
+    pretrain_results = ray.get(reqs)
+    for r_model, r_cv, r_dict in pretrain_results:
+        logger.info('model=%s cv=%d top1_train=%.4f top1_valid=%.4f' % (r_model, r_cv+1, r_dict['top1_train'], r_dict['top1_valid']))
+    logger.info('processed in %.4f secs' % w.pause('train_no_aug'))
+
+    if args.until == 1:
+        sys.exit(0)
+
+    logger.info('----- Search Test-Time Augmentation Policies -----')
+    w.start(tag='search')
+
+    ops = augment_list(False)
+    space = {}
+    for i in range(args.num_policy):
+        for j in range(args.num_op):
+            space['policy_%d_%d' % (i, j)] = hp.choice('policy_%d_%d' % (i, j), list(range(0, len(ops))))
+            space['prob_%d_%d' % (i, j)] = hp.uniform('prob_%d_ %d' % (i, j), 0.0, 1.0)
+            space['level_%d_%d' % (i, j)] = hp.uniform('level_%d_ %d' % (i, j), 0.0, 1.0)
+
+    final_policy_set = []
+    total_computation = 0
+    reward_attr = 'top1_valid'      # top1_valid or minus_loss
+    for _ in range(1):  # run multiple times.
+        for cv_fold in range(cv_num):
+            name = "search_%s_%s_fold%d_ratio%.1f" % (C.get()['dataset'], C.get()['model']['type'], cv_fold, args.cv_ratio)
+            print(name)
+            register_trainable(name, lambda augs, rpt: eval_tta(copy.deepcopy(copied_c), augs, rpt))
+            algo = HyperOptSearch(space, max_concurrent=4*20, reward_attr=reward_attr)
+
+            exp_config = {
+                name: {
+                    'run': name,
+                    'num_samples': 4 if args.smoke_test else args.num_search,
+                    'resources_per_trial': {'gpu': 1},
+                    'stop': {'training_iteration': args.num_policy},
+                    'config': {
+                        'dataroot': args.dataroot, 'save_path': paths[cv_fold],
+                        'cv_ratio_test': args.cv_ratio, 'cv_fold': cv_fold,
+                        'num_op': args.num_op, 'num_policy': args.num_policy
+                    },
+                }
+            }
+            results = run_experiments(exp_config, search_alg=algo, scheduler=None, verbose=0, queue_trials=True, resume=args.resume, raise_on_failed_trial=False)
+            print()
+            results = [x for x in results if x.last_result is not None]
+            results = sorted(results, key=lambda x: x.last_result[reward_attr], reverse=True)
+
+            # calculate computation usage
+            for result in results:
+                total_computation += result.last_result['elapsed_time']
+
+            for result in results[:num_result_per_cv]:
+                final_policy = policy_decoder(result.config, args.num_policy, args.num_op)
+                logger.info('loss=%.12f top1_valid=%.4f %s' % (result.last_result['minus_loss'], result.last_result['top1_valid'], final_policy))
+
+                final_policy = remove_deplicates(final_policy)
+                final_policy_set.extend(final_policy)
+
+    logger.info(json.dumps(final_policy_set))
+    logger.info('final_policy=%d' % len(final_policy_set))
+    logger.info('processed in %.4f secs, gpu hours=%.4f' % (w.pause('search'), total_computation / 3600.))
+    logger.info('----- Train with Augmentations model=%s dataset=%s aug=%s ratio(test)=%.1f -----' % (C.get()['model']['type'], C.get()['dataset'], C.get()['aug'], args.cv_ratio))
+    w.start(tag='train_aug')
+
+    num_experiments = 5
+    default_path = [_get_path(C.get()['dataset'], C.get()['model']['type'], 'ratio%.1f_default%d' % (args.cv_ratio, _)) for _ in range(num_experiments)]
+    augment_path = [_get_path(C.get()['dataset'], C.get()['model']['type'], 'ratio%.1f_augment%d' % (args.cv_ratio, _)) for _ in range(num_experiments)]
+    reqs = [train_model.remote(copy.deepcopy(copied_c), args.dataroot, C.get()['aug'], 0.0, 0, save_path=default_path[_], skip_exist=True) for _ in range(num_experiments)] + \
+        [train_model.remote(copy.deepcopy(copied_c), args.dataroot, final_policy_set, 0.0, 0, save_path=augment_path[_]) for _ in range(num_experiments)]
+
+    tqdm_epoch = tqdm(range(C.get()['epoch']))
+    is_done = False
+    for epoch in tqdm_epoch:
+        while True:
+            epochs = OrderedDict()
+            for exp_idx in range(num_experiments):
+                try:
+                    if os.path.exists(default_path[exp_idx]):
+                        latest_ckpt = torch.load(default_path[exp_idx])
+                        epochs['default_exp%d' % (exp_idx + 1)] = latest_ckpt['epoch']
+                except:
+                    pass
+                try:
+                    if os.path.exists(augment_path[exp_idx]):
+                        latest_ckpt = torch.load(augment_path[exp_idx])
+                        epochs['augment_exp%d' % (exp_idx + 1)] = latest_ckpt['epoch']
+                except:
+                    pass
+
+            tqdm_epoch.set_postfix(epochs)
+            if len(epochs) == num_experiments*2 and min(epochs.values()) >= C.get()['epoch']:
+                is_done = True
+            if len(epochs) == num_experiments*2 and min(epochs.values()) >= epoch:
+                break
+            time.sleep(10)
+        if is_done:
+            break
+
+    logger.info('getting results...')
+    final_results = ray.get(reqs)
+
+    for train_mode in ['default', 'augment']:
+        avg = 0.
+        for _ in range(num_experiments):
+            r_model, r_cv, r_dict = final_results.pop(0)
+            logger.info('[%s] top1_train=%.4f top1_test=%.4f' % (train_mode, r_dict['top1_train'], r_dict['top1_test']))
+            avg += r_dict['top1_test']
+        avg /= num_experiments
+        logger.info('[%s] top1_test average=%.4f (#experiments=%d)' % (train_mode, avg, num_experiments))
+    logger.info('processed in %.4f secs' % w.pause('train_aug'))
+
+    logger.info(w)

+import torch
+from torch.optim.optimizer import Optimizer
+
+
+class RMSpropTF(Optimizer):
+    r"""Implements RMSprop algorithm.
+    Reimplement original formulation to match TF rmsprop
+    Proposed by G. Hinton in his
+    `course <http://www.cs.toronto.edu/~tijmen/csc321/slides/lecture_slides_lec6.pdf>`_.
+    The centered version first appears in `Generating Sequences
+    With Recurrent Neural Networks <https://arxiv.org/pdf/1308.0850v5.pdf>`_.
+    The implementation here takes the square root of the gradient average before
+    adding epsilon (note that TensorFlow interchanges these two operations). The effective
+    learning rate is thus :math:`\alpha/(\sqrt{v + \epsilon})` where :math:`\alpha` from :math:`\alpha/(\sqrt{v} + \epsilon)` where :math:`\alpha`
+    is the scheduled learning rate and :math:`v` is the weighted moving average
+    of the squared gradient.
+    Arguments:
+        params (iterable): iterable of parameters to optimize or dicts defining
+            parameter groups
+        lr (float, optional): learning rate (default: 1e-2)
+        momentum (float, optional): momentum factor (default: 0)
+        alpha (float, optional): smoothing constant (default: 0.99)
+        eps (float, optional): term added to the denominator to improve
+            numerical stability (default: 1e-8)
+        centered (bool, optional) : if ``True``, compute the centered RMSProp,
+            the gradient is normalized by an estimation of its variance
+        weight_decay (float, optional): weight decay (L2 penalty) (default: 0)
+    """
+
+    def __init__(self, params, lr=1e-2, alpha=0.99, eps=1e-8, momentum=0, weight_decay=0.0):
+        if not 0.0 <= lr:
+            raise ValueError("Invalid learning rate: {}".format(lr))
+        if not 0.0 <= eps:
+            raise ValueError("Invalid epsilon value: {}".format(eps))
+        if not 0.0 < momentum:
+            raise ValueError("Invalid momentum value: {}".format(momentum))
+        if not 0.0 <= alpha:
+            raise ValueError("Invalid alpha value: {}".format(alpha))
+        assert momentum > 0.0
+
+        defaults = dict(lr=lr, momentum=momentum, alpha=alpha, eps=eps, weight_decay=weight_decay)
+        super(RMSpropTF, self).__init__(params, defaults)
+        self.initialized = False
+
+    def __setstate__(self, state):
+        super(RMSpropTF, self).__setstate__(state)
+        for group in self.param_groups:
+            group.setdefault('momentum', 0)
+
+    def load_state_dict(self, state_dict):
+        super(RMSpropTF, self).load_state_dict(state_dict)
+        self.initialized = True
+
+    def step(self, closure=None):
+        """Performs a single optimization step.
+        We modified pytorch's RMSProp to be same as Tensorflow's
+        See : https://github.com/tensorflow/tensorflow/blob/master/tensorflow/core/kernels/training_ops.cc#L485
+
+        Arguments:
+            closure (callable, optional): A closure that reevaluates the model
+                and returns the loss.
+        """
+        loss = None
+        if closure is not None:
+            loss = closure()
+
+        for group in self.param_groups:
+            for p in group['params']:
+                if p.grad is None:
+                    continue
+                grad = p.grad.data
+                if grad.is_sparse:
+                    raise RuntimeError('RMSprop does not support sparse gradients')
+                state = self.state[p]
+
+                # State initialization
+                if len(state) == 0:
+                    assert not self.initialized
+                    state['step'] = 0
+                    state['ms'] = torch.ones_like(p.data)  #, memory_format=torch.preserve_format)
+                    state['mom'] = torch.zeros_like(p.data)  #, memory_format=torch.preserve_format)
+
+                # weight decay -----
+                if group['weight_decay'] > 0:
+                    grad = grad.add(group['weight_decay'], p.data)
+
+                rho = group['alpha']
+                ms = state['ms']
+                mom = state['mom']
+                state['step'] += 1
+
+                # ms.mul_(rho).addcmul_(1 - rho, grad, grad)
+                ms.add_(torch.mul(grad, grad).add_(-ms) * (1. - rho))
+                assert group['momentum'] > 0
+
+                # new rmsprop
+                mom.mul_(group['momentum']).addcdiv_(group['lr'], grad, (ms + group['eps']).sqrt())
+
+                p.data.add_(-1.0, mom)
+
+        return loss

+import torch
+from torch.nn import BatchNorm2d
+from torch.nn.parameter import Parameter
+import torch.distributed as dist
+from torch import nn
+
+
+class TpuBatchNormalization(nn.Module):
+    # Ref : https://github.com/tensorflow/tpu/blob/master/models/official/efficientnet/utils.py#L113
+    def __init__(self, num_features, eps=1e-5, momentum=0.1, affine=True,
+                 track_running_stats=True):
+        super(TpuBatchNormalization, self).__init__()   # num_features, eps, momentum, affine, track_running_stats)
+
+        self.weight = Parameter(torch.ones(num_features))
+        self.bias = Parameter(torch.zeros(num_features))
+
+        self.register_buffer('running_mean', torch.zeros(num_features))
+        self.register_buffer('running_var', torch.ones(num_features))
+        self.register_buffer('num_batches_tracked', torch.tensor(0, dtype=torch.long))
+
+        self.eps = eps
+        self.momentum = momentum
+
+    def _reduce_avg(self, t):
+        dist.all_reduce(t, dist.ReduceOp.SUM)
+        t.mul_(1. / dist.get_world_size())
+
+    def forward(self, input):
+        if not self.training or not dist.is_initialized():
+            bn = (input - self.running_mean.view(1, self.running_mean.shape[0], 1, 1)) / \
+                 (torch.sqrt(self.running_var.view(1, self.running_var.shape[0], 1, 1) + self.eps))
+            # print(self.weight.shape, self.bias.shape)
+            return bn.mul(self.weight.view(1, self.weight.shape[0], 1, 1)).add(self.bias.view(1, self.bias.shape[0], 1, 1))
+
+        shard_mean, shard_invstd = torch.batch_norm_stats(input, self.eps)
+        shard_vars = (1. / shard_invstd) ** 2 - self.eps
+
+        shard_square_of_mean = torch.mul(shard_mean, shard_mean)
+        shard_mean_of_square = shard_vars + shard_square_of_mean
+
+        group_mean = shard_mean.clone().detach()
+        self._reduce_avg(group_mean)
+        group_mean_of_square = shard_mean_of_square.clone().detach()
+        self._reduce_avg(group_mean_of_square)
+        group_vars = group_mean_of_square - torch.mul(group_mean, group_mean)
+
+        group_mean = group_mean.detach()
+        group_vars = group_vars.detach()
+
+        # print(self.running_mean.shape, self.running_var.shape)
+        self.running_mean.mul_(1. - self.momentum).add_(group_mean.mul(self.momentum))
+        self.running_var.mul_(1. - self.momentum).add_(group_vars.mul(self.momentum))
+        self.num_batches_tracked.add_(1)
+
+        # print(input.shape, group_mean.view(1, group_mean.shape[0], 1, 1).shape, group_vars.view(1, group_vars.shape[0], 1, 1).shape, self.eps)
+        bn = (input - group_mean.view(1, group_mean.shape[0], 1, 1)) / (torch.sqrt(group_vars.view(1, group_vars.shape[0], 1, 1) + self.eps))
+        # print(self.weight.shape, self.bias.shape)
+        return bn.mul(self.weight.view(1, self.weight.shape[0], 1, 1)).add(self.bias.view(1, self.bias.shape[0], 1, 1))

+import sys
+
+
+sys.path.append('/data/private/fast-autoaugment-public')    # TODO
+
+import itertools
+import json
+import logging
+import math
+import os
+from collections import OrderedDict
+
+import torch
+from torch import nn, optim
+from torch.nn.parallel.data_parallel import DataParallel
+from torch.nn.parallel import DistributedDataParallel
+import torch.distributed as dist
+
+from tqdm import tqdm
+from theconf import Config as C, ConfigArgumentParser
+
+from FastAutoAugment.common import get_logger, EMA, add_filehandler
+from FastAutoAugment.data import get_dataloaders
+from FastAutoAugment.lr_scheduler import adjust_learning_rate_resnet
+from FastAutoAugment.metrics import accuracy, Accumulator, CrossEntropyLabelSmooth
+from FastAutoAugment.networks import get_model, num_class
+from FastAutoAugment.tf_port.rmsprop import RMSpropTF
+from FastAutoAugment.aug_mixup import CrossEntropyMixUpLabelSmooth, mixup
+from warmup_scheduler import GradualWarmupScheduler
+
+logger = get_logger('Fast AutoAugment')
+logger.setLevel(logging.INFO)
+
+
+def run_epoch(model, loader, loss_fn, optimizer, desc_default='', epoch=0, writer=None, verbose=1, scheduler=None, is_master=True, ema=None, wd=0.0, tqdm_disabled=False):
+    if verbose:
+        loader = tqdm(loader, disable=tqdm_disabled)
+        loader.set_description('[%s %04d/%04d]' % (desc_default, epoch, C.get()['epoch']))
+
+    params_without_bn = [params for name, params in model.named_parameters() if not ('_bn' in name or '.bn' in name)]
+
+    loss_ema = None
+    metrics = Accumulator()
+    cnt = 0
+    total_steps = len(loader)
+    steps = 0
+    for data, label in loader:
+        steps += 1
+        data, label = data.cuda(), label.cuda()
+
+        if C.get().conf.get('mixup', 0.0) <= 0.0 or optimizer is None:
+            preds = model(data)
+            loss = loss_fn(preds, label)
+        else:   # mixup
+            data, targets, shuffled_targets, lam = mixup(data, label, C.get()['mixup'])
+            preds = model(data)
+            loss = loss_fn(preds, targets, shuffled_targets, lam)
+            del shuffled_targets, lam
+
+        if optimizer:
+            loss += wd * (1. / 2.) * sum([torch.sum(p ** 2) for p in params_without_bn])
+            loss.backward()
+            if C.get()['optimizer']['clip'] > 0:
+                nn.utils.clip_grad_norm_(model.parameters(), C.get()['optimizer']['clip'])
+            optimizer.step()
+            optimizer.zero_grad()
+
+            if ema is not None:
+                ema(model, (epoch - 1) * total_steps + steps)
+
+        top1, top5 = accuracy(preds, label, (1, 5))
+        metrics.add_dict({
+            'loss': loss.item() * len(data),
+            'top1': top1.item() * len(data),
+            'top5': top5.item() * len(data),
+        })
+        cnt += len(data)
+        if loss_ema:
+            loss_ema = loss_ema * 0.9 + loss.item() * 0.1
+        else:
+            loss_ema = loss.item()
+        if verbose:
+            postfix = metrics / cnt
+            if optimizer:
+                postfix['lr'] = optimizer.param_groups[0]['lr']
+            postfix['loss_ema'] = loss_ema
+            loader.set_postfix(postfix)
+
+        if scheduler is not None:
+            scheduler.step(epoch - 1 + float(steps) / total_steps)
+
+        del preds, loss, top1, top5, data, label
+
+    if tqdm_disabled and verbose:
+        if optimizer:
+            logger.info('[%s %03d/%03d] %s lr=%.6f', desc_default, epoch, C.get()['epoch'], metrics / cnt, optimizer.param_groups[0]['lr'])
+        else:
+            logger.info('[%s %03d/%03d] %s', desc_default, epoch, C.get()['epoch'], metrics / cnt)
+
+    metrics /= cnt
+    if optimizer:
+        metrics.metrics['lr'] = optimizer.param_groups[0]['lr']
+    if verbose:
+        for key, value in metrics.items():
+            writer.add_scalar(key, value, epoch)
+    return metrics
+
+
+def train_and_eval(tag, dataroot, test_ratio=0.0, cv_fold=0, reporter=None, metric='last', save_path=None, only_eval=False, local_rank=-1, evaluation_interval=5):
+    total_batch = C.get()["batch"]
+    if local_rank >= 0:
+        dist.init_process_group(backend='nccl', init_method='env://', world_size=int(os.environ['WORLD_SIZE']))
+        device = torch.device('cuda', local_rank)
+        torch.cuda.set_device(device)
+
+        C.get()['lr'] *= dist.get_world_size()
+        logger.info(f'local batch={C.get()["batch"]} world_size={dist.get_world_size()} ----> total batch={C.get()["batch"] * dist.get_world_size()}')
+        total_batch = C.get()["batch"] * dist.get_world_size()
+
+    is_master = local_rank < 0 or dist.get_rank() == 0
+    if is_master:
+        add_filehandler(logger, args.save + '.log')
+
+    if not reporter:
+        reporter = lambda **kwargs: 0
+
+    max_epoch = C.get()['epoch']
+    trainsampler, trainloader, validloader, testloader_ = get_dataloaders(C.get()['dataset'], C.get()['batch'], dataroot, test_ratio, split_idx=cv_fold, multinode=(local_rank >= 0))
+
+    # create a model & an optimizer
+    model = get_model(C.get()['model'], num_class(C.get()['dataset']), local_rank=local_rank)
+    model_ema = get_model(C.get()['model'], num_class(C.get()['dataset']), local_rank=-1)
+    model_ema.eval()
+
+    criterion_ce = criterion = CrossEntropyLabelSmooth(num_class(C.get()['dataset']), C.get().conf.get('lb_smooth', 0))
+    if C.get().conf.get('mixup', 0.0) > 0.0:
+        criterion = CrossEntropyMixUpLabelSmooth(num_class(C.get()['dataset']), C.get().conf.get('lb_smooth', 0))
+    if C.get()['optimizer']['type'] == 'sgd':
+        optimizer = optim.SGD(
+            model.parameters(),
+            lr=C.get()['lr'],
+            momentum=C.get()['optimizer'].get('momentum', 0.9),
+            weight_decay=0.0,
+            nesterov=C.get()['optimizer'].get('nesterov', True)
+        )
+    elif C.get()['optimizer']['type'] == 'rmsprop':
+        optimizer = RMSpropTF(
+            model.parameters(),
+            lr=C.get()['lr'],
+            weight_decay=0.0,
+            alpha=0.9, momentum=0.9,
+            eps=0.001
+        )
+    else:
+        raise ValueError('invalid optimizer type=%s' % C.get()['optimizer']['type'])
+
+    lr_scheduler_type = C.get()['lr_schedule'].get('type', 'cosine')
+    if lr_scheduler_type == 'cosine':
+        scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=C.get()['epoch'], eta_min=0.)
+    elif lr_scheduler_type == 'resnet':
+        scheduler = adjust_learning_rate_resnet(optimizer)
+    elif lr_scheduler_type == 'efficientnet':
+        scheduler = torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda=lambda x: 0.97 ** int((x + C.get()['lr_schedule']['warmup']['epoch']) / 2.4))
+    else:
+        raise ValueError('invalid lr_schduler=%s' % lr_scheduler_type)
+
+    if C.get()['lr_schedule'].get('warmup', None) and C.get()['lr_schedule']['warmup']['epoch'] > 0:
+        scheduler = GradualWarmupScheduler(
+            optimizer,
+            multiplier=C.get()['lr_schedule']['warmup']['multiplier'],
+            total_epoch=C.get()['lr_schedule']['warmup']['epoch'],
+            after_scheduler=scheduler
+        )
+
+    if not tag or not is_master:
+        from FastAutoAugment.metrics import SummaryWriterDummy as SummaryWriter
+        logger.warning('tag not provided, no tensorboard log.')
+    else:
+        from tensorboardX import SummaryWriter
+    writers = [SummaryWriter(log_dir='./logs/%s/%s' % (tag, x)) for x in ['train', 'valid', 'test']]
+
+    if C.get()['optimizer']['ema'] > 0.0 and is_master:
+        # https://discuss.pytorch.org/t/how-to-apply-exponential-moving-average-decay-for-variables/10856/4?u=ildoonet
+        ema = EMA(C.get()['optimizer']['ema'])
+    else:
+        ema = None
+
+    result = OrderedDict()
+    epoch_start = 1
+    if save_path != 'test.pth':     # and is_master: --> should load all data(not able to be broadcasted)
+        if save_path and os.path.exists(save_path):
+            logger.info('%s file found. loading...' % save_path)
+            data = torch.load(save_path)
+            key = 'model' if 'model' in data else 'state_dict'
+
+            if 'epoch' not in data:
+                model.load_state_dict(data)
+            else:
+                logger.info('checkpoint epoch@%d' % data['epoch'])
+                if not isinstance(model, (DataParallel, DistributedDataParallel)):
+                    model.load_state_dict({k.replace('module.', ''): v for k, v in data[key].items()})
+                else:
+                    model.load_state_dict({k if 'module.' in k else 'module.'+k: v for k, v in data[key].items()})
+                logger.info('optimizer.load_state_dict+')
+                optimizer.load_state_dict(data['optimizer'])
+                if data['epoch'] < C.get()['epoch']:
+                    epoch_start = data['epoch']
+                else:
+                    only_eval = True
+                if ema is not None:
+                    ema.shadow = data.get('ema', {}) if isinstance(data.get('ema', {}), dict) else data['ema'].state_dict()
+            del data
+        else:
+            logger.info('"%s" file not found. skip to pretrain weights...' % save_path)
+            if only_eval:
+                logger.warning('model checkpoint not found. only-evaluation mode is off.')
+            only_eval = False
+
+    if local_rank >= 0:
+        for name, x in model.state_dict().items():
+            dist.broadcast(x, 0)
+        logger.info(f'multinode init. local_rank={dist.get_rank()} is_master={is_master}')
+        torch.cuda.synchronize()
+
+    tqdm_disabled = bool(os.environ.get('TASK_NAME', '')) and local_rank != 0  # KakaoBrain Environment
+
+    if only_eval:
+        logger.info('evaluation only+')
+        model.eval()
+        rs = dict()
+        rs['train'] = run_epoch(model, trainloader, criterion, None, desc_default='train', epoch=0, writer=writers[0], is_master=is_master)
+
+        with torch.no_grad():
+            rs['valid'] = run_epoch(model, validloader, criterion, None, desc_default='valid', epoch=0, writer=writers[1], is_master=is_master)
+            rs['test'] = run_epoch(model, testloader_, criterion, None, desc_default='*test', epoch=0, writer=writers[2], is_master=is_master)
+            if ema is not None and len(ema) > 0:
+                model_ema.load_state_dict({k.replace('module.', ''): v for k, v in ema.state_dict().items()})
+                rs['valid'] = run_epoch(model_ema, validloader, criterion_ce, None, desc_default='valid(EMA)', epoch=0, writer=writers[1], verbose=is_master, tqdm_disabled=tqdm_disabled)
+                rs['test'] = run_epoch(model_ema, testloader_, criterion_ce, None, desc_default='*test(EMA)', epoch=0, writer=writers[2], verbose=is_master, tqdm_disabled=tqdm_disabled)
+        for key, setname in itertools.product(['loss', 'top1', 'top5'], ['train', 'valid', 'test']):
+            if setname not in rs:
+                continue
+            result['%s_%s' % (key, setname)] = rs[setname][key]
+        result['epoch'] = 0
+        return result
+
+    # train loop
+    best_top1 = 0
+    for epoch in range(epoch_start, max_epoch + 1):
+        if local_rank >= 0:
+            trainsampler.set_epoch(epoch)
+
+        model.train()
+        rs = dict()
+        rs['train'] = run_epoch(model, trainloader, criterion, optimizer, desc_default='train', epoch=epoch, writer=writers[0], verbose=(is_master and local_rank <= 0), scheduler=scheduler, ema=ema, wd=C.get()['optimizer']['decay'], tqdm_disabled=tqdm_disabled)
+        model.eval()
+
+        if math.isnan(rs['train']['loss']):
+            raise Exception('train loss is NaN.')
+
+        if ema is not None and C.get()['optimizer']['ema_interval'] > 0 and epoch % C.get()['optimizer']['ema_interval'] == 0:
+            logger.info(f'ema synced+ rank={dist.get_rank()}')
+            if ema is not None:
+                model.load_state_dict(ema.state_dict())
+            for name, x in model.state_dict().items():
+                # print(name)
+                dist.broadcast(x, 0)
+            torch.cuda.synchronize()
+            logger.info(f'ema synced- rank={dist.get_rank()}')
+
+        if is_master and (epoch % evaluation_interval == 0 or epoch == max_epoch):
+            with torch.no_grad():
+                rs['valid'] = run_epoch(model, validloader, criterion_ce, None, desc_default='valid', epoch=epoch, writer=writers[1], verbose=is_master, tqdm_disabled=tqdm_disabled)
+                rs['test'] = run_epoch(model, testloader_, criterion_ce, None, desc_default='*test', epoch=epoch, writer=writers[2], verbose=is_master, tqdm_disabled=tqdm_disabled)
+
+                if ema is not None:
+                    model_ema.load_state_dict({k.replace('module.', ''): v for k, v in ema.state_dict().items()})
+                    rs['valid'] = run_epoch(model_ema, validloader, criterion_ce, None, desc_default='valid(EMA)', epoch=epoch, writer=writers[1], verbose=is_master, tqdm_disabled=tqdm_disabled)
+                    rs['test'] = run_epoch(model_ema, testloader_, criterion_ce, None, desc_default='*test(EMA)', epoch=epoch, writer=writers[2], verbose=is_master, tqdm_disabled=tqdm_disabled)
+
+            logger.info(
+                f'epoch={epoch} '
+                f'[train] loss={rs["train"]["loss"]:.4f} top1={rs["train"]["top1"]:.4f} '
+                f'[valid] loss={rs["valid"]["loss"]:.4f} top1={rs["valid"]["top1"]:.4f} '
+                f'[test] loss={rs["test"]["loss"]:.4f} top1={rs["test"]["top1"]:.4f} '
+            )
+
+            if metric == 'last' or rs[metric]['top1'] > best_top1:
+                if metric != 'last':
+                    best_top1 = rs[metric]['top1']
+                for key, setname in itertools.product(['loss', 'top1', 'top5'], ['train', 'valid', 'test']):
+                    result['%s_%s' % (key, setname)] = rs[setname][key]
+                result['epoch'] = epoch
+
+                writers[1].add_scalar('valid_top1/best', rs['valid']['top1'], epoch)
+                writers[2].add_scalar('test_top1/best', rs['test']['top1'], epoch)
+
+                reporter(
+                    loss_valid=rs['valid']['loss'], top1_valid=rs['valid']['top1'],
+                    loss_test=rs['test']['loss'], top1_test=rs['test']['top1']
+                )
+
+                # save checkpoint
+                if is_master and save_path:
+                    logger.info('save model@%d to %s, err=%.4f' % (epoch, save_path, 1 - best_top1))
+                    torch.save({
+                        'epoch': epoch,
+                        'log': {
+                            'train': rs['train'].get_dict(),
+                            'valid': rs['valid'].get_dict(),
+                            'test': rs['test'].get_dict(),
+                        },
+                        'optimizer': optimizer.state_dict(),
+                        'model': model.state_dict(),
+                        'ema': ema.state_dict() if ema is not None else None,
+                    }, save_path)
+
+    del model
+
+    result['top1_test'] = best_top1
+    return result
+
+
+if __name__ == '__main__':
+    parser = ConfigArgumentParser(conflict_handler='resolve')
+    parser.add_argument('--tag', type=str, default='')
+    parser.add_argument('--dataroot', type=str, default='/data/private/pretrainedmodels', help='torchvision data folder')
+    parser.add_argument('--save', type=str, default='test.pth')
+    parser.add_argument('--cv-ratio', type=float, default=0.0)
+    parser.add_argument('--cv', type=int, default=0)
+    parser.add_argument('--local_rank', type=int, default=-1)
+    parser.add_argument('--evaluation-interval', type=int, default=5)
+    parser.add_argument('--only-eval', action='store_true')
+    args = parser.parse_args()
+
+    assert (args.only_eval and args.save) or not args.only_eval, 'checkpoint path not provided in evaluation mode.'
+
+    if not args.only_eval:
+        if args.save:
+            logger.info('checkpoint will be saved at %s' % args.save)
+        else:
+            logger.warning('Provide --save argument to save the checkpoint. Without it, training result will not be saved!')
+
+    import time
+    t = time.time()
+    result = train_and_eval(args.tag, args.dataroot, test_ratio=args.cv_ratio, cv_fold=args.cv, save_path=args.save, only_eval=args.only_eval, local_rank=args.local_rank, metric='test', evaluation_interval=args.evaluation_interval)
+    elapsed = time.time() - t
+
+    logger.info('done.')
+    logger.info('model: %s' % C.get()['model'])
+    logger.info('augmentation: %s' % C.get()['aug'])
+    logger.info('\n' + json.dumps(result, indent=4))
+    logger.info('elapsed time: %.3f Hours' % (elapsed / 3600.))
+    logger.info('top1 error in testset: %.4f' % (1. - result['top1_test']))
+    logger.info(args.save)

+import sys
+sys.path.append('/data/private/fast-autoaugment-public')    # TODO
+
+import time
+import os
+import threading
+import six
+from six.moves import queue
+
+from FastAutoAugment import safe_shell_exec
+
+
+def _exec_command(command):
+    host_output = six.StringIO()
+    try:
+        exit_code = safe_shell_exec.execute(command,
+                                            stdout=host_output,
+                                            stderr=host_output)
+        if exit_code != 0:
+            print('Launching task function was not successful:\n{host_output}'.format(host_output=host_output.getvalue()))
+            os._exit(exit_code)
+    finally:
+        host_output.close()
+    return exit_code
+
+
+def execute_function_multithreaded(fn,
+                                   args_list,
+                                   block_until_all_done=True,
+                                   max_concurrent_executions=1000):
+    """
+    Executes fn in multiple threads each with one set of the args in the
+    args_list.
+    :param fn: function to be executed
+    :type fn:
+    :param args_list:
+    :type args_list: list(list)
+    :param block_until_all_done: if is True, function will block until all the
+    threads are done and will return the results of each thread's execution.
+    :type block_until_all_done: bool
+    :param max_concurrent_executions:
+    :type max_concurrent_executions: int
+    :return:
+    If block_until_all_done is False, returns None. If block_until_all_done is
+    True, function returns the dict of results.
+        {
+            index: execution result of fn with args_list[index]
+        }
+    :rtype: dict
+    """
+    result_queue = queue.Queue()
+    worker_queue = queue.Queue()
+
+    for i, arg in enumerate(args_list):
+        arg.append(i)
+        worker_queue.put(arg)
+
+    def fn_execute():
+        while True:
+            try:
+                arg = worker_queue.get(block=False)
+            except queue.Empty:
+                return
+            exec_index = arg[-1]
+            res = fn(*arg[:-1])
+            result_queue.put((exec_index, res))
+
+    threads = []
+    number_of_threads = min(max_concurrent_executions, len(args_list))
+
+    for _ in range(number_of_threads):
+        thread = threading.Thread(target=fn_execute)
+        if not block_until_all_done:
+            thread.daemon = True
+        thread.start()
+        threads.append(thread)
+
+    # Returns the results only if block_until_all_done is set.
+    results = None
+    if block_until_all_done:
+        # Because join() cannot be interrupted by signal, a single join()
+        # needs to be separated into join()s with timeout in a while loop.
+        have_alive_child = True
+        while have_alive_child:
+            have_alive_child = False
+            for t in threads:
+                t.join(0.1)
+                if t.is_alive():
+                    have_alive_child = True
+
+        results = {}
+        while not result_queue.empty():
+            item = result_queue.get()
+            results[item[0]] = item[1]
+
+        if len(results) != len(args_list):
+            raise RuntimeError(
+                'Some threads for func {func} did not complete '
+                'successfully.'.format(func=fn.__name__))
+    return results
+
+
+if __name__ == '__main__':
+    import argparse
+
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--host', type=str)
+    parser.add_argument('--num-gpus', type=int, default=4)
+    parser.add_argument('--master', type=str, default='task1')
+    parser.add_argument('--port', type=int, default=1958)
+    parser.add_argument('-c', '--conf', type=str)
+    parser.add_argument('--args', type=str, default='')
+
+    args = parser.parse_args()
+
+    try:
+        hosts = ['task%d' % (x + 1) for x in range(int(args.host))]
+    except:
+        hosts = args.host.split(',')
+
+    cwd = os.getcwd()
+    command_list = []
+    for node_rank, host in enumerate(hosts):
+        ssh_cmd = f'ssh -t -t -o StrictHostKeyChecking=no {host} -p 22 ' \
+                  f'\'bash -O huponexit -c "cd {cwd} && ' \
+                  f'python -m torch.distributed.launch --nproc_per_node={args.num_gpus} --nnodes={len(hosts)} ' \
+                  f'--master_addr={args.master} --master_port={args.port} --node_rank={node_rank} ' \
+                  f'FastAutoAugment/train.py -c {args.conf} {args.args}"' \
+                  '\''
+        print(ssh_cmd)
+
+        command_list.append([ssh_cmd])
+
+    execute_function_multithreaded(_exec_command,
+                                   command_list[1:],
+                                   block_until_all_done=False)
+
+    print(command_list[0])
+
+    while True:
+        time.sleep(1)
+
+    # thread = threading.Thread(target=safe_shell_exec.execute, args=(command_list[0][0],))
+    # thread.start()
+    # thread.join()
+
+    # while True:
+    #     time.sleep(1)

+MIT License
+
+Copyright (c) 2019 Ildoo Kim
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

+# Fast AutoAugment **(Accepted at NeurIPS 2019)**
+
+Official [Fast AutoAugment](https://arxiv.org/abs/1905.00397) implementation in PyTorch.
+
+- Fast AutoAugment learns augmentation policies using a more efficient search strategy based on density matching.
+- Fast AutoAugment speeds up the search time by orders of magnitude while maintaining the comparable performances.
+
+<p align="center">
+<img src="etc/search.jpg" height=350>
+</p>
+
+## Results
+
+### CIFAR-10 / 100
+
+Search : **3.5 GPU Hours (1428x faster than AutoAugment)**, WResNet-40x2 on Reduced CIFAR-10
+
+| Model(CIFAR-10)         | Baseline   | Cutout     | AutoAugment | Fast AutoAugment<br/>(transfer/direct) |   |
+|-------------------------|------------|------------|-------------|------------------|----|
+| Wide-ResNet-40-2        | 5.3        | 4.1        | 3.7         | 3.6 / 3.7        | [Download](https://arena.kakaocdn.net/brainrepo/fast-autoaugment/cifar10_wresnet40x2_top1_3.52.pth) |
+| Wide-ResNet-28-10       | 3.9        | 3.1        | 2.6         | 2.7 / 2.7        | [Download](https://arena.kakaocdn.net/brainrepo/fast-autoaugment/cifar10_wresnet28x10_top1.pth) |
+| Shake-Shake(26 2x32d)   | 3.6        | 3.0        | 2.5         | 2.7 / 2.5        | [Download](https://arena.kakaocdn.net/brainrepo/fast-autoaugment/cifar10_shake26_2x32d_top1_2.68.pth) |
+| Shake-Shake(26 2x96d)   | 2.9        | 2.6        | 2.0         | 2.0 / 2.0        | [Download](https://arena.kakaocdn.net/brainrepo/fast-autoaugment/cifar10_shake26_2x96d_top1_1.97.pth) |
+| Shake-Shake(26 2x112d)  | 2.8        | 2.6        | 1.9         | 2.0 / 1.9        | [Download](https://arena.kakaocdn.net/brainrepo/fast-autoaugment/cifar10_shake26_2x112d_top1_2.04.pth) |
+| PyramidNet+ShakeDrop    | 2.7        | 2.3        | 1.5         | 1.8 / 1.7        | [Download](https://arena.kakaocdn.net/brainrepo/fast-autoaugment/cifar10_pyramid272_top1_1.44.pth) |
+
+| Model(CIFAR-100)      | Baseline   | Cutout     | AutoAugment | Fast AutoAugment<br/>(transfer/direct) |    |
+|-----------------------|------------|------------|-------------|------------------|----|
+| Wide-ResNet-40-2      | 26.0       | 25.2       | 20.7        | 20.7 / 20.6      | [Download](https://arena.kakaocdn.net/brainrepo/fast-autoaugment/cifar100_wresnet40x2_top1_20.43.pth) |
+| Wide-ResNet-28-10     | 18.8       | 18.4       | 17.1        | 17.3 / 17.3      | [Download](https://arena.kakaocdn.net/brainrepo/fast-autoaugment/cifar100_wresnet28x10_top1_17.17.pth) |
+| Shake-Shake(26 2x96d) | 17.1       | 16.0       | 14.3        | 14.9 / 14.6      | [Download](https://arena.kakaocdn.net/brainrepo/fast-autoaugment/cifar100_shake26_2x96d_top1_15.15.pth) |
+| PyramidNet+ShakeDrop  | 14.0       | 12.2       | 10.7        | 11.9 / 11.7      | [Download](https://arena.kakaocdn.net/brainrepo/fast-autoaugment/cifar100_pyramid272_top1_11.74.pth) |
+
+### ImageNet
+
+Search : **450 GPU Hours (33x faster than AutoAugment)**, ResNet-50 on Reduced ImageNet
+
+| Model      | Baseline   | AutoAugment | Fast AutoAugment<br/>(Top1/Top5) |    |
+|------------|------------|-------------|------------------|----|
+| ResNet-50  | 23.7 / 6.9 | 22.4 / 6.2  | **22.4 / 6.3**   | [Download](https://arena.kakaocdn.net/brainrepo/fast-autoaugment/imagenet_resnet50_top1_22.2.pth) |
+| ResNet-200 | 21.5 / 5.8 | 20.0 / 5.0  | **19.4 / 4.7**   | [Download](https://arena.kakaocdn.net/brainrepo/fast-autoaugment/imagenet_resnet200_top1_19.4.pth) |
+
+Notes
+* We evaluated resnet-50 and resnet-200 with resolution of 224 and 320, respectively. According to the original resnet paper, resnet 200 was tested with the resolution of 320. Also our resnet-200 baseline's performance was similar when we use the resolution.
+* But with recent our code clean-up and bugfixes, we've found that the baseline performs similar to the baseline even using 224x224.
+* When we use 224x224, resnet-200 performs **20.0 / 5.2**. Download link for the trained model is [here](https://arena.kakaocdn.net/brainrepo/fast-autoaugment/imagenet_resnet200_res224.pth).
+
+We have conducted additional experiments with EfficientNet.
+
+| Model | Baseline   | AutoAugment |   | Our Baseline(Batch) | +Fast AA |
+|-------|------------|-------------|---|---------------------|----------|
+| B0    | 23.2       | 22.7        |   | 22.96               | 22.68    |
+
+### SVHN Test
+
+Search : **1.5 GPU Hours**
+
+|                                  | Baseline | AutoAug / Our | Fast AutoAugment  |
+|----------------------------------|---------:|--------------:|--------:|
+| Wide-Resnet28x10                 | 1.5      | 1.1           | 1.1     |
+
+## Run
+
+We conducted experiments under
+
+- python 3.6.9
+- pytorch 1.2.0, torchvision 0.4.0, cuda10
+
+### Search a augmentation policy
+
+Please read ray's document to construct a proper ray cluster : https://github.com/ray-project/ray, and run search.py with the master's redis address.
+
+```
+$ python search.py -c confs/wresnet40x2_cifar10_b512.yaml --dataroot ... --redis ...
+```
+
+### Train a model with found policies
+
+You can train network architectures on CIFAR-10 / 100 and ImageNet with our searched policies.
+
+- fa_reduced_cifar10 : reduced CIFAR-10(4k images), WResNet-40x2
+- fa_reduced_imagenet : reduced ImageNet(50k images, 120 classes), ResNet-50
+
+```
+$ export PYTHONPATH=$PYTHONPATH:$PWD
+$ python FastAutoAugment/train.py -c confs/wresnet40x2_cifar10_b512.yaml --aug fa_reduced_cifar10 --dataset cifar10
+$ python FastAutoAugment/train.py -c confs/wresnet40x2_cifar10_b512.yaml --aug fa_reduced_cifar10 --dataset cifar100
+$ python FastAutoAugment/train.py -c confs/wresnet28x10_cifar10_b512.yaml --aug fa_reduced_cifar10 --dataset cifar10
+$ python FastAutoAugment/train.py -c confs/wresnet28x10_cifar10_b512.yaml --aug fa_reduced_cifar10 --dataset cifar100
+...
+$ python FastAutoAugment/train.py -c confs/resnet50_b512.yaml --aug fa_reduced_imagenet
+$ python FastAutoAugment/train.py -c confs/resnet200_b512.yaml --aug fa_reduced_imagenet
+```
+
+By adding --only-eval and --save arguments, you can test trained models without training.
+
+If you want to train with multi-gpu/node, use `torch.distributed.launch` such as
+
+```bash
+$ python -m torch.distributed.launch --nproc_per_node={num_gpu_per_node} --nnodes={num_node} --master_addr={master} --master_port={master_port} --node_rank={0,1,2,...,num_node} FastAutoAugment/train.py -c confs/efficientnet_b4.yaml --aug fa_reduced_imagenet
+```
+
+## Citation
+
+If you use this code in your research, please cite our [paper](https://arxiv.org/abs/1905.00397).
+
+```
+@inproceedings{lim2019fast,
+  title={Fast AutoAugment},
+  author={Lim, Sungbin and Kim, Ildoo and Kim, Taesup and Kim, Chiheon and Kim, Sungwoong},
+  booktitle={Advances in Neural Information Processing Systems (NeurIPS)},
+  year={2019}
+}
+```
+
+## Contact for Issues
+- Ildoo Kim, ildoo.kim@kakaobrain.com
+
+## References & Opensources
+
+We increase the batch size and adapt the learning rate accordingly to boost the training. Otherwise, we set other hyperparameters equal to AutoAugment if possible. For the unknown hyperparameters, we follow values from the original references or we tune them to match baseline performances.
+
+- **ResNet** : [paper1](https://arxiv.org/abs/1512.03385), [paper2](https://arxiv.org/abs/1603.05027), [code](https://github.com/osmr/imgclsmob/tree/master/pytorch/pytorchcv/models)
+- **PyramidNet** : [paper](https://arxiv.org/abs/1610.02915), [code](https://github.com/dyhan0920/PyramidNet-PyTorch)
+- **Wide-ResNet** : [code](https://github.com/meliketoy/wide-resnet.pytorch)
+- **Shake-Shake** : [code](https://github.com/owruby/shake-shake_pytorch)
+- **ShakeDrop Regularization** : [paper](https://arxiv.org/abs/1802.02375), [code](https://github.com/owruby/shake-drop_pytorch)
+- **AutoAugment** : [code](https://github.com/tensorflow/models/tree/master/research/autoaugment)
+- **Ray** : [code](https://github.com/ray-project/ray)
+- **HyperOpt** : [code](https://github.com/hyperopt/hyperopt)

+model:
+  type: efficientnet-b0
+  condconv_num_expert: 1  # if this is greater than 1(eg. 4), it activates condconv.
+dataset: imagenet
+aug: fa_reduced_imagenet
+cutout: 0
+batch: 128    # per gpu
+epoch: 350
+lr: 0.008     # 0.256 for 4096 batch
+lr_schedule:
+  type: 'efficientnet'
+  warmup:
+    multiplier: 1
+    epoch: 5
+optimizer:
+  type: rmsprop
+  decay: 0.00001
+  clip: 0
+  ema: 0.9999
+  ema_interval: -1
+lb_smooth: 0.1
\ No newline at end of file

+model:
+  type: efficientnet-b0
+  condconv_num_expert: 8  # if this is greater than 1(eg. 4), it activates condconv.
+dataset: imagenet
+aug: fa_reduced_imagenet
+cutout: 0
+batch: 128    # per gpu
+epoch: 350
+lr: 0.008     # 0.256 for 4096 batch
+lr_schedule:
+  type: 'efficientnet'
+  warmup:
+    multiplier: 1
+    epoch: 5
+optimizer:
+  type: rmsprop
+  decay: 0.00001
+  clip: 0
+  ema: 0.9999
+  ema_interval: -1
+lb_smooth: 0.1
+mixup: 0.2

+model:
+  type: efficientnet-b1
+  condconv_num_expert: 1  # if this is greater than 1(eg. 4), it activates condconv.
+dataset: imagenet
+aug: fa_reduced_imagenet
+cutout: 0
+batch: 128    # per gpu
+epoch: 350
+lr: 0.008     # 0.256 for 4096 batch
+lr_schedule:
+  type: 'efficientnet'
+  warmup:
+    multiplier: 1
+    epoch: 5
+optimizer:
+  type: rmsprop
+  decay: 0.00001
+  clip: 0
+  ema: 0.9999
+  ema_interval: -1
+lb_smooth: 0.1

+model:
+  type: efficientnet-b2
+  condconv_num_expert: 1  # if this is greater than 1(eg. 4), it activates condconv.
+dataset: imagenet
+aug: fa_reduced_imagenet
+cutout: 0
+batch: 128    # per gpu
+epoch: 350
+lr: 0.008     # 0.256 for 4096 batch
+lr_schedule:
+  type: 'efficientnet'
+  warmup:
+    multiplier: 1
+    epoch: 5
+optimizer:
+  type: rmsprop
+  decay: 0.00001
+  clip: 0
+  ema: 0.9999
+  ema_interval: -1
+lb_smooth: 0.1

+model:
+  type: efficientnet-b3
+  condconv_num_expert: 1  # if this is greater than 1(eg. 4), it activates condconv.
+dataset: imagenet
+aug: fa_reduced_imagenet
+cutout: 0
+batch: 64    # per gpu
+epoch: 350
+lr: 0.004     # 0.256 for 4096 batch
+lr_schedule:
+  type: 'efficientnet'
+  warmup:
+    multiplier: 1
+    epoch: 5
+optimizer:
+  type: rmsprop
+  decay: 0.00001
+  clip: 0
+  ema: 0.9999
+  ema_interval: -1
+lb_smooth: 0.1

+model:
+  type: efficientnet-b4
+  condconv_num_expert: 1  # if this is greater than 1(eg. 4), it activates condconv.
+dataset: imagenet
+aug: fa_reduced_imagenet
+cutout: 0
+batch: 32    # per gpu
+epoch: 350
+lr: 0.002     # 0.256 for 4096 batch
+lr_schedule:
+  type: 'efficientnet'
+  warmup:
+    multiplier: 1
+    epoch: 5
+optimizer:
+  type: rmsprop
+  decay: 0.00001
+  clip: 0
+  ema: 0.9999
+  ema_interval: -1
+lb_smooth: 0.1

+model:
+  type: pyramid
+  depth: 272
+  alpha: 200
+  bottleneck: True
+dataset: cifar10
+aug: fa_reduced_cifar10
+cutout: 16
+batch: 64
+epoch: 1800
+lr: 0.05
+lr_schedule:
+  type: 'cosine'
+  warmup:
+    multiplier: 1
+    epoch: 5
+optimizer:
+  type: sgd
+  nesterov: True
+  decay: 0.00005

+model:
+  type: resnet200
+dataset: imagenet
+aug: fa_reduced_imagenet
+cutout: 0
+batch: 64
+epoch: 270
+lr: 0.025
+lr_schedule:
+  type: 'resnet'
+  warmup:
+    multiplier: 1
+    epoch: 5
+optimizer:
+  type: sgd
+  nesterov: True
+  decay: 0.0001
+  clip: 0

+model:
+  type: resnet50
+dataset: imagenet
+aug: fa_reduced_imagenet
+cutout: 0
+batch: 128
+epoch: 270
+lr: 0.05
+lr_schedule:
+  type: 'resnet'
+  warmup:
+    multiplier: 1
+    epoch: 5
+optimizer:
+  type: sgd
+  nesterov: True
+  decay: 0.0001
+  clip: 0
+  ema: 0

+model:
+  type: resnet50
+dataset: imagenet
+aug: fa_reduced_imagenet
+cutout: 0
+batch: 128
+epoch: 270
+lr: 0.05
+lr_schedule:
+  type: 'resnet'
+  warmup:
+    multiplier: 1
+    epoch: 5
+optimizer:
+  type: sgd
+  nesterov: True
+  decay: 0.0001
+  clip: 0
+  ema: 0
+#lb_smooth: 0.1
+mixup: 0.2

+model:
+  type: shakeshake26_2x112d
+dataset: cifar10
+aug: fa_reduced_cifar10
+cutout: 16
+batch: 128
+epoch: 1800
+lr: 0.01
+lr_schedule:
+  type: 'cosine'
+  warmup:
+    multiplier: 1
+    epoch: 5
+optimizer:
+  type: sgd
+  nesterov: True
+  decay: 0.002

+model:
+  type: shakeshake26_2x32d
+dataset: cifar10
+aug: fa_reduced_cifar10
+cutout: 16
+batch: 128
+epoch: 1800
+lr: 0.01
+lr_schedule:
+  type: 'cosine'
+  warmup:
+    multiplier: 1
+    epoch: 5
+optimizer:
+  type: sgd
+  nesterov: True
+  decay: 0.001

+model:
+  type: shakeshake26_2x96d
+dataset: cifar10
+aug: fa_reduced_cifar10
+cutout: 16
+batch: 128
+epoch: 1800
+lr: 0.01
+lr_schedule:
+  type: 'cosine'
+  warmup:
+    multiplier: 1
+    epoch: 5
+optimizer:
+  type: sgd
+  nesterov: True
+  decay: 0.001

+model:
+  type: wresnet28_10
+dataset: cifar10
+aug: fa_reduced_cifar10
+cutout: 16
+batch: 128
+epoch: 200
+lr: 0.1
+lr_schedule:
+  type: 'cosine'
+  warmup:
+    multiplier: 1
+    epoch: 5
+optimizer:
+  type: sgd
+  nesterov: True
+  decay: 0.0005
\ No newline at end of file

+model:
+  type: wresnet28_10
+dataset: svhn
+aug: fa_reduced_svhn
+cutout: 20
+batch: 128
+epoch: 200
+lr: 0.01
+lr_schedule:
+  type: 'cosine'
+  warmup:
+    multiplier: 1
+    epoch: 5
+optimizer:
+  type: sgd
+  nesterov: True
+  decay: 0.0005
\ No newline at end of file

+model:
+  type: wresnet40_2
+dataset: cifar10
+aug: fa_reduced_cifar10
+cutout: 16
+batch: 128
+epoch: 200
+lr: 0.1
+lr_schedule:
+  type: 'cosine'
+  warmup:
+    multiplier: 1
+    epoch: 5
+optimizer:
+  type: sgd
+  nesterov: True
+  decay: 0.0002

+git+https://github.com/wbaek/theconf
+git+https://github.com/ildoonet/pytorch-gradual-warmup-lr.git@08f7d5e
+git+https://github.com/ildoonet/pystopwatch2.git
+git+https://github.com/hyperopt/hyperopt.git
+git+https://github.com/kakaobrain/torchlars
+
+pretrainedmodels
+tqdm
+tensorboardx
+sklearn
+ray
+matplotlib
+psutil
+requests
\ No newline at end of file

+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import pandas as pd"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/html": [
+       "<div>\n",
+       "<style scoped>\n",
+       "    .dataframe tbody tr th:only-of-type {\n",
+       "        vertical-align: middle;\n",
+       "    }\n",
+       "\n",
+       "    .dataframe tbody tr th {\n",
+       "        vertical-align: top;\n",
+       "    }\n",
+       "\n",
+       "    .dataframe thead th {\n",
+       "        text-align: right;\n",
+       "    }\n",
+       "</style>\n",
+       "<table border=\"1\" class=\"dataframe\">\n",
+       "  <thead>\n",
+       "    <tr style=\"text-align: right;\">\n",
+       "      <th></th>\n",
+       "      <th>id</th>\n",
+       "      <th>Label</th>\n",
+       "      <th>Subject</th>\n",
+       "      <th>Date</th>\n",
+       "      <th>Gender</th>\n",
+       "      <th>Age</th>\n",
+       "      <th>mmse</th>\n",
+       "      <th>ageAtEntry</th>\n",
+       "      <th>cdr</th>\n",
+       "      <th>commun</th>\n",
+       "      <th>...</th>\n",
+       "      <th>memory</th>\n",
+       "      <th>orient</th>\n",
+       "      <th>perscare</th>\n",
+       "      <th>apoe</th>\n",
+       "      <th>sumbox</th>\n",
+       "      <th>acsparnt</th>\n",
+       "      <th>height</th>\n",
+       "      <th>weight</th>\n",
+       "      <th>primStudy</th>\n",
+       "      <th>acsStudy</th>\n",
+       "    </tr>\n",
+       "  </thead>\n",
+       "  <tbody>\n",
+       "    <tr>\n",
+       "      <th>0</th>\n",
+       "      <td>/@WEBAPP/images/r.gif</td>\n",
+       "      <td>OAS30001_ClinicalData_d3025</td>\n",
+       "      <td>OAS30001</td>\n",
+       "      <td>NaN</td>\n",
+       "      <td>female</td>\n",
+       "      <td>NaN</td>\n",
+       "      <td>30.0</td>\n",
+       "      <td>65.149895</td>\n",
+       "      <td>0.0</td>\n",
+       "      <td>0.0</td>\n",
+       "      <td>...</td>\n",
+       "      <td>0.0</td>\n",
+       "      <td>0.0</td>\n",
+       "      <td>0.0</td>\n",
+       "      <td>23.0</td>\n",
+       "      <td>0.0</td>\n",
+       "      <td>NaN</td>\n",
+       "      <td>64.0</td>\n",
+       "      <td>180.0</td>\n",
+       "      <td>NaN</td>\n",
+       "      <td>NaN</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>1</th>\n",
+       "      <td>/@WEBAPP/images/r.gif</td>\n",
+       "      <td>OAS30001_ClinicalData_d3977</td>\n",
+       "      <td>OAS30001</td>\n",
+       "      <td>NaN</td>\n",
+       "      <td>female</td>\n",
+       "      <td>NaN</td>\n",
+       "      <td>29.0</td>\n",
+       "      <td>65.149895</td>\n",
+       "      <td>0.0</td>\n",
+       "      <td>0.0</td>\n",
+       "      <td>...</td>\n",
+       "      <td>0.0</td>\n",
+       "      <td>0.0</td>\n",
+       "      <td>0.0</td>\n",
+       "      <td>23.0</td>\n",
+       "      <td>0.0</td>\n",
+       "      <td>NaN</td>\n",
+       "      <td>NaN</td>\n",
+       "      <td>NaN</td>\n",
+       "      <td>NaN</td>\n",
+       "      <td>NaN</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>2</th>\n",
+       "      <td>/@WEBAPP/images/r.gif</td>\n",
+       "      <td>OAS30001_ClinicalData_d3332</td>\n",
+       "      <td>OAS30001</td>\n",
+       "      <td>NaN</td>\n",
+       "      <td>female</td>\n",
+       "      <td>NaN</td>\n",
+       "      <td>30.0</td>\n",
+       "      <td>65.149895</td>\n",
+       "      <td>0.0</td>\n",
+       "      <td>0.0</td>\n",
+       "      <td>...</td>\n",
+       "      <td>0.0</td>\n",
+       "      <td>0.0</td>\n",
+       "      <td>0.0</td>\n",
+       "      <td>23.0</td>\n",
+       "      <td>0.0</td>\n",
+       "      <td>NaN</td>\n",
+       "      <td>63.0</td>\n",
+       "      <td>185.0</td>\n",
+       "      <td>NaN</td>\n",
+       "      <td>NaN</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>3</th>\n",
+       "      <td>/@WEBAPP/images/r.gif</td>\n",
+       "      <td>OAS30001_ClinicalData_d0000</td>\n",
+       "      <td>OAS30001</td>\n",
+       "      <td>NaN</td>\n",
+       "      <td>female</td>\n",
+       "      <td>NaN</td>\n",
+       "      <td>28.0</td>\n",
+       "      <td>65.149895</td>\n",
+       "      <td>0.0</td>\n",
+       "      <td>0.0</td>\n",
+       "      <td>...</td>\n",
+       "      <td>0.0</td>\n",
+       "      <td>0.0</td>\n",
+       "      <td>0.0</td>\n",
+       "      <td>23.0</td>\n",
+       "      <td>0.0</td>\n",
+       "      <td>NaN</td>\n",
+       "      <td>NaN</td>\n",
+       "      <td>NaN</td>\n",
+       "      <td>NaN</td>\n",
+       "      <td>NaN</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>4</th>\n",
+       "      <td>/@WEBAPP/images/r.gif</td>\n",
+       "      <td>OAS30001_ClinicalData_d1456</td>\n",
+       "      <td>OAS30001</td>\n",
+       "      <td>NaN</td>\n",
+       "      <td>female</td>\n",
+       "      <td>NaN</td>\n",
+       "      <td>30.0</td>\n",
+       "      <td>65.149895</td>\n",
+       "      <td>0.0</td>\n",
+       "      <td>0.0</td>\n",
+       "      <td>...</td>\n",
+       "      <td>0.0</td>\n",
+       "      <td>0.0</td>\n",
+       "      <td>0.0</td>\n",
+       "      <td>23.0</td>\n",
+       "      <td>0.0</td>\n",
+       "      <td>NaN</td>\n",
+       "      <td>63.0</td>\n",
+       "      <td>173.0</td>\n",
+       "      <td>NaN</td>\n",
+       "      <td>NaN</td>\n",
+       "    </tr>\n",
+       "  </tbody>\n",
+       "</table>\n",
+       "<p>5 rows × 27 columns</p>\n",
+       "</div>"
+      ],
+      "text/plain": [
+       "                      id                        Label   Subject  Date  Gender  \\\n",
+       "0  /@WEBAPP/images/r.gif  OAS30001_ClinicalData_d3025  OAS30001   NaN  female   \n",
+       "1  /@WEBAPP/images/r.gif  OAS30001_ClinicalData_d3977  OAS30001   NaN  female   \n",
+       "2  /@WEBAPP/images/r.gif  OAS30001_ClinicalData_d3332  OAS30001   NaN  female   \n",
+       "3  /@WEBAPP/images/r.gif  OAS30001_ClinicalData_d0000  OAS30001   NaN  female   \n",
+       "4  /@WEBAPP/images/r.gif  OAS30001_ClinicalData_d1456  OAS30001   NaN  female   \n",
+       "\n",
+       "   Age  mmse  ageAtEntry  cdr  commun  ... memory orient perscare  apoe  \\\n",
+       "0  NaN  30.0   65.149895  0.0     0.0  ...    0.0    0.0      0.0  23.0   \n",
+       "1  NaN  29.0   65.149895  0.0     0.0  ...    0.0    0.0      0.0  23.0   \n",
+       "2  NaN  30.0   65.149895  0.0     0.0  ...    0.0    0.0      0.0  23.0   \n",
+       "3  NaN  28.0   65.149895  0.0     0.0  ...    0.0    0.0      0.0  23.0   \n",
+       "4  NaN  30.0   65.149895  0.0     0.0  ...    0.0    0.0      0.0  23.0   \n",
+       "\n",
+       "  sumbox  acsparnt  height  weight  primStudy  acsStudy  \n",
+       "0    0.0       NaN    64.0   180.0        NaN       NaN  \n",
+       "1    0.0       NaN     NaN     NaN        NaN       NaN  \n",
+       "2    0.0       NaN    63.0   185.0        NaN       NaN  \n",
+       "3    0.0       NaN     NaN     NaN        NaN       NaN  \n",
+       "4    0.0       NaN    63.0   173.0        NaN       NaN  \n",
+       "\n",
+       "[5 rows x 27 columns]"
+      ]
+     },
+     "execution_count": 3,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "all_data = pd.read_csv(\"..\\data\\ADRC clinical data_all.csv\")\n",
+    "\n",
+    "all_data.head()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 18,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Subject\n",
+      "OAS30001    0.0\n",
+      "OAS30002    0.0\n",
+      "OAS30003    0.0\n",
+      "OAS30004    0.0\n",
+      "OAS30005    0.0\n",
+      "           ... \n",
+      "OAS31168    0.0\n",
+      "OAS31169    3.0\n",
+      "OAS31170    2.0\n",
+      "OAS31171    2.0\n",
+      "OAS31172    0.0\n",
+      "Name: cdr, Length: 1098, dtype: float64\n"
+     ]
+    }
+   ],
+   "source": [
+    "ad = all_data.groupby(['Subject'])['cdr'].max()\n",
+    "print(ad)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 21,
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "'OAS30001'"
+      ]
+     },
+     "execution_count": 21,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "ad.index[0]"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 22,
+   "metadata": {},
+   "outputs": [
+    {
+     "ename": "SyntaxError",
+     "evalue": "unexpected EOF while parsing (<ipython-input-22-b8a078b72aca>, line 5)",
+     "output_type": "error",
+     "traceback": [
+      "\u001b[1;36m  File \u001b[1;32m\"<ipython-input-22-b8a078b72aca>\"\u001b[1;36m, line \u001b[1;32m5\u001b[0m\n\u001b[1;33m    #print(filtered)\u001b[0m\n\u001b[1;37m                    ^\u001b[0m\n\u001b[1;31mSyntaxError\u001b[0m\u001b[1;31m:\u001b[0m unexpected EOF while parsing\n"
+     ]
+    }
+   ],
+   "source": [
+    "filtered = []\n",
+    "for i, val in enumerate(ad):\n",
+    "    if ad[i] == 0:\n",
+    "        filtered.append(ad.index[i])\n",
+    "#print(filtered)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 23,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "df_filtered = pd.DataFrame(filtered)\n",
+    "df_filtered.to_csv('..\\data\\ADRC clinical data_normal.csv')"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "ML",
+   "language": "python",
+   "name": "ml"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.7.4"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}

+%load('..\data\MICCAI_BraTS_2019_Data_Training\name_mapping.csv')
+
+inputheader = '..\data\MICCAI_BraTS_2019_Data_Training\HGG\';
+outfolder = '..\testfolder\';
+id = 'BraTS19_2013_2_1';
+
+type = 'flair.nii';
+filename = strcat(id,'_', type); % BraTS19_2013_2_1_flair.nii
+flair_path = strcat(inputheader, id, '\', filename,'\', filename);
+%disp(path);
+flair = niftiread(flair_path); %size 240x240x155 
+cp_flair = flair;
+
+type = 'seg.nii';
+filename = strcat(id,'_', type); % BraTS19_2013_2_1_seg.nii 
+seg_path = strcat(inputheader, id, '\', filename, '\', filename);
+seg = niftiread(seg_path);
+
+[x,y,z] = size(seg);
+
+% copy flair, segment flair data
+
+cp_flair(seg == 0) = 0;
+
+% save a segmented data
+type = 'seg_flair.nii';
+filename = strcat(id,'_', type); % BraTS19_2013_2_1_seg_flair.nii 
+outpath = strcat(outfolder, filename);
+%niftiwrite(cp_flair, outpath);
+
+
+%whos seg
+
+
+%extract = seg(84, :, 86);
+
+
+
+
+% cp84 = cp_flair(84,:,86);
+% flair84 = flair(84,:, 86);
+
+
+
+%[flair,info] = ReadData3D(filename)
+% whos flair
+
+%volumeViewer(flair);
\ No newline at end of file

+inputheader = '..\data\MICCAI_BraTS_2019_Data_Training\HGG\';
+outfolder = strcat('..\data\MICCAI_BraTS_2019_Data_Training\HGG_seg_flair\');
+
+files = dir(inputheader);
+id = {files.name};
+%  files + dir  dir 
+dirFlag = [files.isdir] & ~strcmp(id, '.') & ~strcmp(id, '..'); 
+subFolders = files(dirFlag);
+disp(length(subFolders));
+
+% for k = 1 : length(subFolders)
+%   fprintf('Sub folder #%d = %s\n', k, subFolders(k).name);
+% end
+
+for i = 1 : length(subFolders)
+    
+    id = subFolders(i).name;
+    fprintf('Sub folder #%d = %s\n', i, id);
+
+    type = 'flair.nii';
+    filename = strcat(id,'_', type); % BraTS19_2013_2_1_flair.nii
+    flair_path = strcat(inputheader, id, '\', filename,'\', filename);
+    flair = niftiread(flair_path); %size 240x240x155 
+    cp_flair = flair;
+
+    type = 'seg.nii';
+    filename = strcat(id,'_', type); % BraTS19_2013_2_1_seg.nii 
+    seg_path = strcat(inputheader, id, '\', filename, '\', filename);
+    seg = niftiread(seg_path);
+
+    [x,y,z] = size(seg);
+
+    % copy flair, segment flair data
+
+    cp_flair(seg == 0) = 0;
+
+    % save a segmented data
+    type = 'seg_flair.nii';
+    filename = strcat(id,'_', type); % BraTS19_2013_2_1_seg_flair.nii 
+    outpath = strcat(outfolder, filename);
+    niftiwrite(cp_flair, outpath);
+
+end

+inputheader = '..\data\MICCAI_BraTS_2019_Data_Training\HGG\';
+outfolder = '..\testfolder\';
+id = 'BraTS19_2013_2_1';
+
+type = 'seg_flair.nii';
+filename = strcat(id,'_', type); % BraTS19_2013_2_1_flair.nii
+flair_path = strcat(outfolder, filename);
+%disp(path);
+ffffff = niftiread(flair_path); %size 240x240x155 
+
+fff84 = ffffff(84, :, 86);
\ No newline at end of file