Initial

.gitignore

@@ -0,0 +1,351 @@
+
+# Created by https://www.toptal.com/developers/gitignore/api/python,pycharm,jupyternotebooks,images
+# Edit at https://www.toptal.com/developers/gitignore?templates=python,pycharm,jupyternotebooks,images
+
+.idea/
+models/InfoSCC-GAN/
+models/BigGAN/
+models/CVAE
+*.csv
+
+
+### Images ###
+# JPEG
+*.jpg
+*.jpeg
+*.jpe
+*.jif
+*.jfif
+*.jfi
+
+# JPEG 2000
+*.jp2
+*.j2k
+*.jpf
+*.jpx
+*.jpm
+*.mj2
+
+# JPEG XR
+*.jxr
+*.hdp
+*.wdp
+
+# Graphics Interchange Format
+*.gif
+
+# RAW
+*.raw
+
+# Web P
+*.webp
+
+# Portable Network Graphics
+*.png
+
+# Animated Portable Network Graphics
+*.apng
+
+# Multiple-image Network Graphics
+*.mng
+
+# Tagged Image File Format
+*.tiff
+*.tif
+
+# Scalable Vector Graphics
+*.svg
+*.svgz
+
+# Portable Document Format
+*.pdf
+
+# X BitMap
+*.xbm
+
+# BMP
+*.bmp
+*.dib
+
+# ICO
+*.ico
+
+# 3D Images
+*.3dm
+*.max
+
+### JupyterNotebooks ###
+# gitignore template for Jupyter Notebooks
+# website: http://jupyter.org/
+
+.ipynb_checkpoints
+*/.ipynb_checkpoints/*
+
+# IPython
+profile_default/
+ipython_config.py
+
+# Remove previous ipynb_checkpoints
+#   git rm -r .ipynb_checkpoints/
+
+### PyCharm ###
+# Covers JetBrains IDEs: IntelliJ, RubyMine, PhpStorm, AppCode, PyCharm, CLion, Android Studio, WebStorm and Rider
+# Reference: https://intellij-support.jetbrains.com/hc/en-us/articles/206544839
+
+# User-specific stuff
+.idea/**/workspace.xml
+.idea/**/tasks.xml
+.idea/**/usage.statistics.xml
+.idea/**/dictionaries
+.idea/**/shelf
+
+# AWS User-specific
+.idea/**/aws.xml
+
+# Generated files
+.idea/**/contentModel.xml
+
+# Sensitive or high-churn files
+.idea/**/dataSources/
+.idea/**/dataSources.ids
+.idea/**/dataSources.local.xml
+.idea/**/sqlDataSources.xml
+.idea/**/dynamic.xml
+.idea/**/uiDesigner.xml
+.idea/**/dbnavigator.xml
+
+# Gradle
+.idea/**/gradle.xml
+.idea/**/libraries
+
+# Gradle and Maven with auto-import
+# When using Gradle or Maven with auto-import, you should exclude module files,
+# since they will be recreated, and may cause churn.  Uncomment if using
+# auto-import.
+# .idea/artifacts
+# .idea/compiler.xml
+# .idea/jarRepositories.xml
+# .idea/modules.xml
+# .idea/*.iml
+# .idea/modules
+# *.iml
+# *.ipr
+
+# CMake
+cmake-build-*/
+
+# Mongo Explorer plugin
+.idea/**/mongoSettings.xml
+
+# File-based project format
+*.iws
+
+# IntelliJ
+out/
+
+# mpeltonen/sbt-idea plugin
+.idea_modules/
+
+# JIRA plugin
+atlassian-ide-plugin.xml
+
+# Cursive Clojure plugin
+.idea/replstate.xml
+
+# SonarLint plugin
+.idea/sonarlint/
+
+# Crashlytics plugin (for Android Studio and IntelliJ)
+com_crashlytics_export_strings.xml
+crashlytics.properties
+crashlytics-build.properties
+fabric.properties
+
+# Editor-based Rest Client
+.idea/httpRequests
+
+# Android studio 3.1+ serialized cache file
+.idea/caches/build_file_checksums.ser
+
+### PyCharm Patch ###
+# Comment Reason: https://github.com/joeblau/gitignore.io/issues/186#issuecomment-215987721
+
+# *.iml
+# modules.xml
+# .idea/misc.xml
+# *.ipr
+
+# Sonarlint plugin
+# https://plugins.jetbrains.com/plugin/7973-sonarlint
+.idea/**/sonarlint/
+
+# SonarQube Plugin
+# https://plugins.jetbrains.com/plugin/7238-sonarqube-community-plugin
+.idea/**/sonarIssues.xml
+
+# Markdown Navigator plugin
+# https://plugins.jetbrains.com/plugin/7896-markdown-navigator-enhanced
+.idea/**/markdown-navigator.xml
+.idea/**/markdown-navigator-enh.xml
+.idea/**/markdown-navigator/
+
+# Cache file creation bug
+# See https://youtrack.jetbrains.com/issue/JBR-2257
+.idea/$CACHE_FILE$
+
+# CodeStream plugin
+# https://plugins.jetbrains.com/plugin/12206-codestream
+.idea/codestream.xml
+
+### Python ###
+# 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/
+share/python-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/
+.nox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+*.py,cover
+.hypothesis/
+.pytest_cache/
+cover/
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+local_settings.py
+db.sqlite3
+db.sqlite3-journal
+
+# Flask stuff:
+instance/
+.webassets-cache
+
+# Scrapy stuff:
+.scrapy
+
+# Sphinx documentation
+docs/_build/
+
+# PyBuilder
+.pybuilder/
+target/
+
+# Jupyter Notebook
+
+# IPython
+
+# pyenv
+#   For a library or package, you might want to ignore these files since the code is
+#   intended to run in multiple environments; otherwise, check them in:
+# .python-version
+
+# pipenv
+#   According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
+#   However, in case of collaboration, if having platform-specific dependencies or dependencies
+#   having no cross-platform support, pipenv may install dependencies that don't work, or not
+#   install all needed dependencies.
+#Pipfile.lock
+
+# poetry
+#   Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control.
+#   This is especially recommended for binary packages to ensure reproducibility, and is more
+#   commonly ignored for libraries.
+#   https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control
+#poetry.lock
+
+# PEP 582; used by e.g. github.com/David-OConnor/pyflow
+__pypackages__/
+
+# Celery stuff
+celerybeat-schedule
+celerybeat.pid
+
+# 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/
+.dmypy.json
+dmypy.json
+
+# Pyre type checker
+.pyre/
+
+# pytype static type analyzer
+.pytype/
+
+# Cython debug symbols
+cython_debug/
+
+# PyCharm
+#  JetBrains specific template is maintained in a separate JetBrains.gitignore that can
+#  be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore
+#  and can be added to the global gitignore or merged into this file.  For a more nuclear
+#  option (not recommended) you can uncomment the following to ignore the entire idea folder.
+#.idea/
+
+# End of https://www.toptal.com/developers/gitignore/api/python,pycharm,jupyternotebooks,images

app.py

@@ -0,0 +1,20 @@
+import streamlit as st
+
+# Custom imports
+from src.app import MultiPage
+from src.app import explore_infoscc_gan, explore_biggan, explore_cvae, compare_models
+
+# Create an instance of the app
+app = MultiPage()
+
+# Title of the main page
+st.title('Galaxy Zoo generation')
+
+# Add all your applications (pages) here
+app.add_page('Compare models', compare_models.app)
+app.add_page('Explore BigGAN', explore_biggan.app)
+app.add_page('Explore cVAE', explore_cvae.app)
+app.add_page('Explore InfoSCC-GAN', explore_infoscc_gan.app)
+
+# The main app
+app.run()

requirements.txt

@@ -0,0 +1,7 @@
+-f https://download.pytorch.org/whl/torch_stable.html
+gdown
+googledrivedownloader==0.4
+pandas==1.4.1
+streamlit==1.7.0
+torch==1.9.1+cpu
+torchvision==0.10.1+cpu

src/app/__init__.py

@@ -0,0 +1 @@
+from .multipage import MultiPage

src/app/compare_models.py

@@ -0,0 +1,292 @@
+from pathlib import Path
+import math
+
+import streamlit as st
+import numpy as np
+
+import torch
+import torch.nn.functional as F
+
+import src.app.params as params
+from src.app.questions import q1, q1_options, q2, q2_options, q3, q3_options, q4, q4_options, q5, q5_options, \
+    q6, q6_options, q7, q7_options, q8, q8_options, q9, q9_options, q10, q10_options, q11, q11_options
+from src.models import ConditionalGenerator as InfoSCC_GAN
+from src.models.big.BigGAN2 import Generator as BigGAN2Generator
+from src.models import ConditionalDecoder as cVAE
+from src.data import get_labels_train, make_galaxy_labels_hierarchical
+from src.utils import download_file, sample_labels
+
+
+device = params.device
+bs = 10  # number of images to generate each model
+n_cols = int(math.sqrt(bs))
+size = params.size
+n_layers = int(math.log2(size) - 2)
+
+# manual labels
+q1_out = [0] * len(q1_options)
+q2_out = [0] * len(q2_options)
+q3_out = [0] * len(q3_options)
+q4_out = [0] * len(q4_options)
+q5_out = [0] * len(q5_options)
+q6_out = [0] * len(q6_options)
+q7_out = [0] * len(q7_options)
+q8_out = [0] * len(q8_options)
+q9_out = [0] * len(q9_options)
+q10_out = [0] * len(q10_options)
+q11_out = [0] * len(q11_options)
+
+
+def clear_out(elems=None):
+    global q1_out, q2_out, q3_out, q4_out, q5_out, q6_out, q6_out, q7_out, q8_out, q9_out, q10_out, q11_out
+
+    if elems is None:
+        elems = list(range(1, 12))
+
+    if 1 in elems:
+        q1_out = [0] * len(q1_options)
+    if 2 in elems:
+        q2_out = [0] * len(q2_options)
+    if 3 in elems:
+        q3_out = [0] * len(q3_options)
+    if 4 in elems:
+        q4_out = [0] * len(q4_options)
+    if 5 in elems:
+        q5_out = [0] * len(q5_options)
+    if 6 in elems:
+        q6_out = [0] * len(q6_options)
+    if 7 in elems:
+        q7_out = [0] * len(q7_options)
+    if 8 in elems:
+        q8_out = [0] * len(q8_options)
+    if 9 in elems:
+        q9_out = [0] * len(q9_options)
+    if 10 in elems:
+        q10_out = [0] * len(q10_options)
+    if 11 in elems:
+        q11_out = [0] * len(q11_options)
+
+
+@st.cache(allow_output_mutation=True)
+def load_model(model_type: str):
+
+    print(f'Loading model: {model_type}')
+    if model_type == 'InfoSCC-GAN':
+        g = InfoSCC_GAN(size=params.size,
+                        y_size=params.shape_label,
+                        z_size=params.noise_dim)
+
+        if not Path(params.path_infoscc_gan).exists():
+            download_file(params.drive_id_infoscc_gan, params.path_infoscc_gan)
+
+        ckpt = torch.load(params.path_infoscc_gan, map_location=torch.device('cpu'))
+        g.load_state_dict(ckpt['g_ema'])
+    elif model_type == 'BigGAN':
+        g = BigGAN2Generator()
+
+        if not Path(params.path_biggan).exists():
+            download_file(params.drive_id_biggan, params.path_biggan)
+
+        ckpt = torch.load(params.path_biggan, map_location=torch.device('cpu'))
+        g.load_state_dict(ckpt)
+    elif model_type == 'cVAE':
+        g = cVAE()
+
+        if not Path(params.path_cvae).exists():
+            download_file(params.drive_id_cvae, params.path_cvae)
+
+        ckpt = torch.load(params.path_cvae, map_location=torch.device('cpu'))
+        g.load_state_dict(ckpt)
+    else:
+        raise ValueError('Unsupported model')
+    g = g.eval().to(device=params.device)
+    return g
+
+
+@st.cache
+def get_labels() -> torch.Tensor:
+    path_labels = params.path_labels
+
+    if not Path(path_labels).exists():
+        download_file(params.drive_id_labels, path_labels)
+
+    labels_train = get_labels_train(path_labels)
+    return labels_train
+
+
+def get_eps(n: int) -> torch.Tensor:
+    eps = torch.randn((n, params.dim_z), device=device)
+    return eps
+
+
+def app():
+    global q1_out, q2_out, q3_out, q4_out, q5_out, q6_out, q6_out, q7_out, q8_out, q9_out, q10_out, q11_out
+
+    st.title('Compare models')
+    st.markdown('This demo allows to compare BigGAN, InfoSCC-GAN and cVAE models for conditional galaxy generation.')
+    st.markdown('In each there there are images generated with the same labels by each of the models')
+
+    biggan = load_model('BigGAN')
+    infoscc_gan = load_model('InfoSCC-GAN')
+    cvae = load_model('cVAE')
+    labels_train = get_labels()
+
+    eps = get_eps(bs)  # for BigGAN and cVAE
+    eps_infoscc = infoscc_gan.sample_eps(bs)
+
+    zs = np.array([[0.0] * params.n_basis] * n_layers, dtype=np.float32)
+    zs_torch = torch.from_numpy(zs).unsqueeze(0).repeat(bs, 1, 1).to(device)
+
+    # ========================== Labels ================================
+    st.subheader('Label')
+    st.markdown(r'There are two types of selecting labels: __Random__ - sample random samples from the dataset;'
+                r' __Manual__ - select labels manually (advanced use). When using __Manual__ all of the images will be'
+                r' generated with tha same labels')
+    label_type = st.radio('Label type', options=['Random', 'Manual (Advanced)'])
+    if label_type == 'Random':
+        labels = sample_labels(labels_train, bs).to(device)
+
+        st.markdown(r'Click on __Sample labels__ button to sample random input labels')
+        change_label = st.button('Sample label')
+
+        if change_label:
+            labels = sample_labels(labels_train, bs).to(device)
+    elif label_type == 'Manual (Advanced)':
+        st.markdown('Answer the questions below')
+
+        q1_select_box = st.selectbox(q1, options=q1_options)
+        clear_out()
+        q1_out[q1_options.index(q1_select_box)] = 1
+        # 1
+
+        if q1_select_box == 'Smooth':
+            q7_select_box = st.selectbox(q7, options=q7_options)
+            clear_out([2, 3, 4, 5, 6, 7, 8, 9, 10, 11])
+            q7_out[q7_options.index(q7_select_box)] = 1
+            # 1 - 7
+
+            q6_select_box = st.selectbox(q6, options=q6_options)
+            clear_out([2, 3, 4, 5, 6, 8, 9, 10, 11])
+            q6_out[q6_options.index(q6_select_box)] = 1
+            # 1 - 7 - 6
+
+            if q6_select_box == 'Yes':
+                q8_select_box = st.selectbox(q8, options=q8_options)
+                clear_out([2, 3, 4, 5, 8, 9, 10, 11])
+                q8_out[q8_options.index(q8_select_box)] = 1
+                # 1 - 7 - 6 - 8 - end
+
+        elif q1_select_box == 'Features or disk':
+            q2_select_box = st.selectbox(q2, options=q2_options)
+            clear_out([2, 3, 4, 5, 6, 7, 8, 9, 10, 11])
+            q2_out[q2_options.index(q2_select_box)] = 1
+            # 1 - 2
+
+            if q2_select_box == 'Yes':
+                q9_select_box = st.selectbox(q9, options=q9_options)
+                clear_out([3, 4, 5, 6, 7, 8, 9, 10, 11])
+                q9_out[q9_options.index(q9_select_box)] = 1
+                # 1 - 2 - 9
+
+                q6_select_box = st.selectbox(q6, options=q6_options)
+                clear_out([3, 4, 5, 6, 7, 8, 10, 11])
+                q6_out[q6_options.index(q6_select_box)] = 1
+                # 1 - 2 - 9 - 6
+
+                if q6_select_box == 'Yes':
+                    q8_select_box = st.selectbox(q8, options=q8_options)
+                    clear_out([3, 4, 5, 7, 8, 10, 11])
+                    q8_out[q8_options.index(q8_select_box)] = 1
+                    # 1 - 2 - 9 - 6 - 8
+            else:
+                q3_select_box = st.selectbox(q3, options=q3_options)
+                clear_out([3, 4, 5, 6, 7, 8, 9, 10, 11])
+                q3_out[q3_options.index(q3_select_box)] = 1
+                # 1 - 2 - 3
+
+                q4_select_box = st.selectbox(q4, options=q4_options)
+                clear_out([4, 5, 6, 7, 8, 9, 10, 11])
+                q4_out[q4_options.index(q4_select_box)] = 1
+                # 1 - 2 - 3 - 4
+
+                if q4_select_box == 'Yes':
+                    q10_select_box = st.selectbox(q10, options=q10_options)
+                    clear_out([5, 6, 7, 8, 9, 10, 11])
+                    q10_out[q10_options.index(q10_select_box)] = 1
+                    # 1 - 2 - 3 - 4 - 10
+
+                    q11_select_box = st.selectbox(q11, options=q11_options)
+                    clear_out([5, 6, 7, 8, 9, 11])
+                    q11_out[q11_options.index(q11_select_box)] = 1
+                    # 1 - 2 - 3 - 4 - 10 - 11
+
+                    q5_select_box = st.selectbox(q5, options=q5_options)
+                    clear_out([5, 6, 7, 8, 9])
+                    q5_out[q5_options.index(q5_select_box)] = 1
+                    # 1 - 2 - 3 - 4 - 10 - 11 - 5
+
+                    q6_select_box = st.selectbox(q6, options=q6_options)
+                    clear_out([6, 7, 8, 9])
+                    q6_out[q6_options.index(q6_select_box)] = 1
+                    # 1 - 2 - 3 - 4 - 10 - 11 - 5 - 6
+
+                    if q6_select_box == 'Yes':
+                        q8_select_box = st.selectbox(q8, options=q8_options)
+                        clear_out([7, 8, 9])
+                        q8_out[q8_options.index(q8_select_box)] = 1
+                        # 1 - 2 - 3 - 4 - 10 - 11 - 5 - 6 - 8 - End
+                else:
+                    q5_select_box = st.selectbox(q5, options=q5_options)
+                    clear_out([5, 6, 7, 8, 9, 10, 11])
+                    q5_out[q5_options.index(q5_select_box)] = 1
+                    # 1 - 2 - 3 - 4 - 5
+
+                    q6_select_box = st.selectbox(q6, options=q6_options)
+                    clear_out([6, 7, 8, 9, 10, 11])
+                    q6_out[q6_options.index(q6_select_box)] = 1
+                    # 1 - 2 - 3 - 4 - 5 - 6
+
+                    if q6_select_box == 'Yes':
+                        q8_select_box = st.selectbox(q8, options=q8_options)
+                        clear_out([7, 8, 9, 10, 11])
+                        q8_out[q8_options.index(q8_select_box)] = 1
+                        # 1 - 2 - 3 - 4 - 5 - 6 - 8 - End
+
+        labels = [*q1_out, *q2_out, *q3_out, *q4_out, *q5_out, *q6_out, *q7_out, *q8_out, *q9_out, *q10_out, *q11_out]
+        labels = torch.Tensor(labels).to(device)
+        labels = labels.unsqueeze(0).repeat(bs, 1)
+        labels = make_galaxy_labels_hierarchical(labels)
+        clear_out()
+    # ========================== Labels ================================
+
+    st.subheader('Noise')
+    st.markdown(r'Click on __Change eps__ button to change input $\varepsilon$ latent space')
+    change_eps = st.button('Change eps')
+    if change_eps:
+        eps = get_eps(bs)  # for BigGAN and cVAE
+        eps_infoscc = infoscc_gan.sample_eps(bs)
+
+    with torch.no_grad():
+        imgs_biggan = biggan(eps, labels).squeeze(0).cpu()
+        imgs_infoscc = infoscc_gan(labels, eps_infoscc, zs_torch).squeeze(0).cpu()
+        imgs_cvae = cvae(eps, labels).squeeze(0).cpu()
+
+    if params.upsample:
+        imgs_biggan = F.interpolate(imgs_biggan, (size * 4, size * 4), mode='bicubic')
+        imgs_infoscc = F.interpolate(imgs_infoscc, (size * 4, size * 4), mode='bicubic')
+        imgs_cvae = F.interpolate(imgs_cvae, (size * 4, size * 4), mode='bicubic')
+
+    imgs_biggan = torch.clip(imgs_biggan, 0, 1)
+    imgs_biggan = [(imgs_biggan[i].permute(1, 2, 0).numpy() * 255).astype(np.uint8) for i in range(bs)]
+    imgs_infoscc = [(imgs_infoscc[i].permute(1, 2, 0).numpy() * 127.5 + 127.5).astype(np.uint8) for i in range(bs)]
+    imgs_cvae = [(imgs_cvae[i].permute(1, 2, 0).numpy() * 127.5 + 127.5).astype(np.uint8) for i in range(bs)]
+
+    c1, c2, c3 = st.columns(3)
+    c1.header('BigGAN')
+    c1.image(imgs_biggan, use_column_width=True)
+
+    c2.header('InfoSCC-GAN')
+    c2.image(imgs_infoscc, use_column_width=True)
+
+    c3.header('cVAE')
+    c3.image(imgs_cvae, use_column_width=True)

src/app/explore_biggan.py

@@ -0,0 +1,256 @@
+import math
+from pathlib import Path
+
+import streamlit as st
+import numpy as np
+
+import torch
+import torch.nn.functional as F
+
+import src.app.params as params
+from src.app.questions import q1, q1_options, q2, q2_options, q3, q3_options, q4, q4_options, q5, q5_options, \
+    q6, q6_options, q7, q7_options, q8, q8_options, q9, q9_options, q10, q10_options, q11, q11_options
+from src.models.big.BigGAN2 import Generator as BigGAN2Generator
+from src.data import get_labels_train, make_galaxy_labels_hierarchical
+from src.utils import download_file, sample_labels
+
+
+# global parameters
+device = params.device
+size = params.size
+y_size = shape_label = params.shape_label
+n_channels = params.n_channels
+upsample = params.upsample
+dim_z = params.dim_z
+bs = 16  # number of samples to generate
+n_cols = int(math.sqrt(bs))
+model_path = params.path_biggan
+drive_id = params.drive_id_biggan
+path_labels = params.path_labels
+
+# manual labels
+q1_out = [0] * len(q1_options)
+q2_out = [0] * len(q2_options)
+q3_out = [0] * len(q3_options)
+q4_out = [0] * len(q4_options)
+q5_out = [0] * len(q5_options)
+q6_out = [0] * len(q6_options)
+q7_out = [0] * len(q7_options)
+q8_out = [0] * len(q8_options)
+q9_out = [0] * len(q9_options)
+q10_out = [0] * len(q10_options)
+q11_out = [0] * len(q11_options)
+
+
+def clear_out(elems=None):
+    global q1_out, q2_out, q3_out, q4_out, q5_out, q6_out, q6_out, q7_out, q8_out, q9_out, q10_out, q11_out
+
+    if elems is None:
+        elems = list(range(1, 12))
+
+    if 1 in elems:
+        q1_out = [0] * len(q1_options)
+    if 2 in elems:
+        q2_out = [0] * len(q2_options)
+    if 3 in elems:
+        q3_out = [0] * len(q3_options)
+    if 4 in elems:
+        q4_out = [0] * len(q4_options)
+    if 5 in elems:
+        q5_out = [0] * len(q5_options)
+    if 6 in elems:
+        q6_out = [0] * len(q6_options)
+    if 7 in elems:
+        q7_out = [0] * len(q7_options)
+    if 8 in elems:
+        q8_out = [0] * len(q8_options)
+    if 9 in elems:
+        q9_out = [0] * len(q9_options)
+    if 10 in elems:
+        q10_out = [0] * len(q10_options)
+    if 11 in elems:
+        q11_out = [0] * len(q11_options)
+
+
+@st.cache(allow_output_mutation=True)
+def load_model(model_path: str) -> BigGAN2Generator:
+
+    print(f'Loading model: {model_path}')
+    g = BigGAN2Generator()
+    ckpt = torch.load(model_path, map_location=torch.device('cpu'))
+    g.load_state_dict(ckpt)
+    g.eval().to(device)
+    return g
+
+
+def get_eps(n: int) -> torch.Tensor:
+    eps = torch.randn((n, dim_z), device=device)
+    return eps
+
+
+@st.cache
+def get_labels() -> torch.Tensor:
+    if not Path(path_labels).exists():
+        download_file(params.drive_id_labels, path_labels)
+
+    labels_train = get_labels_train(path_labels)
+    return labels_train
+
+
+def app():
+    global q1_out, q2_out, q3_out, q4_out, q5_out, q6_out, q6_out, q7_out, q8_out, q9_out, q10_out, q11_out
+
+    st.title('Explore BigGAN')
+    st.markdown('This demo shows BigGAN for conditional galaxy generation')
+
+    if not Path(model_path).exists():
+        download_file(drive_id, model_path)
+
+    model = load_model(model_path)
+    eps = get_eps(bs)
+    labels_train = get_labels()
+
+    # ========================== Labels ================================
+    st.subheader('Label')
+    st.markdown(r'There are two types of selecting labels: __Random__ - sample random samples from the dataset;'
+                r' __Manual__ - select labels manually (advanced use). When using __Manual__ all of the images will be'
+                r' generated with tha same labels')
+    label_type = st.radio('Label type', options=['Random', 'Manual (Advanced)'])
+    if label_type == 'Random':
+        labels = sample_labels(labels_train, bs).to(device)
+
+        st.markdown(r'Click on __Sample labels__ button to sample random input labels')
+        change_label = st.button('Sample label')
+
+        if change_label:
+            labels = sample_labels(labels_train, bs).to(device)
+    elif label_type == 'Manual (Advanced)':
+        st.markdown('Answer the questions below')
+
+        q1_select_box = st.selectbox(q1, options=q1_options)
+        clear_out()
+        q1_out[q1_options.index(q1_select_box)] = 1
+        # 1
+
+        if q1_select_box == 'Smooth':
+            q7_select_box = st.selectbox(q7, options=q7_options)
+            clear_out([2, 3, 4, 5, 6, 7, 8, 9, 10, 11])
+            q7_out[q7_options.index(q7_select_box)] = 1
+            # 1 - 7
+
+            q6_select_box = st.selectbox(q6, options=q6_options)
+            clear_out([2, 3, 4, 5, 6, 8, 9, 10, 11])
+            q6_out[q6_options.index(q6_select_box)] = 1
+            # 1 - 7 - 6
+
+            if q6_select_box == 'Yes':
+                q8_select_box = st.selectbox(q8, options=q8_options)
+                clear_out([2, 3, 4, 5, 8, 9, 10, 11])
+                q8_out[q8_options.index(q8_select_box)] = 1
+                # 1 - 7 - 6 - 8 - end
+
+        elif q1_select_box == 'Features or disk':
+            q2_select_box = st.selectbox(q2, options=q2_options)
+            clear_out([2, 3, 4, 5, 6, 7, 8, 9, 10, 11])
+            q2_out[q2_options.index(q2_select_box)] = 1
+            # 1 - 2
+
+            if q2_select_box == 'Yes':
+                q9_select_box = st.selectbox(q9, options=q9_options)
+                clear_out([3, 4, 5, 6, 7, 8, 9, 10, 11])
+                q9_out[q9_options.index(q9_select_box)] = 1
+                # 1 - 2 - 9
+
+                q6_select_box = st.selectbox(q6, options=q6_options)
+                clear_out([3, 4, 5, 6, 7, 8, 10, 11])
+                q6_out[q6_options.index(q6_select_box)] = 1
+                # 1 - 2 - 9 - 6
+
+                if q6_select_box == 'Yes':
+                    q8_select_box = st.selectbox(q8, options=q8_options)
+                    clear_out([3, 4, 5, 7, 8, 10, 11])
+                    q8_out[q8_options.index(q8_select_box)] = 1
+                    # 1 - 2 - 9 - 6 - 8
+            else:
+                q3_select_box = st.selectbox(q3, options=q3_options)
+                clear_out([3, 4, 5, 6, 7, 8, 9, 10, 11])
+                q3_out[q3_options.index(q3_select_box)] = 1
+                # 1 - 2 - 3
+
+                q4_select_box = st.selectbox(q4, options=q4_options)
+                clear_out([4, 5, 6, 7, 8, 9, 10, 11])
+                q4_out[q4_options.index(q4_select_box)] = 1
+                # 1 - 2 - 3 - 4
+
+                if q4_select_box == 'Yes':
+                    q10_select_box = st.selectbox(q10, options=q10_options)
+                    clear_out([5, 6, 7, 8, 9, 10, 11])
+                    q10_out[q10_options.index(q10_select_box)] = 1
+                    # 1 - 2 - 3 - 4 - 10
+
+                    q11_select_box = st.selectbox(q11, options=q11_options)
+                    clear_out([5, 6, 7, 8, 9, 11])
+                    q11_out[q11_options.index(q11_select_box)] = 1
+                    # 1 - 2 - 3 - 4 - 10 - 11
+
+                    q5_select_box = st.selectbox(q5, options=q5_options)
+                    clear_out([5, 6, 7, 8, 9])
+                    q5_out[q5_options.index(q5_select_box)] = 1
+                    # 1 - 2 - 3 - 4 - 10 - 11 - 5
+
+                    q6_select_box = st.selectbox(q6, options=q6_options)
+                    clear_out([6, 7, 8, 9])
+                    q6_out[q6_options.index(q6_select_box)] = 1
+                    # 1 - 2 - 3 - 4 - 10 - 11 - 5 - 6
+
+                    if q6_select_box == 'Yes':
+                        q8_select_box = st.selectbox(q8, options=q8_options)
+                        clear_out([7, 8, 9])
+                        q8_out[q8_options.index(q8_select_box)] = 1
+                        # 1 - 2 - 3 - 4 - 10 - 11 - 5 - 6 - 8 - End
+                else:
+                    q5_select_box = st.selectbox(q5, options=q5_options)
+                    clear_out([5, 6, 7, 8, 9, 10, 11])
+                    q5_out[q5_options.index(q5_select_box)] = 1
+                    # 1 - 2 - 3 - 4 - 5
+
+                    q6_select_box = st.selectbox(q6, options=q6_options)
+                    clear_out([6, 7, 8, 9, 10, 11])
+                    q6_out[q6_options.index(q6_select_box)] = 1
+                    # 1 - 2 - 3 - 4 - 5 - 6
+
+                    if q6_select_box == 'Yes':
+                        q8_select_box = st.selectbox(q8, options=q8_options)
+                        clear_out([7, 8, 9, 10, 11])
+                        q8_out[q8_options.index(q8_select_box)] = 1
+                        # 1 - 2 - 3 - 4 - 5 - 6 - 8 - End
+
+        labels = [*q1_out, *q2_out, *q3_out, *q4_out, *q5_out, *q6_out, *q7_out, *q8_out, *q9_out, *q10_out, *q11_out]
+        labels = torch.Tensor(labels).to(device)
+        labels = labels.unsqueeze(0).repeat(bs, 1)
+        labels = make_galaxy_labels_hierarchical(labels)
+        clear_out()
+    # ========================== Labels ================================
+
+    st.subheader('Noise')
+    st.markdown(r'Click on __Change eps__ button to change input $\varepsilon$ latent space')
+    change_eps = st.button('Change eps')
+    if change_eps:
+        eps = get_eps(bs)
+
+    with torch.no_grad():
+        imgs = model(eps, labels)
+
+    if upsample:
+        imgs = F.interpolate(imgs, (size * 4, size * 4), mode='bicubic')
+
+    imgs = torch.clip(imgs, 0, 1)
+    imgs = [(imgs[i].permute(1, 2, 0).numpy() * 255).astype(np.uint8) for i in range(bs)]
+
+    counter = 0
+    for r in range(bs // n_cols):
+        cols = st.columns(n_cols)
+
+        for c in range(n_cols):
+            cols[c].image(imgs[counter])
+            counter += 1

src/app/explore_cvae.py

@@ -0,0 +1,255 @@
+import math
+from pathlib import Path
+
+import streamlit as st
+import numpy as np
+
+import torch
+import torch.nn.functional as F
+
+import src.app.params as params
+from src.app.questions import q1, q1_options, q2, q2_options, q3, q3_options, q4, q4_options, q5, q5_options, \
+    q6, q6_options, q7, q7_options, q8, q8_options, q9, q9_options, q10, q10_options, q11, q11_options
+from src.models import ConditionalDecoder
+from src.data import get_labels_train, make_galaxy_labels_hierarchical
+from src.utils import download_file, sample_labels
+
+
+# global parameters
+device = params.device
+size = params.size
+y_size = shape_label = params.shape_label
+n_channels = params.n_channels
+upsample = params.upsample
+dim_z = params.dim_z
+bs = 16  # number of samples to generate
+n_cols = int(math.sqrt(bs))
+model_path = params.path_cvae
+drive_id = params.drive_id_cvae
+path_labels = params.path_labels
+
+# manual labels
+q1_out = [0] * len(q1_options)
+q2_out = [0] * len(q2_options)
+q3_out = [0] * len(q3_options)
+q4_out = [0] * len(q4_options)
+q5_out = [0] * len(q5_options)
+q6_out = [0] * len(q6_options)
+q7_out = [0] * len(q7_options)
+q8_out = [0] * len(q8_options)
+q9_out = [0] * len(q9_options)
+q10_out = [0] * len(q10_options)
+q11_out = [0] * len(q11_options)
+
+
+def clear_out(elems=None):
+    global q1_out, q2_out, q3_out, q4_out, q5_out, q6_out, q6_out, q7_out, q8_out, q9_out, q10_out, q11_out
+
+    if elems is None:
+        elems = list(range(1, 12))
+
+    if 1 in elems:
+        q1_out = [0] * len(q1_options)
+    if 2 in elems:
+        q2_out = [0] * len(q2_options)
+    if 3 in elems:
+        q3_out = [0] * len(q3_options)
+    if 4 in elems:
+        q4_out = [0] * len(q4_options)
+    if 5 in elems:
+        q5_out = [0] * len(q5_options)
+    if 6 in elems:
+        q6_out = [0] * len(q6_options)
+    if 7 in elems:
+        q7_out = [0] * len(q7_options)
+    if 8 in elems:
+        q8_out = [0] * len(q8_options)
+    if 9 in elems:
+        q9_out = [0] * len(q9_options)
+    if 10 in elems:
+        q10_out = [0] * len(q10_options)
+    if 11 in elems:
+        q11_out = [0] * len(q11_options)
+
+
+@st.cache(allow_output_mutation=True)
+def load_model(model_path: str) -> ConditionalDecoder:
+
+    print(f'Loading model: {model_path}')
+    g = ConditionalDecoder()
+    ckpt = torch.load(model_path, map_location=torch.device('cpu'))
+    g.load_state_dict(ckpt)
+    g.eval().to(device)
+    return g
+
+
+def get_eps(n: int) -> torch.Tensor:
+    eps = torch.randn((n, dim_z), device=device)
+    return eps
+
+
+@st.cache
+def get_labels() -> torch.Tensor:
+    if not Path(path_labels).exists():
+        download_file(params.drive_id_labels, path_labels)
+
+    labels_train = get_labels_train(path_labels)
+    return labels_train
+
+
+def app():
+    global q1_out, q2_out, q3_out, q4_out, q5_out, q6_out, q6_out, q7_out, q8_out, q9_out, q10_out, q11_out
+
+    st.title('Explore cVAE')
+    st.markdown('This demo shows cVAE for conditional galaxy generation')
+
+    if not Path(model_path).exists():
+        download_file(drive_id, model_path)
+
+    model = load_model(model_path)
+    eps = get_eps(bs)
+    labels_train = get_labels()
+
+    # ========================== Labels ================================
+    st.subheader('Label')
+    st.markdown(r'There are two types of selecting labels: __Random__ - sample random samples from the dataset;'
+                r' __Manual__ - select labels manually (advanced use). When using __Manual__ all of the images will be'
+                r' generated with tha same labels')
+    label_type = st.radio('Label type', options=['Random', 'Manual (Advanced)'])
+    if label_type == 'Random':
+        labels = sample_labels(labels_train, bs).to(device)
+
+        st.markdown(r'Click on __Sample labels__ button to sample random input labels')
+        change_label = st.button('Sample label')
+
+        if change_label:
+            labels = sample_labels(labels_train, bs).to(device)
+    elif label_type == 'Manual (Advanced)':
+        st.markdown('Answer the questions below')
+
+        q1_select_box = st.selectbox(q1, options=q1_options)
+        clear_out()
+        q1_out[q1_options.index(q1_select_box)] = 1
+        # 1
+
+        if q1_select_box == 'Smooth':
+            q7_select_box = st.selectbox(q7, options=q7_options)
+            clear_out([2, 3, 4, 5, 6, 7, 8, 9, 10, 11])
+            q7_out[q7_options.index(q7_select_box)] = 1
+            # 1 - 7
+
+            q6_select_box = st.selectbox(q6, options=q6_options)
+            clear_out([2, 3, 4, 5, 6, 8, 9, 10, 11])
+            q6_out[q6_options.index(q6_select_box)] = 1
+            # 1 - 7 - 6
+
+            if q6_select_box == 'Yes':
+                q8_select_box = st.selectbox(q8, options=q8_options)
+                clear_out([2, 3, 4, 5, 8, 9, 10, 11])
+                q8_out[q8_options.index(q8_select_box)] = 1
+                # 1 - 7 - 6 - 8 - end
+
+        elif q1_select_box == 'Features or disk':
+            q2_select_box = st.selectbox(q2, options=q2_options)
+            clear_out([2, 3, 4, 5, 6, 7, 8, 9, 10, 11])
+            q2_out[q2_options.index(q2_select_box)] = 1
+            # 1 - 2
+
+            if q2_select_box == 'Yes':
+                q9_select_box = st.selectbox(q9, options=q9_options)
+                clear_out([3, 4, 5, 6, 7, 8, 9, 10, 11])
+                q9_out[q9_options.index(q9_select_box)] = 1
+                # 1 - 2 - 9
+
+                q6_select_box = st.selectbox(q6, options=q6_options)
+                clear_out([3, 4, 5, 6, 7, 8, 10, 11])
+                q6_out[q6_options.index(q6_select_box)] = 1
+                # 1 - 2 - 9 - 6
+
+                if q6_select_box == 'Yes':
+                    q8_select_box = st.selectbox(q8, options=q8_options)
+                    clear_out([3, 4, 5, 7, 8, 10, 11])
+                    q8_out[q8_options.index(q8_select_box)] = 1
+                    # 1 - 2 - 9 - 6 - 8
+            else:
+                q3_select_box = st.selectbox(q3, options=q3_options)
+                clear_out([3, 4, 5, 6, 7, 8, 9, 10, 11])
+                q3_out[q3_options.index(q3_select_box)] = 1
+                # 1 - 2 - 3
+
+                q4_select_box = st.selectbox(q4, options=q4_options)
+                clear_out([4, 5, 6, 7, 8, 9, 10, 11])
+                q4_out[q4_options.index(q4_select_box)] = 1
+                # 1 - 2 - 3 - 4
+
+                if q4_select_box == 'Yes':
+                    q10_select_box = st.selectbox(q10, options=q10_options)
+                    clear_out([5, 6, 7, 8, 9, 10, 11])
+                    q10_out[q10_options.index(q10_select_box)] = 1
+                    # 1 - 2 - 3 - 4 - 10
+
+                    q11_select_box = st.selectbox(q11, options=q11_options)
+                    clear_out([5, 6, 7, 8, 9, 11])
+                    q11_out[q11_options.index(q11_select_box)] = 1
+                    # 1 - 2 - 3 - 4 - 10 - 11
+
+                    q5_select_box = st.selectbox(q5, options=q5_options)
+                    clear_out([5, 6, 7, 8, 9])
+                    q5_out[q5_options.index(q5_select_box)] = 1
+                    # 1 - 2 - 3 - 4 - 10 - 11 - 5
+
+                    q6_select_box = st.selectbox(q6, options=q6_options)
+                    clear_out([6, 7, 8, 9])
+                    q6_out[q6_options.index(q6_select_box)] = 1
+                    # 1 - 2 - 3 - 4 - 10 - 11 - 5 - 6
+
+                    if q6_select_box == 'Yes':
+                        q8_select_box = st.selectbox(q8, options=q8_options)
+                        clear_out([7, 8, 9])
+                        q8_out[q8_options.index(q8_select_box)] = 1
+                        # 1 - 2 - 3 - 4 - 10 - 11 - 5 - 6 - 8 - End
+                else:
+                    q5_select_box = st.selectbox(q5, options=q5_options)
+                    clear_out([5, 6, 7, 8, 9, 10, 11])
+                    q5_out[q5_options.index(q5_select_box)] = 1
+                    # 1 - 2 - 3 - 4 - 5
+
+                    q6_select_box = st.selectbox(q6, options=q6_options)
+                    clear_out([6, 7, 8, 9, 10, 11])
+                    q6_out[q6_options.index(q6_select_box)] = 1
+                    # 1 - 2 - 3 - 4 - 5 - 6
+
+                    if q6_select_box == 'Yes':
+                        q8_select_box = st.selectbox(q8, options=q8_options)
+                        clear_out([7, 8, 9, 10, 11])
+                        q8_out[q8_options.index(q8_select_box)] = 1
+                        # 1 - 2 - 3 - 4 - 5 - 6 - 8 - End
+
+        labels = [*q1_out, *q2_out, *q3_out, *q4_out, *q5_out, *q6_out, *q7_out, *q8_out, *q9_out, *q10_out, *q11_out]
+        labels = torch.Tensor(labels).to(device)
+        labels = labels.unsqueeze(0).repeat(bs, 1)
+        labels = make_galaxy_labels_hierarchical(labels)
+        clear_out()
+    # ========================== Labels ================================
+
+    st.subheader('Noise')
+    st.markdown(r'Click on __Change eps__ button to change input $\varepsilon$ latent space')
+    change_eps = st.button('Change eps')
+    if change_eps:
+        eps = get_eps(bs)
+
+    with torch.no_grad():
+        imgs = model(eps, labels)
+
+    if upsample:
+        imgs = F.interpolate(imgs, (size * 4, size * 4), mode='bicubic')
+
+    imgs = [(imgs[i].permute(1, 2, 0).numpy() * 127.5 + 127.5).astype(np.uint8) for i in range(bs)]
+
+    counter = 0
+    for r in range(bs // n_cols):
+        cols = st.columns(n_cols)
+
+        for c in range(n_cols):
+            cols[c].image(imgs[counter])
+            counter += 1

src/app/explore_infoscc_gan.py

@@ -0,0 +1,263 @@
+from pathlib import Path
+import math
+
+import numpy as np
+import streamlit as st
+
+import torch
+import torch.nn.functional as F
+
+import src.app.params as params
+from src.app.questions import q1, q1_options, q2, q2_options, q3, q3_options, q4, q4_options, q5, q5_options, \
+    q6, q6_options, q7, q7_options, q8, q8_options, q9, q9_options, q10, q10_options, q11, q11_options
+from src.models import ConditionalGenerator
+from src.data import get_labels_train, make_galaxy_labels_hierarchical
+from src.utils import download_file, sample_labels
+
+# global parameters
+device = params.device
+size = params.size
+y_size = params.shape_label
+n_channels = params.n_channels
+upsample = params.upsample
+z_size = noise_dim = params.noise_dim
+n_layers = int(math.log2(size) - 2)
+n_basis = params.n_basis
+y_type = params.y_type
+bs = 16  # number of samples to generate
+n_cols = int(math.sqrt(bs))
+model_path = params.path_infoscc_gan  # path to the model
+drive_id = params.drive_id_infoscc_gan   # google drive id of the model
+path_labels = params.path_labels
+
+# manual labels
+q1_out = [0] * len(q1_options)
+q2_out = [0] * len(q2_options)
+q3_out = [0] * len(q3_options)
+q4_out = [0] * len(q4_options)
+q5_out = [0] * len(q5_options)
+q6_out = [0] * len(q6_options)
+q7_out = [0] * len(q7_options)
+q8_out = [0] * len(q8_options)
+q9_out = [0] * len(q9_options)
+q10_out = [0] * len(q10_options)
+q11_out = [0] * len(q11_options)
+
+
+def clear_out(elems=None):
+    global q1_out, q2_out, q3_out, q4_out, q5_out, q6_out, q6_out, q7_out, q8_out, q9_out, q10_out, q11_out
+
+    if elems is None:
+        elems = list(range(1, 12))
+
+    if 1 in elems:
+        q1_out = [0] * len(q1_options)
+    if 2 in elems:
+        q2_out = [0] * len(q2_options)
+    if 3 in elems:
+        q3_out = [0] * len(q3_options)
+    if 4 in elems:
+        q4_out = [0] * len(q4_options)
+    if 5 in elems:
+        q5_out = [0] * len(q5_options)
+    if 6 in elems:
+        q6_out = [0] * len(q6_options)
+    if 7 in elems:
+        q7_out = [0] * len(q7_options)
+    if 8 in elems:
+        q8_out = [0] * len(q8_options)
+    if 9 in elems:
+        q9_out = [0] * len(q9_options)
+    if 10 in elems:
+        q10_out = [0] * len(q10_options)
+    if 11 in elems:
+        q11_out = [0] * len(q11_options)
+
+
+@st.cache(allow_output_mutation=True)
+def load_model(model_path: str) -> ConditionalGenerator:
+
+    print(f'Loading model: {model_path}')
+    g_ema = ConditionalGenerator(size, y_size, z_size, n_channels, n_basis, noise_dim)
+    ckpt = torch.load(model_path, map_location=torch.device('cpu'))
+    g_ema.load_state_dict(ckpt['g_ema'])
+    g_ema.eval().to(device)
+    return g_ema
+
+
+@st.cache
+def get_labels() -> torch.Tensor:
+    if not Path(path_labels).exists():
+        download_file(params.drive_id_labels, path_labels)
+    labels_train = get_labels_train(path_labels)
+    return labels_train
+
+
+def app():
+    global q1_out, q2_out, q3_out, q4_out, q5_out, q6_out, q6_out, q7_out, q8_out, q9_out, q10_out, q11_out
+
+    st.title('Explore InfoSCC-GAN')
+    st.markdown('This demo shows InfoSCC-GAN for conditional galaxy generation')
+    st.subheader(r'<- Use sidebar to explore $z_1, ..., z_k$ latent variables')
+
+    if not Path(model_path).exists():
+        download_file(drive_id, model_path)
+
+    model = load_model(model_path)
+    eps = model.sample_eps(bs).to(device)
+    labels_train = get_labels()
+
+    # get zs
+    zs = np.array([[0.0] * n_basis] * n_layers, dtype=np.float32)
+
+    for l in range(n_layers):
+        st.sidebar.markdown(f'## Layer: {l}')
+        for d in range(n_basis):
+            zs[l][d] = st.sidebar.slider(f'Dimension: {d}', key=f'{l}{d}',
+                                         min_value=-5., max_value=5., value=0., step=0.1)
+
+    # ========================== Labels ================================
+    st.subheader('Label')
+    st.markdown(r'There are two types of selecting labels: __Random__ - sample random samples from the dataset;'
+                r' __Manual__ - select labels manually (advanced use). When using __Manual__ all of the images will be'
+                r' generated with tha same labels')
+    label_type = st.radio('Label type', options=['Random', 'Manual (Advanced)'])
+    if label_type == 'Random':
+        labels = sample_labels(labels_train, bs).to(device)
+
+        st.markdown(r'Click on __Sample labels__ button to sample random input labels')
+        change_label = st.button('Sample label')
+
+        if change_label:
+            labels = sample_labels(labels_train, bs).to(device)
+    elif label_type == 'Manual (Advanced)':
+        st.markdown('Answer the questions below')
+
+        q1_select_box = st.selectbox(q1, options=q1_options)
+        clear_out()
+        q1_out[q1_options.index(q1_select_box)] = 1
+        # 1
+
+        if q1_select_box == 'Smooth':
+            q7_select_box = st.selectbox(q7, options=q7_options)
+            clear_out([2, 3, 4, 5, 6, 7, 8, 9, 10, 11])
+            q7_out[q7_options.index(q7_select_box)] = 1
+            # 1 - 7
+
+            q6_select_box = st.selectbox(q6, options=q6_options)
+            clear_out([2, 3, 4, 5, 6, 8, 9, 10, 11])
+            q6_out[q6_options.index(q6_select_box)] = 1
+            # 1 - 7 - 6
+
+            if q6_select_box == 'Yes':
+                q8_select_box = st.selectbox(q8, options=q8_options)
+                clear_out([2, 3, 4, 5, 8, 9, 10, 11])
+                q8_out[q8_options.index(q8_select_box)] = 1
+                # 1 - 7 - 6 - 8 - end
+
+        elif q1_select_box == 'Features or disk':
+            q2_select_box = st.selectbox(q2, options=q2_options)
+            clear_out([2, 3, 4, 5, 6, 7, 8, 9, 10, 11])
+            q2_out[q2_options.index(q2_select_box)] = 1
+            # 1 - 2
+
+            if q2_select_box == 'Yes':
+                q9_select_box = st.selectbox(q9, options=q9_options)
+                clear_out([3, 4, 5, 6, 7, 8, 9, 10, 11])
+                q9_out[q9_options.index(q9_select_box)] = 1
+                # 1 - 2 - 9
+
+                q6_select_box = st.selectbox(q6, options=q6_options)
+                clear_out([3, 4, 5, 6, 7, 8, 10, 11])
+                q6_out[q6_options.index(q6_select_box)] = 1
+                # 1 - 2 - 9 - 6
+
+                if q6_select_box == 'Yes':
+                    q8_select_box = st.selectbox(q8, options=q8_options)
+                    clear_out([3, 4, 5, 7, 8, 10, 11])
+                    q8_out[q8_options.index(q8_select_box)] = 1
+                    # 1 - 2 - 9 - 6 - 8
+            else:
+                q3_select_box = st.selectbox(q3, options=q3_options)
+                clear_out([3, 4, 5, 6, 7, 8, 9, 10, 11])
+                q3_out[q3_options.index(q3_select_box)] = 1
+                # 1 - 2 - 3
+
+                q4_select_box = st.selectbox(q4, options=q4_options)
+                clear_out([4, 5, 6, 7, 8, 9, 10, 11])
+                q4_out[q4_options.index(q4_select_box)] = 1
+                # 1 - 2 - 3 - 4
+
+                if q4_select_box == 'Yes':
+                    q10_select_box = st.selectbox(q10, options=q10_options)
+                    clear_out([5, 6, 7, 8, 9, 10, 11])
+                    q10_out[q10_options.index(q10_select_box)] = 1
+                    # 1 - 2 - 3 - 4 - 10
+
+                    q11_select_box = st.selectbox(q11, options=q11_options)
+                    clear_out([5, 6, 7, 8, 9, 11])
+                    q11_out[q11_options.index(q11_select_box)] = 1
+                    # 1 - 2 - 3 - 4 - 10 - 11
+
+                    q5_select_box = st.selectbox(q5, options=q5_options)
+                    clear_out([5, 6, 7, 8, 9])
+                    q5_out[q5_options.index(q5_select_box)] = 1
+                    # 1 - 2 - 3 - 4 - 10 - 11 - 5
+
+                    q6_select_box = st.selectbox(q6, options=q6_options)
+                    clear_out([6, 7, 8, 9])
+                    q6_out[q6_options.index(q6_select_box)] = 1
+                    # 1 - 2 - 3 - 4 - 10 - 11 - 5 - 6
+
+                    if q6_select_box == 'Yes':
+                        q8_select_box = st.selectbox(q8, options=q8_options)
+                        clear_out([7, 8, 9])
+                        q8_out[q8_options.index(q8_select_box)] = 1
+                        # 1 - 2 - 3 - 4 - 10 - 11 - 5 - 6 - 8 - End
+                else:
+                    q5_select_box = st.selectbox(q5, options=q5_options)
+                    clear_out([5, 6, 7, 8, 9, 10, 11])
+                    q5_out[q5_options.index(q5_select_box)] = 1
+                    # 1 - 2 - 3 - 4 - 5
+
+                    q6_select_box = st.selectbox(q6, options=q6_options)
+                    clear_out([6, 7, 8, 9, 10, 11])
+                    q6_out[q6_options.index(q6_select_box)] = 1
+                    # 1 - 2 - 3 - 4 - 5 - 6
+
+                    if q6_select_box == 'Yes':
+                        q8_select_box = st.selectbox(q8, options=q8_options)
+                        clear_out([7, 8, 9, 10, 11])
+                        q8_out[q8_options.index(q8_select_box)] = 1
+                        # 1 - 2 - 3 - 4 - 5 - 6 - 8 - End
+
+        labels = [*q1_out, *q2_out, *q3_out, *q4_out, *q5_out, *q6_out, *q7_out, *q8_out, *q9_out, *q10_out, *q11_out]
+        labels = torch.Tensor(labels).to(device)
+        labels = labels.unsqueeze(0).repeat(bs, 1)
+        labels = make_galaxy_labels_hierarchical(labels)
+        clear_out()
+    # ========================== Labels ================================
+
+    st.subheader('Noise')
+    st.markdown(r'Click on __Change eps__ button to change input $\varepsilon$ latent space')
+    change_eps = st.button('Change eps')
+    if change_eps:
+        eps = model.sample_eps(bs).to(device)
+
+    zs_torch = torch.from_numpy(zs).unsqueeze(0).repeat(bs, 1, 1).to(device)
+
+    with torch.no_grad():
+        imgs = model(labels, eps, zs_torch).squeeze(0).cpu()
+
+    if upsample:
+        imgs = F.interpolate(imgs, (size * 4, size * 4), mode='bicubic')
+
+    imgs = [(imgs[i].permute(1, 2, 0).numpy() * 127.5 + 127.5).astype(np.uint8) for i in range(bs)]
+
+    counter = 0
+    for r in range(bs // n_cols):
+        cols = st.columns(n_cols)
+
+        for c in range(n_cols):
+            cols[c].image(imgs[counter])
+            counter += 1

src/app/multipage.py

@@ -0,0 +1,41 @@
+"""
+This file is the framework for generating multiple Streamlit applications
+through an object oriented framework.
+"""
+
+# Import necessary libraries
+import streamlit as st
+
+
+# Define the multipage class to manage the multiple apps in our program
+class MultiPage:
+    """Framework for combining multiple streamlit applications."""
+
+    def __init__(self) -> None:
+        """Constructor class to generate a list which will store all our applications as an instance variable."""
+        self.pages = []
+
+    def add_page(self, title, func) -> None:
+        """Class Method to Add pages to the project
+        Args:
+            title ([str]): The title of page which we are adding to the list of apps
+
+            func: Python function to render this page in Streamlit
+        """
+
+        self.pages.append({
+
+                "title": title,
+                "function": func
+            })
+
+    def run(self):
+        # Drodown to select the page to run
+        page = st.sidebar.selectbox(
+            'App Navigation',
+            self.pages,
+            format_func=lambda page: page['title']
+        )
+
+        # run the app function
+        page['function']()

src/app/params.py

@@ -0,0 +1,25 @@
+"""
+This file contains list of global parameters for the Galaxy Zoo generation app
+"""
+
+device = 'cpu'
+size = 64   # generated image size
+shape_label = 37  # shape of the input label
+n_channels = 3  # number of color channels in image
+upsample = True   # if true, generated images will be upsampled
+noise_dim = 512  # noise size in InfoSCC-GAN
+n_basis = 6   # size of additional z vectors in InfoSCC-GAN
+y_type = 'real'  # type of labels in InfoSCC-GAN
+dim_z = 128  # z vector size in BigGAN and cVAE
+
+path_infoscc_gan = './models/InfoSCC-GAN/generator.pt'
+drive_id_infoscc_gan = '1_kIujc497OH0ZJ7PNPwS5_otNlS7jMLI'
+
+path_biggan = './models/BigGAN/generator.pth'
+drive_id_biggan = '1sMSDdnQ5GjHcno5knHTDSKAKhhoHh_4z'
+
+path_cvae = './models/CVAE/generator.pth'
+drive_id_cvae = '17FmLvhwXq8PQMrD1CtjqyoAy5BobYMTE'
+
+path_labels = './data/training_solutions_rev1.csv'
+drive_id_labels = '1dzsB_HdGtmSHE4pCppamISpFaJBfPF7E'

src/app/questions.py

@@ -0,0 +1,36 @@
+"""
+This file contains question and options for the manual labeling
+"""
+
+q1 = 'Is the object a smooth galaxy, a galaxy with features/disk or a star?'
+q1_options = ['Smooth', 'Features or disk', 'Star or artifact']
+
+q2 = 'Is it edge-on? '
+q2_options = ['Yes', 'No']
+
+q3 = 'Is there a bar?'
+q3_options = ['Yes', 'No']
+
+q4 = 'Is there a spiral pattern?'
+q4_options = ['Yes', 'No']
+
+q5 = 'How prominent is the central bulge?'
+q5_options = ['No bulge', 'Just noticeable', 'Obvious', 'Dominant']
+
+q6 = 'Is there anything "odd" about the galaxy?'
+q6_options = ['Yes', 'No']
+
+q7 = 'How round is the smooth galaxy?'
+q7_options = ['Completely round', 'In between', 'Cigar-shaped']
+
+q8 = 'What is the odd feature?'
+q8_options = ['Ring', 'Lens or are', 'Disturbed', 'Irregular', 'Other', 'Merger', 'Dust lane']
+
+q9 = 'What shape is the bulge in the edge-on galaxy?'
+q9_options = ['Rounded', 'Boxy', 'No bulge']
+
+q10 = 'How tightly wound are the spiral arms?'
+q10_options = ['Tight', 'Medium', 'Loose']
+
+q11 = 'How many spiral arms are there?'
+q11_options = ['1', '2', '3', '4', 'more than four', 'can`t tell']

src/data/__init__.py

@@ -0,0 +1,2 @@
+from .data import get_labels_train
+from .labels import make_galaxy_labels_hierarchical

src/data/data.py

@@ -0,0 +1,10 @@
+from pandas import read_csv
+
+import torch
+
+
+def get_labels_train(file_galaxy_labels) -> torch.Tensor:
+    df_galaxy_labels = read_csv(file_galaxy_labels)
+    labels_train = df_galaxy_labels[df_galaxy_labels.columns[1:]].values
+    labels_train = torch.from_numpy(labels_train).float()
+    return labels_train

src/data/labels.py

@@ -0,0 +1,63 @@
+import numpy as np
+
+import torch
+
+
+class_groups = {
+    # group : indices (assuming 0th position is id)
+    0: (),
+    1: (1, 2, 3),
+    2: (4, 5),
+    3: (6, 7),
+    4: (8, 9),
+    5: (10, 11, 12, 13),
+    6: (14, 15),
+    7: (16, 17, 18),
+    8: (19, 20, 21, 22, 23, 24, 25),
+    9: (26, 27, 28),
+    10: (29, 30, 31),
+    11: (32, 33, 34, 35, 36, 37),
+}
+
+
+class_groups_indices = {g: np.array(ixs)-1 for g, ixs in class_groups.items()}
+
+
+hierarchy = {
+    # group : parent (group, label)
+    2: (1, 1),
+    3: (2, 1),
+    4: (2, 1),
+    5: (2, 1),
+    7: (1, 0),
+    8: (6, 0),
+    9: (2, 0),
+    10: (4, 0),
+    11: (4, 0),
+}
+
+
+def make_galaxy_labels_hierarchical(labels: torch.Tensor) -> torch.Tensor:
+    """ transform groups of galaxy label probabilities to follow the hierarchical order defined in galaxy zoo
+    more info here: https://www.kaggle.com/c/galaxy-zoo-the-galaxy-challenge/overview/the-galaxy-zoo-decision-tree
+    labels is a NxL torch tensor, where N is the batch size and L is the number of labels,
+    all labels should be > 1
+    the indices of label groups are listed in class_groups_indices
+
+    Return
+    ------
+    hierarchical_labels : NxL torch tensor, where L is the total number of labels
+    """
+    shift = labels.shape[1] > 37 ## in case the id is included at 0th position, shift indices accordingly
+    index = lambda i: class_groups_indices[i] + shift
+
+    for i in range(1, 12):
+        ## normalize probabilities to 1
+        norm = torch.sum(labels[:, index(i)], dim=1, keepdims=True)
+        norm[norm == 0] += 1e-4  ## add small number to prevent NaNs dividing by zero, yet keep track of gradient
+        labels[:, index(i)] /= norm
+        ## renormalize according to hierarchical structure
+        if i not in [1, 6]:
+            parent_group_label = labels[:, index(hierarchy[i][0])]
+            labels[:, index(i)] *= parent_group_label[:, hierarchy[i][1]].unsqueeze(-1)
+    return labels

src/models/__init__.py

@@ -0,0 +1,2 @@
+from .infoscc_gan import ConditionalGenerator
+from .cvae import ConditionalDecoder

src/models/big/BigGAN2.py

@@ -0,0 +1,469 @@
+import functools
+
+import torch
+import torch.nn as nn
+from torch.nn import init
+import torch.optim as optim
+import torch.nn.functional as F
+
+import src.models.big.layers as layers
+from src.models.parameter import labels_dim, parameter
+from src.models.neuralnetwork import NeuralNetwork
+
+
+# Architectures for G
+# Attention is passed in in the format '32_64' to mean applying an attention
+# block at both resolution 32x32 and 64x64. Just '64' will apply at 64x64.
+def G_arch(ch=64, attention='64', ksize='333333', dilation='111111'):
+  arch = {}
+  arch[512] = {'in_channels' :  [ch * item for item in [16, 16, 8, 8, 4, 2, 1]],
+               'out_channels' : [ch * item for item in [16,  8, 8, 4, 2, 1, 1]],
+               'upsample' : [True] * 7,
+               'resolution' : [8, 16, 32, 64, 128, 256, 512],
+               'attention' : {2**i: (2**i in [int(item) for item in attention.split('_')])
+                              for i in range(3,10)}}
+  arch[256] = {'in_channels' :  [ch * item for item in [16, 16, 8, 8, 4, 2]],
+               'out_channels' : [ch * item for item in [16,  8, 8, 4, 2, 1]],
+               'upsample' : [True] * 6,
+               'resolution' : [8, 16, 32, 64, 128, 256],
+               'attention' : {2**i: (2**i in [int(item) for item in attention.split('_')])
+                              for i in range(3,9)}}
+  arch[128] = {'in_channels' :  [ch * item for item in [16, 16, 8, 4, 2]],
+               'out_channels' : [ch * item for item in [16, 8, 4, 2, 1]],
+               'upsample' : [True] * 5,
+               'resolution' : [8, 16, 32, 64, 128],
+               'attention' : {2**i: (2**i in [int(item) for item in attention.split('_')])
+                              for i in range(3,8)}}
+  arch[64]  = {'in_channels' :  [ch * item for item in [16, 16, 8, 4]],
+               'out_channels' : [ch * item for item in [16, 8, 4, 2]],
+               'upsample' : [True] * 4,
+               'resolution' : [8, 16, 32, 64],
+               'attention' : {2**i: (2**i in [int(item) for item in attention.split('_')])
+                              for i in range(3,7)}}
+  arch[32]  = {'in_channels' :  [ch * item for item in [4, 4, 4]],
+               'out_channels' : [ch * item for item in [4, 4, 4]],
+               'upsample' : [True] * 3,
+               'resolution' : [8, 16, 32],
+               'attention' : {2**i: (2**i in [int(item) for item in attention.split('_')])
+                              for i in range(3,6)}}
+
+  return arch
+
+class Generator(NeuralNetwork):
+  def __init__(self, G_ch=64, dim_z=128, bottom_width=4, resolution=64, labels_dim=labels_dim,
+               G_kernel_size=3, G_attn='64', n_classes=1,
+               num_G_SVs=1, num_G_SV_itrs=1,
+               G_shared=True, shared_dim=0, hier=False,
+               cross_replica=False, mybn=False,
+               G_activation=nn.ReLU(inplace=False),
+               G_lr=5e-5, G_B1=0.0, G_B2=0.999, adam_eps=1e-8,
+               BN_eps=1e-5, SN_eps=1e-12, G_mixed_precision=False, G_fp16=False,
+               G_init='ortho', skip_init=False, no_optim=False,
+               G_param='SN', norm_style='bn',
+               **kwargs):
+    super(Generator, self).__init__()
+    # Channel width mulitplier
+    self.ch = G_ch
+    # Dimensionality of the latent space
+    self.dim_z = dim_z
+    # The initial spatial dimensions
+    self.bottom_width = bottom_width
+    # Resolution of the output
+    self.resolution = resolution
+    # Kernel size?
+    self.kernel_size = G_kernel_size
+    # Attention?
+    self.attention = G_attn
+    # number of classes, for use in categorical conditional generation
+    self.n_classes = n_classes
+    # Use shared embeddings?
+    self.G_shared = G_shared
+    # Dimensionality of the shared embedding? Unused if not using G_shared
+    self.shared_dim = shared_dim if shared_dim > 0 else dim_z
+    # Hierarchical latent space?
+    self.hier = hier
+    # Cross replica batchnorm?
+    self.cross_replica = cross_replica
+    # Use my batchnorm?
+    self.mybn = mybn
+    # nonlinearity for residual blocks
+    self.activation = G_activation
+    # Initialization style
+    self.init = G_init
+    # Parameterization style
+    self.G_param = G_param
+    # Normalization style
+    self.norm_style = norm_style
+    # Epsilon for BatchNorm?
+    self.BN_eps = BN_eps
+    # Epsilon for Spectral Norm?
+    self.SN_eps = SN_eps
+    # fp16?
+    self.fp16 = G_fp16
+    # Architecture dict
+    self.arch = G_arch(self.ch, self.attention)[resolution]
+
+    # If using hierarchical latents, adjust z
+    if self.hier:
+      # Number of places z slots into
+      self.num_slots = len(self.arch['in_channels']) + 1
+      self.z_chunk_size = (self.dim_z // self.num_slots)
+      # Recalculate latent dimensionality for even splitting into chunks
+      self.dim_z = self.z_chunk_size *  self.num_slots
+    else:
+      self.num_slots = 1
+      self.z_chunk_size = 0
+
+    # Which convs, batchnorms, and linear layers to use
+    if self.G_param == 'SN':
+      self.which_conv = functools.partial(layers.SNConv2d,
+                          kernel_size=3, padding=1,
+                          num_svs=num_G_SVs, num_itrs=num_G_SV_itrs,
+                          eps=self.SN_eps)
+      self.which_linear = functools.partial(layers.SNLinear,
+                          num_svs=num_G_SVs, num_itrs=num_G_SV_itrs,
+                          eps=self.SN_eps)
+    else:
+      self.which_conv = functools.partial(nn.Conv2d, kernel_size=3, padding=1)
+      self.which_linear = nn.Linear
+      
+    # We use a non-spectral-normed embedding here regardless;
+    # For some reason applying SN to G's embedding seems to randomly cripple G
+    self.which_embedding = nn.Embedding
+    bn_linear = (functools.partial(self.which_linear, bias=False) if self.G_shared
+                 else self.which_embedding)
+    self.which_bn = functools.partial(layers.ccbn,
+                          which_linear=bn_linear,
+                          cross_replica=self.cross_replica,
+                          mybn=self.mybn,
+                          input_size=(self.shared_dim + self.z_chunk_size if self.G_shared
+                                      else self.n_classes),
+                          norm_style=self.norm_style,
+                          eps=self.BN_eps)
+
+
+    # Prepare model
+    # prepare label input
+    self.transform_label_layer = torch.nn.Linear(labels_dim, 128)
+    # If not using shared embeddings, self.shared is just a passthrough
+    self.shared = (self.which_embedding(n_classes, self.shared_dim) if G_shared 
+                    else layers.identity())
+    # First linear layer
+    self.linear = self.which_linear(self.dim_z // self.num_slots,
+                                    self.arch['in_channels'][0] * (self.bottom_width **2))
+
+    # self.blocks is a doubly-nested list of modules, the outer loop intended
+    # to be over blocks at a given resolution (resblocks and/or self-attention)
+    # while the inner loop is over a given block
+    self.blocks = []
+    for index in range(len(self.arch['out_channels'])):
+      self.blocks += [[layers.GBlock(in_channels=self.arch['in_channels'][index],
+                             out_channels=self.arch['out_channels'][index],
+                             which_conv=self.which_conv,
+                             which_bn=self.which_bn,
+                             activation=self.activation,
+                             upsample=(functools.partial(F.interpolate, scale_factor=2)
+                                       if self.arch['upsample'][index] else None))]]
+
+      # If attention on this block, attach it to the end
+      if self.arch['attention'][self.arch['resolution'][index]]:
+        print('Adding attention layer in G at resolution %d' % self.arch['resolution'][index])
+        self.blocks[-1] += [layers.Attention(self.arch['out_channels'][index], self.which_conv)]
+
+    # Turn self.blocks into a ModuleList so that it's all properly registered.
+    self.blocks = nn.ModuleList([nn.ModuleList(block) for block in self.blocks])
+
+    # output layer: batchnorm-relu-conv.
+    # Consider using a non-spectral conv here
+    self.output_layer = nn.Sequential(layers.bn(self.arch['out_channels'][-1],
+                                                cross_replica=self.cross_replica,
+                                                mybn=self.mybn),
+                                    self.activation,
+                                    self.which_conv(self.arch['out_channels'][-1], 3))
+
+    # Initialize weights. Optionally skip init for testing.
+    if not skip_init:
+      self.init_weights()
+
+    # Set up optimizer
+    # If this is an EMA copy, no need for an optim, so just return now
+    if no_optim:
+      return
+    self.lr, self.B1, self.B2, self.adam_eps = G_lr, G_B1, G_B2, adam_eps
+    if G_mixed_precision:
+      print('Using fp16 adam in G...')
+      import utils
+      self.optim = utils.Adam16(params=self.parameters(), lr=self.lr,
+                           betas=(self.B1, self.B2), weight_decay=0,
+                           eps=self.adam_eps)
+    else:
+      self.optim = optim.Adam(params=self.parameters(), lr=self.lr,
+                           betas=(self.B1, self.B2), weight_decay=0,
+                           eps=self.adam_eps)
+
+    # LR scheduling, left here for forward compatibility
+    # self.lr_sched = {'itr' : 0}# if self.progressive else {}
+    # self.j = 0
+    self.set_optimizer(parameter.optimizer, lr=parameter.learning_rate, betas=parameter.betas)
+
+  # Initialize
+  def init_weights(self):
+    self.param_count = 0
+    for module in self.modules():
+      if (isinstance(module, nn.Conv2d) 
+          or isinstance(module, nn.Linear) 
+          or isinstance(module, nn.Embedding)):
+        if self.init == 'ortho':
+          init.orthogonal_(module.weight)
+        elif self.init == 'N02':
+          init.normal_(module.weight, 0, 0.02)
+        elif self.init in ['glorot', 'xavier']:
+          init.xavier_uniform_(module.weight)
+        else:
+          print('Init style not recognized...')
+        self.param_count += sum([p.data.nelement() for p in module.parameters()])
+    print('Param count for G''s initialized parameters: %d' % self.param_count)
+
+
+  def transform_labels(self, labels):
+    """ prepore labels for input to generator """
+    return self.transform_label_layer(labels)
+
+
+  # Note on this forward function: we pass in a y vector which has
+  # already been passed through G.shared to enable easy class-wise
+  # interpolation later. If we passed in the one-hot and then ran it through
+  # G.shared in this forward function, it would be harder to handle.
+  def forward(self, z, y):
+    # If hierarchical, concatenate zs and ys
+    y = self.transform_labels(y)
+    if self.hier:
+      zs = torch.split(z, self.z_chunk_size, 1)
+      z = zs[0]
+      ys = [torch.cat([y, item], 1) for item in zs[1:]]
+    else:
+      ys = [y] * len(self.blocks)
+      
+    # First linear layer
+    h = self.linear(z)
+    # Reshape
+    h = h.view(h.size(0), -1, self.bottom_width, self.bottom_width)
+    
+    # Loop over blocks
+    for index, blocklist in enumerate(self.blocks):
+      # Second inner loop in case block has multiple layers
+      for block in blocklist:
+        h = block(h, ys[index])
+        
+    # Apply batchnorm-relu-conv-tanh at output
+    return torch.sigmoid(self.output_layer(h))
+#    return torch.tanh(self.output_layer(h))
+
+
+# Discriminator architecture, same paradigm as G's above
+def D_arch(ch=64, attention='64',ksize='333333', dilation='111111'):
+  arch = {}
+  arch[256] = {'in_channels' :  [3] + [ch*item for item in [1, 2, 4, 8, 8, 16]],
+               'out_channels' : [item * ch for item in [1, 2, 4, 8, 8, 16, 16]],
+               'downsample' : [True] * 6 + [False],
+               'resolution' : [128, 64, 32, 16, 8, 4, 4 ],
+               'attention' : {2**i: 2**i in [int(item) for item in attention.split('_')]
+                              for i in range(2,8)}}
+  arch[128] = {'in_channels' :  [3] + [ch*item for item in [1, 2, 4, 8, 16]],
+               'out_channels' : [item * ch for item in [1, 2, 4, 8, 16, 16]],
+               'downsample' : [True] * 5 + [False],
+               'resolution' : [64, 32, 16, 8, 4, 4],
+               'attention' : {2**i: 2**i in [int(item) for item in attention.split('_')]
+                              for i in range(2,8)}}
+  arch[64]  = {'in_channels' :  [3] + [ch*item for item in [1, 2, 4, 8]],
+               'out_channels' : [item * ch for item in [1, 2, 4, 8, 16]],
+               'downsample' : [True] * 4 + [False],
+               'resolution' : [32, 16, 8, 4, 4],
+               'attention' : {2**i: 2**i in [int(item) for item in attention.split('_')]
+                              for i in range(2,7)}}
+  arch[32]  = {'in_channels' :  [3] + [item * ch for item in [4, 4, 4]],
+               'out_channels' : [item * ch for item in [4, 4, 4, 4]],
+               'downsample' : [True, True, False, False],
+               'resolution' : [16, 16, 16, 16],
+               'attention' : {2**i: 2**i in [int(item) for item in attention.split('_')]
+                              for i in range(2,6)}}
+  return arch
+
+class Discriminator(NeuralNetwork):
+
+  def __init__(self, D_ch=64, D_wide=True, resolution=64, labels_dim=labels_dim,
+               D_kernel_size=3, D_attn='64', n_classes=1,
+               num_D_SVs=1, num_D_SV_itrs=1, D_activation=nn.ReLU(inplace=False),
+               D_lr=2e-4, D_B1=0.0, D_B2=0.999, adam_eps=1e-8,
+               SN_eps=1e-12, output_dim=1, D_mixed_precision=False, D_fp16=False,
+               D_init='ortho', skip_init=False, D_param='SN', **kwargs):
+    super(Discriminator, self).__init__()
+    # Width multiplier
+    self.ch = D_ch
+    # Use Wide D as in BigGAN and SA-GAN or skinny D as in SN-GAN?
+    self.D_wide = D_wide
+    # Resolution
+    self.resolution = resolution
+    # Kernel size
+    self.kernel_size = D_kernel_size
+    # Attention?
+    self.attention = D_attn
+    # Number of classes
+    self.n_classes = n_classes
+    # Activation
+    self.activation = D_activation
+    # Initialization style
+    self.init = D_init
+    # Parameterization style
+    self.D_param = D_param
+    # Epsilon for Spectral Norm?
+    self.SN_eps = SN_eps
+    # Fp16?
+    self.fp16 = D_fp16
+    # Architecture
+    self.arch = D_arch(self.ch, self.attention)[resolution]
+
+    # Which convs, batchnorms, and linear layers to use
+    # No option to turn off SN in D right now
+    if self.D_param == 'SN':
+      self.which_conv = functools.partial(layers.SNConv2d,
+                          kernel_size=3, padding=1,
+                          num_svs=num_D_SVs, num_itrs=num_D_SV_itrs,
+                          eps=self.SN_eps)
+      self.which_linear = functools.partial(layers.SNLinear,
+                          num_svs=num_D_SVs, num_itrs=num_D_SV_itrs,
+                          eps=self.SN_eps)
+      self.which_embedding = functools.partial(layers.SNEmbedding,
+                              num_svs=num_D_SVs, num_itrs=num_D_SV_itrs,
+                              eps=self.SN_eps)
+    # Prepare model
+    # prepare label input
+    self.transform_label_layer = torch.nn.Linear(labels_dim, 1024)
+    # self.blocks is a doubly-nested list of modules, the outer loop intended
+    # to be over blocks at a given resolution (resblocks and/or self-attention)
+    self.blocks = []
+    for index in range(len(self.arch['out_channels'])):
+      self.blocks += [[layers.DBlock(in_channels=self.arch['in_channels'][index],
+                       out_channels=self.arch['out_channels'][index],
+                       which_conv=self.which_conv,
+                       wide=self.D_wide,
+                       activation=self.activation,
+                       preactivation=(index > 0),
+                       downsample=(nn.AvgPool2d(2) if self.arch['downsample'][index] else None))]]
+      # If attention on this block, attach it to the end
+      if self.arch['attention'][self.arch['resolution'][index]]:
+        print('Adding attention layer in D at resolution %d' % self.arch['resolution'][index])
+        self.blocks[-1] += [layers.Attention(self.arch['out_channels'][index],
+                                             self.which_conv)]
+    # Turn self.blocks into a ModuleList so that it's all properly registered.
+    self.blocks = nn.ModuleList([nn.ModuleList(block) for block in self.blocks])
+    # Linear output layer. The output dimension is typically 1, but may be
+    # larger if we're e.g. turning this into a VAE with an inference output
+    self.linear = self.which_linear(self.arch['out_channels'][-1], output_dim)
+    # Embedding for projection discrimination
+    self.embed = self.which_embedding(self.n_classes, self.arch['out_channels'][-1])
+
+    # Initialize weights
+    if not skip_init:
+      self.init_weights()
+
+    # Set up optimizer
+    self.lr, self.B1, self.B2, self.adam_eps = D_lr, D_B1, D_B2, adam_eps
+    if D_mixed_precision:
+      print('Using fp16 adam in D...')
+      import utils
+      self.optim = utils.Adam16(params=self.parameters(), lr=self.lr,
+                             betas=(self.B1, self.B2), weight_decay=0, eps=self.adam_eps)
+    else:
+      self.optim = optim.Adam(params=self.parameters(), lr=self.lr,
+                             betas=(self.B1, self.B2), weight_decay=0, eps=self.adam_eps)
+    # LR scheduling, left here for forward compatibility
+    # self.lr_sched = {'itr' : 0}# if self.progressive else {}
+    # self.j = 0
+    self.set_optimizer(parameter.optimizer, lr=parameter.learning_rate*3, betas=parameter.betas)
+
+  # Initialize
+  def init_weights(self):
+    self.param_count = 0
+    for module in self.modules():
+      if (isinstance(module, nn.Conv2d)
+          or isinstance(module, nn.Linear)
+          or isinstance(module, nn.Embedding)):
+        if self.init == 'ortho':
+          init.orthogonal_(module.weight)
+        elif self.init == 'N02':
+          init.normal_(module.weight, 0, 0.02)
+        elif self.init in ['glorot', 'xavier']:
+          init.xavier_uniform_(module.weight)
+        else:
+          print('Init style not recognized...')
+        self.param_count += sum([p.data.nelement() for p in module.parameters()])
+    print('Param count for D''s initialized parameters: %d' % self.param_count)
+
+  def transform_labels(self, labels):
+    """ prepore labels for input to discriminator """
+    return self.transform_label_layer(labels)
+
+
+  def forward(self, x, y=None):
+    # Stick x into h for cleaner for loops without flow control
+    h = x
+    # Loop over blocks
+    for index, blocklist in enumerate(self.blocks):
+      for block in blocklist:
+        h = block(h)
+    # Apply global sum pooling as in SN-GAN
+    h = torch.sum(self.activation(h), [2, 3])
+    # Get initial class-unconditional output
+    out = self.linear(h)
+    # Get projection of final featureset onto class vectors and add to evidence
+    y = self.transform_labels(y)
+    out = out + torch.sum(y * h, 1, keepdim=True)
+#    out = out + torch.sum(self.embed(y) * h, 1, keepdim=True)  ## use y = torch.tensor(0)
+    return out
+
+# Parallelized G_D to minimize cross-gpu communication
+# Without this, Generator outputs would get all-gathered and then rebroadcast.
+class G_D(nn.Module):
+  def __init__(self, G, D):
+    super(G_D, self).__init__()
+    self.G = G
+    self.D = D
+
+  def forward(self, z, gy, x=None, dy=None, train_G=False, return_G_z=False,
+              split_D=False):              
+    # If training G, enable grad tape
+    with torch.set_grad_enabled(train_G):
+      # Get Generator output given noise
+      G_z = self.G(z, self.G.shared(gy))
+      # Cast as necessary
+      if self.G.fp16 and not self.D.fp16:
+        G_z = G_z.float()
+      if self.D.fp16 and not self.G.fp16:
+        G_z = G_z.half()
+    # Split_D means to run D once with real data and once with fake,
+    # rather than concatenating along the batch dimension.
+    if split_D:
+      D_fake = self.D(G_z, gy)
+      if x is not None:
+        D_real = self.D(x, dy)
+        return D_fake, D_real
+      else:
+        if return_G_z:
+          return D_fake, G_z
+        else:
+          return D_fake
+    # If real data is provided, concatenate it with the Generator's output
+    # along the batch dimension for improved efficiency.
+    else:
+      D_input = torch.cat([G_z, x], 0) if x is not None else G_z
+      D_class = torch.cat([gy, dy], 0) if dy is not None else gy
+      # Get Discriminator output
+      D_out = self.D(D_input, D_class)
+      if x is not None:
+        return torch.split(D_out, [G_z.shape[0], x.shape[0]]) # D_fake, D_real
+      else:
+        if return_G_z:
+          return D_out, G_z
+        else:
+          return D_out

src/models/big/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2019 Andy Brock
+
+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.

src/models/big/README.md

@@ -0,0 +1,144 @@
+# BigGAN-PyTorch
+The author's officially unofficial PyTorch BigGAN implementation.
+
+![Dogball? Dogball!](imgs/header_image.jpg?raw=true "Dogball? Dogball!")
+
+
+This repo contains code for 4-8 GPU training of BigGANs from [Large Scale GAN Training for High Fidelity Natural Image Synthesis](https://arxiv.org/abs/1809.11096) by Andrew Brock, Jeff Donahue, and Karen Simonyan.
+
+This code is by Andy Brock and Alex Andonian.
+
+## How To Use This Code
+You will need:
+
+- [PyTorch](https://PyTorch.org/), version 1.0.1
+- tqdm, numpy, scipy, and h5py
+- The ImageNet training set
+
+First, you may optionally prepare a pre-processed HDF5 version of your target dataset for faster I/O. Following this (or not), you'll need the Inception moments needed to calculate FID. These can both be done by modifying and running
+
+```sh
+sh scripts/utils/prepare_data.sh
+```
+
+Which by default assumes your ImageNet training set is downloaded into the root folder `data` in this directory, and will prepare the cached HDF5 at 128x128 pixel resolution.
+
+In the scripts folder, there are multiple bash scripts which will train BigGANs with different batch sizes. This code assumes you do not have access to a full TPU pod, and accordingly
+spoofs mega-batches by using gradient accumulation (averaging grads over multiple minibatches, and only taking an optimizer step after N accumulations). By default, the `launch_BigGAN_bs256x8.sh` script trains a
+full-sized BigGAN model with a batch size of 256 and 8 gradient accumulations, for a total batch size of 2048. On 8xV100 with full-precision training (no Tensor cores), this script takes 15 days to train to 150k iterations.
+
+You will first need to figure out the maximum batch size your setup can support. The pre-trained models provided here were trained on 8xV100 (16GB VRAM each) which can support slightly more than the BS256 used by default.
+Once you've determined this, you should modify the script so that the batch size times the number of gradient accumulations is equal to your desired total batch size (BigGAN defaults to 2048).
+
+Note also that this script uses the `--load_in_mem` arg, which loads the entire (~64GB) I128.hdf5 file into RAM for faster data loading. If you don't have enough RAM to support this (probably 96GB+), remove this argument.
+
+
+## Metrics and Sampling
+![I believe I can fly!](imgs/interp_sample.jpg?raw=true "I believe I can fly!")
+
+During training, this script will output logs with training metrics and test metrics, will save multiple copies (2 most recent and 5 highest-scoring) of the model weights/optimizer params, and will produce samples and interpolations every time it saves weights.
+The logs folder contains scripts to process these logs and plot the results using MATLAB (sorry not sorry).
+
+After training, one can use `sample.py` to produce additional samples and interpolations, test with different truncation values, batch sizes, number of standing stat accumulations, etc. See the `sample_BigGAN_bs256x8.sh` script for an example.
+
+By default, everything is saved to weights/samples/logs/data folders which are assumed to be in the same folder as this repo.
+You can point all of these to a different base folder using the `--base_root` argument, or pick specific locations for each of these with their respective arguments (e.g. `--logs_root`).
+
+We include scripts to run BigGAN-deep, but we have not fully trained a model using them, so consider them untested. Additionally, we include scripts to run a model on CIFAR, and to run SA-GAN (with EMA) and SN-GAN on ImageNet. The SA-GAN code assumes you have 4xTitanX (or equivalent in terms of GPU RAM) and will run with a batch size of 128 and 2 gradient accumulations.
+
+## An Important Note on Inception Metrics
+This repo uses the PyTorch in-built inception network to calculate IS and FID. 
+These scores are different from the scores you would get using the official TF inception code, and are only for monitoring purposes!
+Run sample.py on your model, with the `--sample_npz` argument, then run inception_tf13 to calculate the actual TensorFlow IS. Note that you will need to have TensorFlow 1.3 or earlier installed, as TF1.4+ breaks the original IS code.
+
+## Pretrained models
+![PyTorch Inception Score and FID](imgs/IS_FID.png)
+We include two pretrained model checkpoints (with G, D, the EMA copy of G, the optimizers, and the state dict):
+- The main checkpoint is for a BigGAN trained on ImageNet at 128x128, using BS256 and 8 gradient accumulations, taken just before collapse, with a TF Inception Score of 97.35 +/- 1.79: [LINK](https://drive.google.com/open?id=1nAle7FCVFZdix2--ks0r5JBkFnKw8ctW)
+- An earlier checkpoint of the first model (100k G iters), at high performance but well before collapse, which may be easier to fine-tune: [LINK](https://drive.google.com/open?id=1dmZrcVJUAWkPBGza_XgswSuT-UODXZcO)
+
+
+
+Pretrained models for Places-365 coming soon.
+
+This repo also contains scripts for porting the original TFHub BigGAN Generator weights to PyTorch. See the scripts in the TFHub folder for more details.
+
+## Fine-tuning, Using Your Own Dataset, or Making New Training Functions
+![That's deep, man](imgs/DeepSamples.png?raw=true "Deep Samples")
+
+If you wish to resume interrupted training or fine-tune a pre-trained model, run the same launch script but with the `--resume` argument added. 
+Experiment names are automatically generated from the configuration, but can be overridden using the `--experiment_name` arg (for example, if you wish to fine-tune a model using modified optimizer settings).
+
+To prep your own dataset, you will need to add it to datasets.py and modify the convenience dicts in utils.py (dset_dict, imsize_dict, root_dict, nclass_dict, classes_per_sheet_dict) to have the appropriate metadata for your dataset.
+Repeat the process in prepare_data.sh (optionally produce an HDF5 preprocessed copy, and calculate the Inception Moments for FID).
+
+By default, the training script will save the top 5 best checkpoints as measured by Inception Score. 
+For datasets other than ImageNet, Inception Score can be a very poor measure of quality, so you will likely want to use `--which_best FID` instead.
+
+To use your own training function (e.g. train a BigVAE): either modify train_fns.GAN_training_function or add a new train fn and add it after the `if config['which_train_fn'] == 'GAN':` line in `train.py`.
+
+
+## Neat Stuff
+- We include the full training and metrics logs [here](https://drive.google.com/open?id=1ZhY9Mg2b_S4QwxNmt57aXJ9FOC3ZN1qb) for reference. I've found that one of the hardest things about re-implementing a paper can be checking if the logs line up early in training,
+especially if training takes multiple weeks. Hopefully these will be helpful for future work.
+- We include an accelerated FID calculation--the original scipy version can require upwards of 10 minutes to calculate the matrix sqrt, this version uses an accelerated PyTorch version to calculate it in under a second.
+- We include an accelerated, low-memory consumption ortho reg implementation.
+- By default, we only compute the top singular value (the spectral norm), but this code supports computing more SVs through the `--num_G_SVs` argument.
+
+## Key Differences Between This Code And The Original BigGAN
+- We use the optimizer settings from SA-GAN (G_lr=1e-4, D_lr=4e-4, num_D_steps=1, as opposed to BigGAN's G_lr=5e-5, D_lr=2e-4, num_D_steps=2).
+While slightly less performant, this was the first corner we cut to bring training times down.
+- By default, we do not use Cross-Replica BatchNorm (AKA Synced BatchNorm). 
+The two variants we tried (a custom, naive one and the one included in this repo) have slightly different gradients (albeit identical forward passes) from the built-in BatchNorm, which appear to be sufficient to cripple training.
+- Gradient accumulation means that we update the SV estimates and the BN statistics 8 times more frequently. This means that the BN stats are much closer to standing stats, and that the singular value estimates tend to be more accurate.
+Because of this, we measure metrics by default with G in test mode (using the BatchNorm running stat estimates instead of computing standing stats as in the paper). We do still support standing stats (see the sample.sh scripts).
+This could also conceivably result in gradients from the earlier accumulations being stale, but in practice this does not appear to be a problem.
+- The currently provided pretrained models were not trained with orthogonal regularization. Training without ortho reg seems to increase the probability that models will not be amenable to truncation,
+but it looks like this particular model got a winning ticket. Regardless, we provide two highly optimized (fast and minimal memory consumption) ortho reg implementations which directly compute the ortho reg. gradients.
+
+## A Note On The Design Of This Repo
+This code is designed from the ground up to serve as an extensible, hackable base for further research code. 
+We've put a lot of thought into making sure the abstractions are the *right* thickness for research--not so thick as to be impenetrable, but not so thin as to be useless.
+The key idea is that if you want to experiment with a SOTA setup and make some modification (try out your own new loss function, architecture, self-attention block, etc) you should be able to easily do so just by dropping your code in one or two places, without having to worry about the rest of the codebase.
+Things like the use of self.which_conv and functools.partial in the BigGAN.py model definition were put together with this in mind, as was the design of the Spectral Norm class inheritance.
+
+With that said, this is a somewhat large codebase for a single project. While we tried to be thorough with the comments, if there's something you think could be more clear, better written, or better refactored, please feel free to raise an issue or a pull request.
+
+## Feature Requests
+Want to work on or improve this code? There are a couple things this repo would benefit from, but which don't yet work.
+
+- Synchronized BatchNorm (AKA Cross-Replica BatchNorm). We tried out two variants of this, but for some unknown reason it crippled training each time.
+  We have not tried the [apex](https://github.com/NVIDIA/apex) SyncBN as my school's servers are on ancient NVIDIA drivers that don't support it--apex would probably be a good place to start. 
+- Mixed precision training and making use of Tensor cores. This repo includes a naive mixed-precision Adam implementation which works early in training but leads to early collapse, and doesn't do anything to activate Tensor cores (it just reduces memory consumption).
+  As above, integrating [apex](https://github.com/NVIDIA/apex) into this code and employing its mixed-precision training techniques to take advantage of Tensor cores and reduce memory consumption could yield substantial speed gains.
+
+## Misc Notes
+See [This directory](https://gist.github.com/yrevar/942d3a0ac09ec9e5eb3a) for ImageNet labels.
+
+If you use this code, please cite
+```text
+@inproceedings{
+brock2018large,
+title={Large Scale {GAN} Training for High Fidelity Natural Image Synthesis},
+author={Andrew Brock and Jeff Donahue and Karen Simonyan},
+booktitle={International Conference on Learning Representations},
+year={2019},
+url={https://openreview.net/forum?id=B1xsqj09Fm},
+}
+```
+
+## Acknowledgments
+Thanks to Google for the generous cloud credit donations.
+
+[SyncBN](https://github.com/vacancy/Synchronized-BatchNorm-PyTorch) by Jiayuan Mao and Tete Xiao.
+
+[Progress bar](https://github.com/Lasagne/Recipes/tree/master/papers/densenet) originally from Jan Schlüter.
+
+Test metrics logger from [VoxNet.](https://github.com/dimatura/voxnet)
+
+PyTorch [implementation of cov](https://discuss.PyTorch.org/t/covariance-and-gradient-support/16217/2) from Modar M. Alfadly.
+
+PyTorch [fast Matrix Sqrt](https://github.com/msubhransu/matrix-sqrt) for FID from Tsung-Yu Lin and Subhransu Maji.
+
+TensorFlow Inception Score code from [OpenAI's Improved-GAN.](https://github.com/openai/improved-gan)
+

src/models/big/animal_hash.py

@@ -0,0 +1,439 @@
+c = ['Aardvark', 'Abyssinian', 'Affenpinscher', 'Akbash', 'Akita', 'Albatross',
+     'Alligator', 'Alpaca', 'Angelfish', 'Ant', 'Anteater', 'Antelope', 'Ape',
+     'Armadillo', 'Ass', 'Avocet', 'Axolotl', 'Baboon', 'Badger', 'Balinese',
+     'Bandicoot', 'Barb', 'Barnacle', 'Barracuda', 'Bat', 'Beagle', 'Bear',
+     'Beaver', 'Bee', 'Beetle', 'Binturong', 'Bird', 'Birman', 'Bison',
+     'Bloodhound', 'Boar', 'Bobcat', 'Bombay', 'Bongo', 'Bonobo', 'Booby',
+     'Budgerigar', 'Buffalo', 'Bulldog', 'Bullfrog', 'Burmese', 'Butterfly',
+     'Caiman', 'Camel', 'Capybara', 'Caracal', 'Caribou', 'Cassowary', 'Cat',
+     'Caterpillar', 'Catfish', 'Cattle', 'Centipede', 'Chameleon', 'Chamois',
+     'Cheetah', 'Chicken', 'Chihuahua', 'Chimpanzee', 'Chinchilla', 'Chinook', 
+     'Chipmunk', 'Chough', 'Cichlid', 'Clam', 'Coati', 'Cobra', 'Cockroach',
+     'Cod', 'Collie', 'Coral', 'Cormorant', 'Cougar', 'Cow', 'Coyote', 
+     'Crab', 'Crane', 'Crocodile', 'Crow', 'Curlew', 'Cuscus', 'Cuttlefish',
+     'Dachshund', 'Dalmatian', 'Deer', 'Dhole', 'Dingo', 'Dinosaur', 'Discus',
+     'Dodo', 'Dog', 'Dogball', 'Dogfish', 'Dolphin', 'Donkey', 'Dormouse',
+     'Dove', 'Dragonfly', 'Drever', 'Duck', 'Dugong', 'Dunker', 'Dunlin', 
+     'Eagle', 'Earwig', 'Echidna', 'Eel', 'Eland', 'Elephant', 'ElephantSeal',
+     'Elk', 'Emu', 'Falcon', 'Ferret', 'Finch', 'Fish', 'Flamingo', 'Flounder',
+     'Fly', 'Fossa', 'Fox', 'Frigatebird', 'Frog', 'Galago', 'Gar', 'Gaur', 
+     'Gazelle', 'Gecko', 'Gerbil', 'Gharial', 'GiantPanda', 'Gibbon', 'Giraffe',
+     'Gnat', 'Gnu', 'Goat', 'Goldfinch', 'Goldfish', 'Goose', 'Gopher',
+     'Gorilla', 'Goshawk', 'Grasshopper', 'Greyhound', 'Grouse', 'Guanaco', 
+     'GuineaFowl', 'GuineaPig', 'Gull', 'Guppy', 'Hamster', 'Hare', 'Harrier', 
+     'Havanese', 'Hawk', 'Hedgehog', 'Heron', 'Herring', 'Himalayan', 
+     'Hippopotamus', 'Hornet', 'Horse', 'Human', 'Hummingbird', 'Hyena', 
+     'Ibis', 'Iguana', 'Impala', 'Indri', 'Insect', 'Jackal', 'Jaguar', 
+     'Javanese', 'Jay', 'Jellyfish', 'Kakapo', 'Kangaroo', 'Kingfisher', 
+     'Kiwi', 'Koala', 'KomodoDragon', 'Kouprey', 'Kudu', 'Labradoodle', 
+     'Ladybird', 'Lapwing', 'Lark', 'Lemming', 'Lemur', 'Leopard', 'Liger',
+     'Lion', 'Lionfish', 'Lizard', 'Llama', 'Lobster', 'Locust', 'Loris', 
+     'Louse', 'Lynx', 'Lyrebird', 'Macaw', 'Magpie', 'Mallard', 'Maltese',
+     'Manatee', 'Mandrill', 'Markhor', 'Marten', 'Mastiff', 'Mayfly', 'Meerkat',
+     'Millipede', 'Mink', 'Mole', 'Molly', 'Mongoose', 'Mongrel', 'Monkey',
+     'Moorhen', 'Moose', 'Mosquito', 'Moth', 'Mouse', 'Mule', 'Narwhal',
+     'Neanderthal', 'Newfoundland', 'Newt', 'Nightingale', 'Numbat', 'Ocelot',
+     'Octopus', 'Okapi', 'Olm', 'Opossum', 'Orang-utan', 'Oryx', 'Ostrich', 
+     'Otter', 'Owl', 'Ox', 'Oyster', 'Pademelon', 'Panther', 'Parrot',
+     'Partridge', 'Peacock', 'Peafowl', 'Pekingese', 'Pelican', 'Penguin', 
+     'Persian', 'Pheasant', 'Pig', 'Pigeon', 'Pika', 'Pike', 'Piranha', 
+     'Platypus', 'Pointer', 'Pony', 'Poodle', 'Porcupine', 'Porpoise',
+     'Possum', 'PrairieDog', 'Prawn', 'Puffin', 'Pug', 'Puma', 'Quail', 
+     'Quelea', 'Quetzal', 'Quokka', 'Quoll', 'Rabbit', 'Raccoon', 'Ragdoll', 
+     'Rail', 'Ram', 'Rat', 'Rattlesnake', 'Raven', 'RedDeer', 'RedPanda',
+     'Reindeer', 'Rhinoceros', 'Robin', 'Rook', 'Rottweiler', 'Ruff',
+     'Salamander', 'Salmon', 'SandDollar', 'Sandpiper', 'Saola', 
+     'Sardine', 'Scorpion', 'SeaLion', 'SeaUrchin', 'Seahorse',
+     'Seal', 'Serval', 'Shark', 'Sheep', 'Shrew', 'Shrimp', 'Siamese',
+     'Siberian', 'Skunk', 'Sloth', 'Snail', 'Snake', 'Snowshoe', 'Somali', 
+     'Sparrow', 'Spider', 'Sponge', 'Squid', 'Squirrel', 'Starfish', 'Starling',
+     'Stingray', 'Stinkbug', 'Stoat', 'Stork', 'Swallow', 'Swan', 'Tang', 
+     'Tapir', 'Tarsier', 'Termite', 'Tetra', 'Tiffany', 'Tiger', 'Toad', 
+     'Tortoise', 'Toucan', 'Tropicbird', 'Trout', 'Tuatara', 'Turkey', 
+     'Turtle', 'Uakari', 'Uguisu', 'Umbrellabird', 'Viper', 'Vulture',
+     'Wallaby', 'Walrus', 'Warthog', 'Wasp', 'WaterBuffalo', 'Weasel',
+     'Whale', 'Whippet', 'Wildebeest', 'Wolf', 'Wolverine', 'Wombat', 
+     'Woodcock', 'Woodlouse', 'Woodpecker', 'Worm', 'Wrasse', 'Wren', 
+     'Yak', 'Zebra', 'Zebu', 'Zonkey']
+a = ['able', 'above', 'absent', 'absolute', 'abstract', 'abundant', 'academic',
+     'acceptable', 'accepted', 'accessible', 'accurate', 'accused', 'active', 
+     'actual', 'acute', 'added', 'additional', 'adequate', 'adjacent', 
+     'administrative', 'adorable', 'advanced', 'adverse', 'advisory', 
+     'aesthetic', 'afraid', 'african', 'aggregate', 'aggressive', 'agreeable', 
+     'agreed', 'agricultural', 'alert', 'alive', 'alleged', 'allied', 'alone', 
+     'alright', 'alternative', 'amateur', 'amazing', 'ambitious', 'american', 
+     'amused', 'ancient', 'angry', 'annoyed', 'annual', 'anonymous', 'anxious', 
+     'appalling', 'apparent', 'applicable', 'appropriate', 'arab', 'arbitrary',
+     'architectural', 'armed', 'arrogant', 'artificial', 'artistic', 'ashamed', 
+     'asian', 'asleep', 'assistant', 'associated', 'atomic', 'attractive',
+     'australian', 'automatic', 'autonomous', 'available', 'average',
+     'awake', 'aware', 'awful', 'awkward', 'back', 'bad', 'balanced', 'bare', 
+     'basic', 'beautiful', 'beneficial', 'better', 'bewildered', 'big', 
+     'binding', 'biological', 'bitter', 'bizarre', 'black', 'blank', 'blind', 
+     'blonde', 'bloody', 'blue', 'blushing', 'boiling', 'bold', 'bored', 
+     'boring', 'bottom', 'brainy', 'brave', 'breakable', 'breezy', 'brief', 
+     'bright', 'brilliant', 'british', 'broad', 'broken', 'brown', 'bumpy', 
+     'burning', 'busy', 'calm', 'canadian', 'capable', 'capitalist', 'careful',
+     'casual', 'catholic', 'causal', 'cautious', 'central', 'certain', 
+     'changing', 'characteristic', 'charming', 'cheap', 'cheerful', 'chemical', 
+     'chief', 'chilly', 'chinese', 'chosen', 'christian', 'chronic', 'chubby', 
+     'circular', 'civic', 'civil', 'civilian', 'classic', 'classical', 'clean',
+     'clear', 'clever', 'clinical', 'close', 'closed', 'cloudy', 'clumsy', 
+     'coastal', 'cognitive', 'coherent', 'cold', 'collective', 'colonial', 
+     'colorful', 'colossal', 'coloured', 'colourful', 'combative', 'combined',
+     'comfortable', 'coming', 'commercial', 'common', 'communist', 'compact', 
+     'comparable', 'comparative', 'compatible', 'competent', 'competitive', 
+     'complete', 'complex', 'complicated', 'comprehensive', 'compulsory',
+     'conceptual', 'concerned', 'concrete', 'condemned', 'confident', 
+     'confidential', 'confused', 'conscious', 'conservation', 'conservative',
+     'considerable', 'consistent', 'constant', 'constitutional', 
+     'contemporary', 'content', 'continental', 'continued', 'continuing', 
+     'continuous', 'controlled', 'controversial', 'convenient', 'conventional',
+     'convinced', 'convincing', 'cooing', 'cool', 'cooperative', 'corporate',
+     'correct', 'corresponding', 'costly', 'courageous', 'crazy', 'creative', 
+     'creepy', 'criminal', 'critical', 'crooked', 'crowded', 'crucial', 
+     'crude', 'cruel', 'cuddly', 'cultural', 'curious', 'curly', 'current', 
+     'curved', 'cute', 'daily', 'damaged', 'damp', 'dangerous', 'dark', 'dead',
+     'deaf', 'deafening', 'dear', 'decent', 'decisive', 'deep', 'defeated', 
+     'defensive', 'defiant', 'definite', 'deliberate', 'delicate', 'delicious',
+     'delighted', 'delightful', 'democratic', 'dependent', 'depressed', 
+     'desirable', 'desperate', 'detailed', 'determined', 'developed', 
+     'developing', 'devoted', 'different', 'difficult', 'digital', 'diplomatic', 
+     'direct', 'dirty', 'disabled', 'disappointed', 'disastrous', 
+     'disciplinary', 'disgusted', 'distant', 'distinct', 'distinctive',
+     'distinguished', 'disturbed', 'disturbing', 'diverse', 'divine', 'dizzy', 
+     'domestic', 'dominant', 'double', 'doubtful', 'drab', 'dramatic',
+     'dreadful', 'driving', 'drunk', 'dry', 'dual', 'due', 'dull', 'dusty',
+     'dutch', 'dying', 'dynamic', 'eager', 'early', 'eastern', 'easy', 
+     'economic', 'educational', 'eerie', 'effective', 'efficient', 
+     'elaborate', 'elated', 'elderly', 'eldest', 'electoral', 'electric',
+     'electrical', 'electronic', 'elegant', 'eligible', 'embarrassed',
+     'embarrassing', 'emotional', 'empirical', 'empty', 'enchanting',
+     'encouraging', 'endless', 'energetic', 'english', 'enormous', 
+     'enthusiastic', 'entire', 'entitled', 'envious', 'environmental', 'equal', 
+     'equivalent', 'essential', 'established', 'estimated', 'ethical', 
+     'ethnic', 'european', 'eventual', 'everyday', 'evident', 'evil', 
+     'evolutionary', 'exact', 'excellent', 'exceptional', 'excess', 
+     'excessive', 'excited', 'exciting', 'exclusive', 'existing', 'exotic', 
+     'expected', 'expensive', 'experienced', 'experimental', 'explicit',
+     'extended', 'extensive', 'external', 'extra', 'extraordinary', 'extreme', 
+     'exuberant', 'faint', 'fair', 'faithful', 'familiar', 'famous', 'fancy',
+     'fantastic', 'far', 'fascinating', 'fashionable', 'fast', 'fat', 'fatal', 
+     'favourable', 'favourite', 'federal', 'fellow', 'female', 'feminist', 
+     'few', 'fierce', 'filthy', 'final', 'financial', 'fine', 'firm', 'fiscal', 
+     'fit', 'fixed', 'flaky', 'flat', 'flexible', 'fluffy', 'fluttering', 
+     'flying', 'following', 'fond', 'foolish', 'foreign', 'formal', 
+     'formidable', 'forthcoming', 'fortunate', 'forward', 'fragile', 
+     'frail', 'frantic', 'free', 'french', 'frequent', 'fresh', 'friendly', 
+     'frightened', 'front', 'frozen', 'fucking', 'full', 'full-time', 'fun',
+     'functional', 'fundamental', 'funny', 'furious', 'future', 'fuzzy',
+     'gastric', 'gay', 'general', 'generous', 'genetic', 'gentle', 'genuine',
+     'geographical', 'german', 'giant', 'gigantic', 'given', 'glad',
+     'glamorous', 'gleaming', 'global', 'glorious', 'golden', 'good', 
+     'gorgeous', 'gothic', 'governing', 'graceful', 'gradual', 'grand', 
+     'grateful', 'greasy', 'great', 'greek', 'green', 'grey', 'grieving',
+     'grim', 'gross', 'grotesque', 'growing', 'grubby', 'grumpy', 'guilty',
+     'handicapped', 'handsome', 'happy', 'hard', 'harsh', 'head', 'healthy', 
+     'heavy', 'helpful', 'helpless', 'hidden', 'high', 'high-pitched',
+     'hilarious', 'hissing', 'historic', 'historical', 'hollow', 'holy',
+     'homeless', 'homely', 'hon', 'honest', 'horizontal', 'horrible', 
+     'hostile', 'hot', 'huge', 'human', 'hungry', 'hurt', 'hushed', 'husky',
+     'icy', 'ideal', 'identical', 'ideological', 'ill', 'illegal', 
+     'imaginative', 'immediate', 'immense', 'imperial', 'implicit', 
+     'important', 'impossible', 'impressed', 'impressive', 'improved', 
+     'inadequate', 'inappropriate', 'inc', 'inclined', 'increased', 
+     'increasing', 'incredible', 'independent', 'indian', 'indirect', 
+     'individual', 'industrial', 'inevitable', 'influential', 'informal',
+     'inherent', 'initial', 'injured', 'inland', 'inner', 'innocent', 
+     'innovative', 'inquisitive', 'instant', 'institutional', 'insufficient',
+     'intact', 'integral', 'integrated', 'intellectual', 'intelligent', 
+     'intense', 'intensive', 'interested', 'interesting', 'interim', 
+     'interior', 'intermediate', 'internal', 'international', 'intimate',
+     'invisible', 'involved', 'iraqi', 'irish', 'irrelevant', 'islamic',
+     'isolated', 'israeli', 'italian', 'itchy', 'japanese', 'jealous', 
+     'jewish', 'jittery', 'joint', 'jolly', 'joyous', 'judicial', 'juicy', 
+     'junior', 'just', 'keen', 'key', 'kind', 'known', 'korean', 'labour', 
+     'large', 'large-scale', 'late', 'latin', 'lazy', 'leading', 'left', 
+     'legal', 'legislative', 'legitimate', 'lengthy', 'lesser', 'level', 
+     'lexical', 'liable', 'liberal', 'light', 'like', 'likely', 'limited', 
+     'linear', 'linguistic', 'liquid', 'literary', 'little', 'live', 'lively', 
+     'living', 'local', 'logical', 'lonely', 'long', 'long-term', 'loose', 
+     'lost', 'loud', 'lovely', 'low', 'loyal', 'ltd', 'lucky', 'mad',
+     'magenta', 'magic', 'magnetic', 'magnificent', 'main', 'major', 'male',
+     'mammoth', 'managerial', 'managing', 'manual', 'many', 'marginal', 
+     'marine', 'marked', 'married', 'marvellous', 'marxist', 'mass', 'massive', 
+     'mathematical', 'mature', 'maximum', 'mean', 'meaningful', 'mechanical',
+     'medical', 'medieval', 'melodic', 'melted', 'mental', 'mere', 
+     'metropolitan', 'mid', 'middle', 'middle-class', 'mighty', 'mild',
+     'military', 'miniature', 'minimal', 'minimum', 'ministerial', 'minor', 
+     'miserable', 'misleading', 'missing', 'misty', 'mixed', 'moaning', 
+     'mobile', 'moderate', 'modern', 'modest', 'molecular', 'monetary', 
+     'monthly', 'moral', 'motionless', 'muddy', 'multiple', 'mushy', 
+     'musical', 'mute', 'mutual', 'mysterious', 'naked', 'narrow', 'nasty',
+     'national', 'native', 'natural', 'naughty', 'naval', 'near', 'nearby', 
+     'neat', 'necessary', 'negative', 'neighbouring', 'nervous', 'net', 
+     'neutral', 'new', 'nice', 'nineteenth-century', 'noble', 'noisy', 
+     'normal', 'northern', 'nosy', 'notable', 'novel', 'nuclear', 'numerous',
+     'nursing', 'nutritious', 'nutty', 'obedient', 'objective', 'obliged', 
+     'obnoxious', 'obvious', 'occasional', 'occupational', 'odd', 'official',
+     'ok', 'okay', 'old', 'old-fashioned', 'olympic', 'only', 'open', 
+     'operational', 'opposite', 'optimistic', 'oral', 'orange', 'ordinary', 
+     'organic', 'organisational', 'original', 'orthodox', 'other', 'outdoor', 
+     'outer', 'outrageous', 'outside', 'outstanding', 'overall', 'overseas',
+     'overwhelming', 'painful', 'pale', 'palestinian', 'panicky', 'parallel', 
+     'parental', 'parliamentary', 'part-time', 'partial', 'particular', 
+     'passing', 'passive', 'past', 'patient', 'payable', 'peaceful', 
+     'peculiar', 'perfect', 'permanent', 'persistent', 'personal', 'petite',
+     'philosophical', 'physical', 'pink', 'plain', 'planned', 'plastic',
+     'pleasant', 'pleased', 'poised', 'polish', 'polite', 'political', 'poor', 
+     'popular', 'positive', 'possible', 'post-war', 'potential', 'powerful',
+     'practical', 'precious', 'precise', 'preferred', 'pregnant', 
+     'preliminary', 'premier', 'prepared', 'present', 'presidential', 
+     'pretty', 'previous', 'prickly', 'primary', 'prime', 'primitive', 
+     'principal', 'printed', 'prior', 'private', 'probable', 'productive',
+     'professional', 'profitable', 'profound', 'progressive', 'prominent', 
+     'promising', 'proper', 'proposed', 'prospective', 'protective', 
+     'protestant', 'proud', 'provincial', 'psychiatric', 'psychological',
+     'public', 'puny', 'pure', 'purple', 'purring', 'puzzled', 'quaint', 
+     'qualified', 'quick', 'quickest', 'quiet', 'racial', 'radical', 'rainy',
+     'random', 'rapid', 'rare', 'raspy', 'rational', 'ratty', 'raw', 'ready', 
+     'real', 'realistic', 'rear', 'reasonable', 'recent', 'red', 'reduced',
+     'redundant', 'regional', 'registered', 'regular', 'regulatory', 'related', 
+     'relative', 'relaxed', 'relevant', 'reliable', 'relieved', 'religious',
+     'reluctant', 'remaining', 'remarkable', 'remote', 'renewed',
+     'representative', 'repulsive', 'required', 'resident', 'residential',
+     'resonant', 'respectable', 'respective', 'responsible', 'resulting',
+     'retail', 'retired', 'revolutionary', 'rich', 'ridiculous', 'right',
+     'rigid', 'ripe', 'rising', 'rival', 'roasted', 'robust', 'rolling', 
+     'roman', 'romantic', 'rotten', 'rough', 'round', 'royal', 'rubber',
+     'rude', 'ruling', 'running', 'rural', 'russian', 'sacred', 'sad', 'safe',
+     'salty', 'satisfactory', 'satisfied', 'scared', 'scary', 'scattered',
+     'scientific', 'scornful', 'scottish', 'scrawny', 'screeching', 
+     'secondary', 'secret', 'secure', 'select', 'selected', 'selective', 
+     'selfish', 'semantic', 'senior', 'sensible', 'sensitive', 'separate',
+     'serious', 'severe', 'sexual', 'shaggy', 'shaky', 'shallow', 'shared', 
+     'sharp', 'sheer', 'shiny', 'shivering', 'shocked', 'short', 'short-term', 
+     'shrill', 'shy', 'sick', 'significant', 'silent', 'silky', 'silly', 
+     'similar', 'simple', 'single', 'skilled', 'skinny', 'sleepy', 'slight',
+     'slim', 'slimy', 'slippery', 'slow', 'small', 'smart', 'smiling', 
+     'smoggy', 'smooth', 'so-called', 'social', 'socialist', 'soft', 'solar',
+     'sole', 'solid', 'sophisticated', 'sore', 'sorry', 'sound', 'sour', 
+     'southern', 'soviet', 'spanish', 'spare', 'sparkling', 'spatial', 
+     'special', 'specific', 'specified', 'spectacular', 'spicy', 'spiritual',
+     'splendid', 'spontaneous', 'sporting', 'spotless', 'spotty', 'square', 
+     'squealing', 'stable', 'stale', 'standard', 'static', 'statistical', 
+     'statutory', 'steady', 'steep', 'sticky', 'stiff', 'still', 'stingy',
+     'stormy', 'straight', 'straightforward', 'strange', 'strategic',
+     'strict', 'striking', 'striped', 'strong', 'structural', 'stuck', 
+     'stupid', 'subjective', 'subsequent', 'substantial', 'subtle', 
+     'successful', 'successive', 'sudden', 'sufficient', 'suitable',
+     'sunny', 'super', 'superb', 'superior', 'supporting', 'supposed',
+     'supreme', 'sure', 'surprised', 'surprising', 'surrounding', 
+     'surviving', 'suspicious', 'sweet', 'swift', 'swiss', 'symbolic',
+     'sympathetic', 'systematic', 'tall', 'tame', 'tan', 'tart',
+     'tasteless', 'tasty', 'technical', 'technological', 'teenage', 
+     'temporary', 'tender', 'tense', 'terrible', 'territorial', 'testy',
+     'then', 'theoretical', 'thick', 'thin', 'thirsty', 'thorough', 
+     'thoughtful', 'thoughtless', 'thundering', 'tight', 'tiny', 'tired',
+     'top', 'tory', 'total', 'tough', 'toxic', 'traditional', 'tragic', 
+     'tremendous', 'tricky', 'tropical', 'troubled', 'turkish', 'typical', 
+     'ugliest', 'ugly', 'ultimate', 'unable', 'unacceptable', 'unaware', 
+     'uncertain', 'unchanged', 'uncomfortable', 'unconscious', 'underground',
+     'underlying', 'unemployed', 'uneven', 'unexpected', 'unfair', 
+     'unfortunate', 'unhappy', 'uniform', 'uninterested', 'unique', 'united',
+     'universal', 'unknown', 'unlikely', 'unnecessary', 'unpleasant', 
+     'unsightly', 'unusual', 'unwilling', 'upper', 'upset', 'uptight', 
+     'urban', 'urgent', 'used', 'useful', 'useless', 'usual', 'vague', 
+     'valid', 'valuable', 'variable', 'varied', 'various', 'varying', 'vast',
+     'verbal', 'vertical', 'very', 'victorian', 'victorious', 'video-taped', 
+     'violent', 'visible', 'visiting', 'visual', 'vital', 'vivacious', 
+     'vivid', 'vocational', 'voiceless', 'voluntary', 'vulnerable', 
+     'wandering', 'warm', 'wasteful', 'watery', 'weak', 'wealthy', 'weary', 
+     'wee', 'weekly', 'weird', 'welcome', 'well', 'well-known', 'welsh', 
+     'western', 'wet', 'whispering', 'white', 'whole', 'wicked', 'wide',
+     'wide-eyed', 'widespread', 'wild', 'willing', 'wise', 'witty', 
+     'wonderful', 'wooden', 'working', 'working-class', 'worldwide',
+     'worried', 'worrying', 'worthwhile', 'worthy', 'written', 'wrong',
+     'yellow', 'young', 'yummy', 'zany', 'zealous']
+b = ['abiding', 'accelerating', 'accepting', 'accomplishing', 'achieving', 
+'acquiring', 'acteding', 'activating', 'adapting', 'adding', 'addressing', 
+'administering', 'admiring', 'admiting', 'adopting', 'advising', 'affording', 
+'agreeing', 'alerting', 'alighting', 'allowing', 'altereding', 'amusing', 
+'analyzing', 'announcing', 'annoying', 'answering', 'anticipating', 
+'apologizing', 'appearing', 'applauding', 'applieding', 'appointing',
+ 'appraising', 'appreciating', 'approving', 'arbitrating', 'arguing', 
+ 'arising', 'arranging', 'arresting', 'arriving', 'ascertaining', 'asking', 
+ 'assembling', 'assessing', 'assisting', 'assuring', 'attaching', 'attacking', 
+ 'attaining', 'attempting', 'attending', 'attracting', 'auditeding', 'avoiding',
+ 'awaking', 'backing', 'baking', 'balancing', 'baning', 'banging', 'baring', 
+ 'bating', 'bathing', 'battling', 'bing', 'beaming', 'bearing', 'beating', 
+ 'becoming', 'beging', 'begining', 'behaving', 'beholding', 'belonging', 
+ 'bending', 'beseting', 'beting', 'biding', 'binding', 'biting', 'bleaching',
+ 'bleeding', 'blessing', 'blinding', 'blinking', 'bloting', 'blowing', 
+ 'blushing', 'boasting', 'boiling', 'bolting', 'bombing', 'booking', 
+ 'boring', 'borrowing', 'bouncing', 'bowing', 'boxing', 'braking', 
+ 'branching', 'breaking', 'breathing', 'breeding', 'briefing', 'bringing',
+ 'broadcasting', 'bruising', 'brushing', 'bubbling', 'budgeting', 'building', 
+ 'bumping', 'burning', 'bursting', 'burying', 'busting', 'buying', 'buzing', 
+ 'calculating', 'calling', 'camping', 'caring', 'carrying', 'carving', 
+ 'casting', 'cataloging', 'catching', 'causing', 'challenging', 'changing',
+ 'charging', 'charting', 'chasing', 'cheating', 'checking', 'cheering', 
+ 'chewing', 'choking', 'choosing', 'choping', 'claiming', 'claping', 
+ 'clarifying', 'classifying', 'cleaning', 'clearing', 'clinging', 'cliping',
+ 'closing', 'clothing', 'coaching', 'coiling', 'collecting', 'coloring', 
+ 'combing', 'coming', 'commanding', 'communicating', 'comparing', 'competing',
+ 'compiling', 'complaining', 'completing', 'composing', 'computing',
+ 'conceiving', 'concentrating', 'conceptualizing', 'concerning', 'concluding',
+ 'conducting', 'confessing', 'confronting', 'confusing', 'connecting', 
+ 'conserving', 'considering', 'consisting', 'consolidating', 'constructing',
+ 'consulting', 'containing', 'continuing', 'contracting', 'controling', 
+ 'converting', 'coordinating', 'copying', 'correcting', 'correlating',
+ 'costing', 'coughing', 'counseling', 'counting', 'covering', 'cracking',
+ 'crashing', 'crawling', 'creating', 'creeping', 'critiquing', 'crossing', 
+ 'crushing', 'crying', 'curing', 'curling', 'curving', 'cuting', 'cycling',
+ 'daming', 'damaging', 'dancing', 'daring', 'dealing', 'decaying', 'deceiving',
+ 'deciding', 'decorating', 'defining', 'delaying', 'delegating', 'delighting',
+ 'delivering', 'demonstrating', 'depending', 'describing', 'deserting', 
+ 'deserving', 'designing', 'destroying', 'detailing', 'detecting', 
+ 'determining', 'developing', 'devising', 'diagnosing', 'diging', 
+ 'directing', 'disagreing', 'disappearing', 'disapproving', 'disarming', 
+ 'discovering', 'disliking', 'dispensing', 'displaying', 'disproving',
+ 'dissecting', 'distributing', 'diving', 'diverting', 'dividing', 'doing',
+ 'doubling', 'doubting', 'drafting', 'draging', 'draining', 'dramatizing', 
+ 'drawing', 'dreaming', 'dressing', 'drinking', 'driping', 'driving', 
+ 'dropping', 'drowning', 'druming', 'drying', 'dusting', 'dwelling',
+ 'earning', 'eating', 'editeding', 'educating', 'eliminating',
+ 'embarrassing', 'employing', 'emptying', 'enacteding', 'encouraging',
+ 'ending', 'enduring', 'enforcing', 'engineering', 'enhancing',
+ 'enjoying', 'enlisting', 'ensuring', 'entering', 'entertaining',
+ 'escaping', 'establishing', 'estimating', 'evaluating', 'examining',
+ 'exceeding', 'exciting', 'excusing', 'executing', 'exercising', 'exhibiting',
+ 'existing', 'expanding', 'expecting', 'expediting', 'experimenting', 
+ 'explaining', 'exploding', 'expressing', 'extending', 'extracting', 
+ 'facing', 'facilitating', 'fading', 'failing', 'fancying', 'fastening', 
+ 'faxing', 'fearing', 'feeding', 'feeling', 'fencing', 'fetching', 'fighting', 
+ 'filing', 'filling', 'filming', 'finalizing', 'financing', 'finding',
+ 'firing', 'fiting', 'fixing', 'flaping', 'flashing', 'fleing', 'flinging',
+ 'floating', 'flooding', 'flowing', 'flowering', 'flying', 'folding', 
+ 'following', 'fooling', 'forbiding', 'forcing', 'forecasting', 'foregoing', 
+ 'foreseing', 'foretelling', 'forgeting', 'forgiving', 'forming', 
+ 'formulating', 'forsaking', 'framing', 'freezing', 'frightening', 'frying',
+ 'gathering', 'gazing', 'generating', 'geting', 'giving', 'glowing', 'gluing', 
+ 'going', 'governing', 'grabing', 'graduating', 'grating', 'greasing', 'greeting',
+ 'grinning', 'grinding', 'griping', 'groaning', 'growing', 'guaranteeing',
+ 'guarding', 'guessing', 'guiding', 'hammering', 'handing', 'handling', 
+ 'handwriting', 'hanging', 'happening', 'harassing', 'harming', 'hating',
+ 'haunting', 'heading', 'healing', 'heaping', 'hearing', 'heating', 'helping', 
+ 'hiding', 'hitting', 'holding', 'hooking', 'hoping', 'hopping', 'hovering',
+ 'hugging', 'hmuming', 'hunting', 'hurrying', 'hurting', 'hypothesizing', 
+ 'identifying', 'ignoring', 'illustrating', 'imagining', 'implementing', 
+ 'impressing', 'improving', 'improvising', 'including', 'increasing', 
+ 'inducing', 'influencing', 'informing', 'initiating', 'injecting', 
+ 'injuring', 'inlaying', 'innovating', 'inputing', 'inspecting', 
+ 'inspiring', 'installing', 'instituting', 'instructing', 'insuring', 
+ 'integrating', 'intending', 'intensifying', 'interesting', 
+ 'interfering', 'interlaying', 'interpreting', 'interrupting', 
+ 'interviewing', 'introducing', 'inventing', 'inventorying', 
+ 'investigating', 'inviting', 'irritating', 'itching', 'jailing', 
+ 'jamming', 'jogging', 'joining', 'joking', 'judging', 'juggling', 'jumping',
+ 'justifying', 'keeping', 'kepting', 'kicking', 'killing', 'kissing', 'kneeling',
+ 'kniting', 'knocking', 'knotting', 'knowing', 'labeling', 'landing', 'lasting',
+ 'laughing', 'launching', 'laying', 'leading', 'leaning', 'leaping', 'learning', 
+ 'leaving', 'lecturing', 'leding', 'lending', 'leting', 'leveling', 
+ 'licensing', 'licking', 'lying', 'lifteding', 'lighting', 'lightening',
+ 'liking', 'listing', 'listening', 'living', 'loading', 'locating', 
+ 'locking', 'loging', 'longing', 'looking', 'losing', 'loving', 
+ 'maintaining', 'making', 'maning', 'managing', 'manipulating', 
+ 'manufacturing', 'mapping', 'marching', 'marking', 'marketing',
+ 'marrying', 'matching', 'mating', 'mattering', 'meaning', 'measuring',
+ 'meddling', 'mediating', 'meeting', 'melting', 'melting', 'memorizing',
+ 'mending', 'mentoring', 'milking', 'mining', 'misleading', 'missing',
+ 'misspelling', 'mistaking', 'misunderstanding', 'mixing', 'moaning', 
+ 'modeling', 'modifying', 'monitoring', 'mooring', 'motivating',
+ 'mourning', 'moving', 'mowing', 'muddling', 'muging', 'multiplying', 
+ 'murdering', 'nailing', 'naming', 'navigating', 'needing', 'negotiating', 
+ 'nesting', 'noding', 'nominating', 'normalizing', 'noting', 'noticing', 
+ 'numbering', 'obeying', 'objecting', 'observing', 'obtaining', 'occuring', 
+ 'offending', 'offering', 'officiating', 'opening', 'operating', 'ordering', 
+ 'organizing', 'orienteding', 'originating', 'overcoming', 'overdoing', 
+ 'overdrawing', 'overflowing', 'overhearing', 'overtaking', 'overthrowing',
+ 'owing', 'owning', 'packing', 'paddling', 'painting', 'parking', 'parting', 
+ 'participating', 'passing', 'pasting', 'pating', 'pausing', 'paying',
+ 'pecking', 'pedaling', 'peeling', 'peeping', 'perceiving', 'perfecting', 
+ 'performing', 'permiting', 'persuading', 'phoning', 'photographing',
+ 'picking', 'piloting', 'pinching', 'pining', 'pinpointing', 'pioneering',
+ 'placing', 'planing', 'planting', 'playing', 'pleading', 'pleasing',
+ 'plugging', 'pointing', 'poking', 'polishing', 'poping', 'possessing',
+ 'posting', 'pouring', 'practicing', 'praiseding', 'praying', 'preaching', 
+ 'preceding', 'predicting', 'prefering', 'preparing', 'prescribing', 
+ 'presenting', 'preserving', 'preseting', 'presiding', 'pressing', 
+ 'pretending', 'preventing', 'pricking', 'printing', 'processing', 
+ 'procuring', 'producing', 'professing', 'programing', 'progressing', 
+ 'projecting', 'promising', 'promoting', 'proofreading', 'proposing', 
+ 'protecting', 'proving', 'providing', 'publicizing', 'pulling', 'pumping',
+ 'punching', 'puncturing', 'punishing', 'purchasing', 'pushing', 'puting',
+ 'qualifying', 'questioning', 'queuing', 'quiting', 'racing', 'radiating',
+ 'raining', 'raising', 'ranking', 'rating', 'reaching', 'reading', 
+ 'realigning', 'realizing', 'reasoning', 'receiving', 'recognizing', 
+ 'recommending', 'reconciling', 'recording', 'recruiting', 'reducing', 
+ 'referring', 'reflecting', 'refusing', 'regreting', 'regulating', 
+ 'rehabilitating', 'reigning', 'reinforcing', 'rejecting', 'rejoicing',
+ 'relating', 'relaxing', 'releasing', 'relying', 'remaining', 'remembering',
+ 'reminding', 'removing', 'rendering', 'reorganizing', 'repairing',
+ 'repeating', 'replacing', 'replying', 'reporting', 'representing',
+ 'reproducing', 'requesting', 'rescuing', 'researching', 'resolving', 
+ 'responding', 'restoreding', 'restructuring', 'retiring', 'retrieving',
+ 'returning', 'reviewing', 'revising', 'rhyming', 'riding', 'riding', 
+ 'ringing', 'rinsing', 'rising', 'risking', 'robing', 'rocking', 'rolling',
+ 'roting', 'rubing', 'ruining', 'ruling', 'runing', 'rushing', 'sacking',
+ 'sailing', 'satisfying', 'saving', 'sawing', 'saying', 'scaring', 
+ 'scattering', 'scheduling', 'scolding', 'scorching', 'scraping', 
+ 'scratching', 'screaming', 'screwing', 'scribbling', 'scrubing', 
+ 'sealing', 'searching', 'securing', 'seing', 'seeking', 'selecting', 
+ 'selling', 'sending', 'sensing', 'separating', 'serving', 'servicing', 
+ 'seting', 'settling', 'sewing', 'shading', 'shaking', 'shaping', 
+ 'sharing', 'shaving', 'shearing', 'sheding', 'sheltering', 'shining', 
+ 'shivering', 'shocking', 'shoing', 'shooting', 'shoping', 'showing', 
+ 'shrinking', 'shruging', 'shuting', 'sighing', 'signing', 'signaling',
+ 'simplifying', 'sining', 'singing', 'sinking', 'siping', 'siting',
+ 'sketching', 'skiing', 'skiping', 'slaping', 'slaying', 'sleeping',
+ 'sliding', 'slinging', 'slinking', 'sliping', 'sliting', 'slowing',
+ 'smashing', 'smelling', 'smiling', 'smiting', 'smoking', 'snatching',
+ 'sneaking', 'sneezing', 'sniffing', 'snoring', 'snowing', 'soaking', 
+ 'solving', 'soothing', 'soothsaying', 'sorting', 'sounding', 'sowing', 
+ 'sparing', 'sparking', 'sparkling', 'speaking', 'specifying', 'speeding',
+ 'spelling', 'spending', 'spilling', 'spining', 'spiting', 'spliting',
+ 'spoiling', 'spoting', 'spraying', 'spreading', 'springing', 'sprouting', 
+ 'squashing', 'squeaking', 'squealing', 'squeezing', 'staining', 'stamping',
+ 'standing', 'staring', 'starting', 'staying', 'stealing', 'steering', 
+ 'stepping', 'sticking', 'stimulating', 'stinging', 'stinking', 'stirring', 
+ 'stitching', 'stoping', 'storing', 'straping', 'streamlining', 
+ 'strengthening', 'stretching', 'striding', 'striking', 'stringing', 
+ 'stripping', 'striving', 'stroking', 'structuring', 'studying', 
+ 'stuffing', 'subleting', 'subtracting', 'succeeding', 'sucking', 
+ 'suffering', 'suggesting', 'suiting', 'summarizing', 'supervising',
+ 'supplying', 'supporting', 'supposing', 'surprising', 'surrounding', 
+ 'suspecting', 'suspending', 'swearing', 'sweating', 'sweeping', 'swelling', 
+ 'swimming', 'swinging', 'switching', 'symbolizing', 'synthesizing',
+ 'systemizing', 'tabulating', 'taking', 'talking', 'taming', 'taping', 
+ 'targeting', 'tasting', 'teaching', 'tearing', 'teasing', 'telephoning', 
+ 'telling', 'tempting', 'terrifying', 'testing', 'thanking', 'thawing', 
+ 'thinking', 'thriving', 'throwing', 'thrusting', 'ticking', 'tickling', 
+ 'tying', 'timing', 'tiping', 'tiring', 'touching', 'touring', 'towing',
+ 'tracing', 'trading', 'training', 'transcribing', 'transfering',
+ 'transforming', 'translating', 'transporting', 'traping', 'traveling',
+ 'treading', 'treating', 'trembling', 'tricking', 'triping', 'troting', 
+ 'troubling', 'troubleshooting', 'trusting', 'trying', 'tuging', 'tumbling',
+ 'turning', 'tutoring', 'twisting', 'typing', 'undergoing', 'understanding',
+ 'undertaking', 'undressing', 'unfastening', 'unifying', 'uniting', 
+ 'unlocking', 'unpacking', 'untidying', 'updating', 'upgrading', 
+ 'upholding', 'upseting', 'using', 'utilizing', 'vanishing', 'verbalizing',
+ 'verifying', 'vexing', 'visiting', 'wailing', 'waiting', 'waking', 
+ 'walking', 'wandering', 'wanting', 'warming', 'warning', 'washing', 
+ 'wasting', 'watching', 'watering', 'waving', 'wearing', 'weaving', 
+ 'wedding', 'weeping', 'weighing', 'welcoming', 'wending', 'weting', 
+ 'whining', 'whiping', 'whirling', 'whispering', 'whistling', 'wining', 
+ 'winding', 'winking', 'wiping', 'wishing', 'withdrawing', 'withholding',
+ 'withstanding', 'wobbling', 'wondering', 'working', 'worrying', 'wrapping', 
+ 'wrecking', 'wrestling', 'wriggling', 'wringing', 'writing', 'x-raying',
+ 'yawning', 'yelling', 'zipping', 'zooming']

src/models/big/cheat sheet

@@ -0,0 +1,30 @@
+from big.BigGAN2 import Generator,Discriminator
+
+from big.losses import generator_loss, discriminator_loss
+
+generator = Generator().cuda()
+discriminator = Discriminator().cuda()
+
+label_transformed_fake = label_fc_net(label_fake)
+label_transformed_real = label_fc_net(label_real)
+
+generated_images = generator(decoder_input,label_transformed_fake)
+
+#disc training
+
+prediction_fake = discriminator(generated_images.detach(),label_transformed_fake).view(-1) 
+prediction_real = discriminator(images,label_transformed_real).view(-1) 
+
+d_loss_real,d_loss_fake =  discriminator_loss(prediction_fake,prediction_real)
+
+discriminator.optim.step()
+
+
+#gen training
+
+prediction = discriminator(generated_images,label_transformed_fake).view(-1)
+      
+g_loss = generator_loss( prediction)
+g_loss.backward()
+
+generator.optim.step()

src/models/big/datasets.py

@@ -0,0 +1,362 @@
+''' Datasets
+    This file contains definitions for our CIFAR, ImageFolder, and HDF5 datasets
+'''
+import os
+import os.path
+import sys
+from PIL import Image
+import numpy as np
+from tqdm import tqdm, trange
+
+import torchvision.datasets as dset
+import torchvision.transforms as transforms
+from torchvision.datasets.utils import download_url, check_integrity
+import torch.utils.data as data
+from torch.utils.data import DataLoader
+         
+IMG_EXTENSIONS = ['.jpg', '.jpeg', '.png', '.ppm', '.bmp', '.pgm']
+
+
+def is_image_file(filename):
+    """Checks if a file is an image.
+
+    Args:
+        filename (string): path to a file
+
+    Returns:
+        bool: True if the filename ends with a known image extension
+    """
+    filename_lower = filename.lower()
+    return any(filename_lower.endswith(ext) for ext in IMG_EXTENSIONS)
+
+
+def find_classes(dir):
+    classes = [d for d in os.listdir(dir) if os.path.isdir(os.path.join(dir, d))]
+    classes.sort()
+    class_to_idx = {classes[i]: i for i in range(len(classes))}
+    return classes, class_to_idx
+
+
+def make_dataset(dir, class_to_idx):
+  images = []
+  dir = os.path.expanduser(dir)
+  for target in tqdm(sorted(os.listdir(dir))):
+    d = os.path.join(dir, target)
+    if not os.path.isdir(d):
+      continue
+
+    for root, _, fnames in sorted(os.walk(d)):
+      for fname in sorted(fnames):
+        if is_image_file(fname):
+          path = os.path.join(root, fname)
+          item = (path, class_to_idx[target])
+          images.append(item)
+
+  return images
+
+
+def pil_loader(path):
+    # open path as file to avoid ResourceWarning (https://github.com/python-pillow/Pillow/issues/835)
+  with open(path, 'rb') as f:
+    img = Image.open(f)
+    return img.convert('RGB')
+
+
+def accimage_loader(path):
+  import accimage
+  try:
+    return accimage.Image(path)
+  except IOError:
+    # Potentially a decoding problem, fall back to PIL.Image
+    return pil_loader(path)
+
+
+def default_loader(path):
+  from torchvision import get_image_backend
+  if get_image_backend() == 'accimage':
+    return accimage_loader(path)
+  else:
+    return pil_loader(path)
+
+
+class ImageFolder(data.Dataset):
+  """A generic data loader where the images are arranged in this way: ::
+
+      root/dogball/xxx.png
+      root/dogball/xxy.png
+      root/dogball/xxz.png
+
+      root/cat/123.png
+      root/cat/nsdf3.png
+      root/cat/asd932_.png
+
+  Args:
+      root (string): Root directory path.
+      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).
+      imgs (list): List of (image path, class_index) tuples
+  """
+
+  def __init__(self, root, transform=None, target_transform=None,
+               loader=default_loader, load_in_mem=False, 
+               index_filename='imagenet_imgs.npz', **kwargs):
+    classes, class_to_idx = find_classes(root)
+    # Load pre-computed image directory walk
+    if os.path.exists(index_filename):
+      print('Loading pre-saved Index file %s...' % index_filename)
+      imgs = np.load(index_filename)['imgs']
+    # If first time, walk the folder directory and save the 
+    # results to a pre-computed file.
+    else:
+      print('Generating  Index file %s...' % index_filename)
+      imgs = make_dataset(root, class_to_idx)
+      np.savez_compressed(index_filename, **{'imgs' : imgs})
+    if len(imgs) == 0:
+      raise(RuntimeError("Found 0 images in subfolders of: " + root + "\n"
+                           "Supported image extensions are: " + ",".join(IMG_EXTENSIONS)))
+
+    self.root = root
+    self.imgs = imgs
+    self.classes = classes
+    self.class_to_idx = class_to_idx
+    self.transform = transform
+    self.target_transform = target_transform
+    self.loader = loader
+    self.load_in_mem = load_in_mem
+    
+    if self.load_in_mem:
+      print('Loading all images into memory...')
+      self.data, self.labels = [], []
+      for index in tqdm(range(len(self.imgs))):
+        path, target = imgs[index][0], imgs[index][1]
+        self.data.append(self.transform(self.loader(path)))
+        self.labels.append(target)
+          
+
+  def __getitem__(self, index):
+    """
+    Args:
+        index (int): Index
+
+    Returns:
+        tuple: (image, target) where target is class_index of the target class.
+    """
+    if self.load_in_mem:
+        img = self.data[index]
+        target = self.labels[index]
+    else:
+      path, target = self.imgs[index]
+      img = self.loader(str(path))
+      if self.transform is not None:
+        img = self.transform(img)
+    
+    if self.target_transform is not None:
+      target = self.target_transform(target)
+    
+    # print(img.size(), target)
+    return img, int(target)
+
+  def __len__(self):
+    return len(self.imgs)
+
+  def __repr__(self):
+    fmt_str = 'Dataset ' + self.__class__.__name__ + '\n'
+    fmt_str += '    Number of datapoints: {}\n'.format(self.__len__())
+    fmt_str += '    Root Location: {}\n'.format(self.root)
+    tmp = '    Transforms (if any): '
+    fmt_str += '{0}{1}\n'.format(tmp, self.transform.__repr__().replace('\n', '\n' + ' ' * len(tmp)))
+    tmp = '    Target Transforms (if any): '
+    fmt_str += '{0}{1}'.format(tmp, self.target_transform.__repr__().replace('\n', '\n' + ' ' * len(tmp)))
+    return fmt_str
+        
+
+''' ILSVRC_HDF5: A dataset to support I/O from an HDF5 to avoid
+    having to load individual images all the time. '''
+import h5py as h5
+import torch
+class ILSVRC_HDF5(data.Dataset):
+  def __init__(self, root, transform=None, target_transform=None,
+               load_in_mem=False, train=True,download=False, validate_seed=0,
+               val_split=0, **kwargs): # last four are dummies
+      
+    self.root = root
+    self.num_imgs = len(h5.File(root, 'r')['labels'])
+    
+    # self.transform = transform
+    self.target_transform = target_transform   
+    
+    # Set the transform here
+    self.transform = transform
+    
+    # load the entire dataset into memory? 
+    self.load_in_mem = load_in_mem
+    
+    # If loading into memory, do so now
+    if self.load_in_mem:
+      print('Loading %s into memory...' % root)
+      with h5.File(root,'r') as f:
+        self.data = f['imgs'][:]
+        self.labels = f['labels'][:]
+
+  def __getitem__(self, index):
+    """
+    Args:
+        index (int): Index
+
+    Returns:
+        tuple: (image, target) where target is class_index of the target class.
+    """
+    # If loaded the entire dataset in RAM, get image from memory
+    if self.load_in_mem:
+      img = self.data[index]
+      target = self.labels[index]
+    
+    # Else load it from disk
+    else:
+      with h5.File(self.root,'r') as f:
+        img = f['imgs'][index]
+        target = f['labels'][index]
+    
+   
+    # if self.transform is not None:
+        # img = self.transform(img)
+    # Apply my own transform
+    img = ((torch.from_numpy(img).float() / 255) - 0.5) * 2
+    
+    if self.target_transform is not None:
+      target = self.target_transform(target)
+    
+    return img, int(target)
+
+  def __len__(self):
+      return self.num_imgs
+      # return len(self.f['imgs'])
+
+import pickle
+class CIFAR10(dset.CIFAR10):
+
+  def __init__(self, root, train=True,
+           transform=None, target_transform=None,
+           download=True, validate_seed=0,
+           val_split=0, load_in_mem=True, **kwargs):
+    self.root = os.path.expanduser(root)
+    self.transform = transform
+    self.target_transform = target_transform
+    self.train = train  # training set or test set
+    self.val_split = val_split
+
+    if download:
+      self.download()
+
+    if not self._check_integrity():
+      raise RuntimeError('Dataset not found or corrupted.' +
+                           ' You can use download=True to download it')
+
+    # now load the picked numpy arrays    
+    self.data = []
+    self.labels= []
+    for fentry in self.train_list:
+      f = fentry[0]
+      file = os.path.join(self.root, self.base_folder, f)
+      fo = open(file, 'rb')
+      if sys.version_info[0] == 2:
+        entry = pickle.load(fo)
+      else:
+        entry = pickle.load(fo, encoding='latin1')
+      self.data.append(entry['data'])
+      if 'labels' in entry:
+        self.labels += entry['labels']
+      else:
+        self.labels += entry['fine_labels']
+      fo.close()
+        
+    self.data = np.concatenate(self.data)
+    # Randomly select indices for validation
+    if self.val_split > 0:
+      label_indices = [[] for _ in range(max(self.labels)+1)]
+      for i,l in enumerate(self.labels):
+        label_indices[l] += [i]  
+      label_indices = np.asarray(label_indices)
+      
+      # randomly grab 500 elements of each class
+      np.random.seed(validate_seed)
+      self.val_indices = []           
+      for l_i in label_indices:
+        self.val_indices += list(l_i[np.random.choice(len(l_i), int(len(self.data) * val_split) // (max(self.labels) + 1) ,replace=False)])
+    
+    if self.train=='validate':    
+      self.data = self.data[self.val_indices]
+      self.labels = list(np.asarray(self.labels)[self.val_indices])
+      
+      self.data = self.data.reshape((int(50e3 * self.val_split), 3, 32, 32))
+      self.data = self.data.transpose((0, 2, 3, 1))  # convert to HWC
+    
+    elif self.train:
+      print(np.shape(self.data))
+      if self.val_split > 0:
+        self.data = np.delete(self.data,self.val_indices,axis=0)
+        self.labels = list(np.delete(np.asarray(self.labels),self.val_indices,axis=0))
+          
+      self.data = self.data.reshape((int(50e3 * (1.-self.val_split)), 3, 32, 32))
+      self.data = self.data.transpose((0, 2, 3, 1))  # convert to HWC
+    else:
+      f = self.test_list[0][0]
+      file = os.path.join(self.root, self.base_folder, f)
+      fo = open(file, 'rb')
+      if sys.version_info[0] == 2:
+        entry = pickle.load(fo)
+      else:
+        entry = pickle.load(fo, encoding='latin1')
+      self.data = entry['data']
+      if 'labels' in entry:
+        self.labels = entry['labels']
+      else:
+        self.labels = entry['fine_labels']
+      fo.close()
+      self.data = self.data.reshape((10000, 3, 32, 32))
+      self.data = self.data.transpose((0, 2, 3, 1))  # convert to HWC
+      
+  def __getitem__(self, index):
+    """
+    Args:
+        index (int): Index
+    Returns:
+        tuple: (image, target) where target is index of the target class.
+    """
+    img, target = self.data[index], self.labels[index]
+
+    # doing this so that it is consistent with all other datasets
+    # to return a PIL Image
+    img = Image.fromarray(img)
+
+    if self.transform is not None:
+      img = self.transform(img)
+
+    if self.target_transform is not None:
+      target = self.target_transform(target)
+
+    return img, target
+      
+  def __len__(self):
+      return len(self.data)
+
+
+class CIFAR100(CIFAR10):
+    base_folder = 'cifar-100-python'
+    url = "http://www.cs.toronto.edu/~kriz/cifar-100-python.tar.gz"
+    filename = "cifar-100-python.tar.gz"
+    tgz_md5 = 'eb9058c3a382ffc7106e4002c42a8d85'
+    train_list = [
+        ['train', '16019d7e3df5f24257cddd939b257f8d'],
+    ]
+
+    test_list = [
+        ['test', 'f0ef6b0ae62326f3e7ffdfab6717acfc'],
+    ]

src/models/big/layers.py

@@ -0,0 +1,456 @@
+''' Layers
+    This file contains various layers for the BigGAN models.
+'''
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.nn import Parameter as P
+
+from src.models.big.sync_batchnorm import SynchronizedBatchNorm2d as SyncBN2d
+
+
+# Projection of x onto y
+def proj(x, y):
+  return torch.mm(y, x.t()) * y / torch.mm(y, y.t())
+
+
+# Orthogonalize x wrt list of vectors ys
+def gram_schmidt(x, ys):
+  for y in ys:
+    x = x - proj(x, y)
+  return x
+
+
+# Apply num_itrs steps of the power method to estimate top N singular values.
+def power_iteration(W, u_, update=True, eps=1e-12):
+  # Lists holding singular vectors and values
+  us, vs, svs = [], [], []
+  for i, u in enumerate(u_):
+    # Run one step of the power iteration
+    with torch.no_grad():
+      v = torch.matmul(u, W)
+      # Run Gram-Schmidt to subtract components of all other singular vectors
+      v = F.normalize(gram_schmidt(v, vs), eps=eps)
+      # Add to the list
+      vs += [v]
+      # Update the other singular vector
+      u = torch.matmul(v, W.t())
+      # Run Gram-Schmidt to subtract components of all other singular vectors
+      u = F.normalize(gram_schmidt(u, us), eps=eps)
+      # Add to the list
+      us += [u]
+      if update:
+        u_[i][:] = u
+    # Compute this singular value and add it to the list
+    svs += [torch.squeeze(torch.matmul(torch.matmul(v, W.t()), u.t()))]
+    #svs += [torch.sum(F.linear(u, W.transpose(0, 1)) * v)]
+  return svs, us, vs
+
+
+# Convenience passthrough function
+class identity(nn.Module):
+  def forward(self, input):
+    return input
+ 
+
+# Spectral normalization base class 
+class SN(object):
+  def __init__(self, num_svs, num_itrs, num_outputs, transpose=False, eps=1e-12):
+    # Number of power iterations per step
+    self.num_itrs = num_itrs
+    # Number of singular values
+    self.num_svs = num_svs
+    # Transposed?
+    self.transpose = transpose
+    # Epsilon value for avoiding divide-by-0
+    self.eps = eps
+    # Register a singular vector for each sv
+    for i in range(self.num_svs):
+      self.register_buffer('u%d' % i, torch.randn(1, num_outputs))
+      self.register_buffer('sv%d' % i, torch.ones(1))
+  
+  # Singular vectors (u side)
+  @property
+  def u(self):
+    return [getattr(self, 'u%d' % i) for i in range(self.num_svs)]
+
+  # Singular values; 
+  # note that these buffers are just for logging and are not used in training. 
+  @property
+  def sv(self):
+   return [getattr(self, 'sv%d' % i) for i in range(self.num_svs)]
+   
+  # Compute the spectrally-normalized weight
+  def W_(self):
+    W_mat = self.weight.view(self.weight.size(0), -1)
+    if self.transpose:
+      W_mat = W_mat.t()
+    # Apply num_itrs power iterations
+    for _ in range(self.num_itrs):
+      svs, us, vs = power_iteration(W_mat, self.u, update=self.training, eps=self.eps) 
+    # Update the svs
+    if self.training:
+      with torch.no_grad(): # Make sure to do this in a no_grad() context or you'll get memory leaks!
+        for i, sv in enumerate(svs):
+          self.sv[i][:] = sv     
+    return self.weight / svs[0]
+
+
+# 2D Conv layer with spectral norm
+class SNConv2d(nn.Conv2d, SN):
+  def __init__(self, in_channels, out_channels, kernel_size, stride=1,
+             padding=0, dilation=1, groups=1, bias=True, 
+             num_svs=1, num_itrs=1, eps=1e-12):
+    nn.Conv2d.__init__(self, in_channels, out_channels, kernel_size, stride, 
+                     padding, dilation, groups, bias)
+    SN.__init__(self, num_svs, num_itrs, out_channels, eps=eps)    
+  def forward(self, x):
+    return F.conv2d(x, self.W_(), self.bias, self.stride, 
+                    self.padding, self.dilation, self.groups)
+
+
+# Linear layer with spectral norm
+class SNLinear(nn.Linear, SN):
+  def __init__(self, in_features, out_features, bias=True,
+               num_svs=1, num_itrs=1, eps=1e-12):
+    nn.Linear.__init__(self, in_features, out_features, bias)
+    SN.__init__(self, num_svs, num_itrs, out_features, eps=eps)
+  def forward(self, x):
+    return F.linear(x, self.W_(), self.bias)
+
+
+# Embedding layer with spectral norm
+# We use num_embeddings as the dim instead of embedding_dim here
+# for convenience sake
+class SNEmbedding(nn.Embedding, SN):
+  def __init__(self, num_embeddings, embedding_dim, padding_idx=None, 
+               max_norm=None, norm_type=2, scale_grad_by_freq=False,
+               sparse=False, _weight=None,
+               num_svs=1, num_itrs=1, eps=1e-12):
+    nn.Embedding.__init__(self, num_embeddings, embedding_dim, padding_idx,
+                          max_norm, norm_type, scale_grad_by_freq, 
+                          sparse, _weight)
+    SN.__init__(self, num_svs, num_itrs, num_embeddings, eps=eps)
+  def forward(self, x):
+    return F.embedding(x, self.W_())
+
+
+# A non-local block as used in SA-GAN
+# Note that the implementation as described in the paper is largely incorrect;
+# refer to the released code for the actual implementation.
+class Attention(nn.Module):
+  def __init__(self, ch, which_conv=SNConv2d, name='attention'):
+    super(Attention, self).__init__()
+    # Channel multiplier
+    self.ch = ch
+    self.which_conv = which_conv
+    self.theta = self.which_conv(self.ch, self.ch // 8, kernel_size=1, padding=0, bias=False)
+    self.phi = self.which_conv(self.ch, self.ch // 8, kernel_size=1, padding=0, bias=False)
+    self.g = self.which_conv(self.ch, self.ch // 2, kernel_size=1, padding=0, bias=False)
+    self.o = self.which_conv(self.ch // 2, self.ch, kernel_size=1, padding=0, bias=False)
+    # Learnable gain parameter
+    self.gamma = P(torch.tensor(0.), requires_grad=True)
+  def forward(self, x, y=None):
+    # Apply convs
+    theta = self.theta(x)
+    phi = F.max_pool2d(self.phi(x), [2,2])
+    g = F.max_pool2d(self.g(x), [2,2])    
+    # Perform reshapes
+    theta = theta.view(-1, self. ch // 8, x.shape[2] * x.shape[3])
+    phi = phi.view(-1, self. ch // 8, x.shape[2] * x.shape[3] // 4)
+    g = g.view(-1, self. ch // 2, x.shape[2] * x.shape[3] // 4)
+    # Matmul and softmax to get attention maps
+    beta = F.softmax(torch.bmm(theta.transpose(1, 2), phi), -1)
+    # Attention map times g path
+    o = self.o(torch.bmm(g, beta.transpose(1,2)).view(-1, self.ch // 2, x.shape[2], x.shape[3]))
+    return self.gamma * o + x
+
+
+# Fused batchnorm op
+def fused_bn(x, mean, var, gain=None, bias=None, eps=1e-5):
+  # Apply scale and shift--if gain and bias are provided, fuse them here
+  # Prepare scale
+  scale = torch.rsqrt(var + eps)
+  # If a gain is provided, use it
+  if gain is not None:
+    scale = scale * gain
+  # Prepare shift
+  shift = mean * scale
+  # If bias is provided, use it
+  if bias is not None:
+    shift = shift - bias
+  return x * scale - shift
+  #return ((x - mean) / ((var + eps) ** 0.5)) * gain + bias # The unfused way.
+
+
+# Manual BN
+# Calculate means and variances using mean-of-squares minus mean-squared
+def manual_bn(x, gain=None, bias=None, return_mean_var=False, eps=1e-5):
+  # Cast x to float32 if necessary
+  float_x = x.float()
+  # Calculate expected value of x (m) and expected value of x**2 (m2)  
+  # Mean of x
+  m = torch.mean(float_x, [0, 2, 3], keepdim=True)
+  # Mean of x squared
+  m2 = torch.mean(float_x ** 2, [0, 2, 3], keepdim=True)
+  # Calculate variance as mean of squared minus mean squared.
+  var = (m2 - m **2)
+  # Cast back to float 16 if necessary
+  var = var.type(x.type())
+  m = m.type(x.type())
+  # Return mean and variance for updating stored mean/var if requested  
+  if return_mean_var:
+    return fused_bn(x, m, var, gain, bias, eps), m.squeeze(), var.squeeze()
+  else:
+    return fused_bn(x, m, var, gain, bias, eps)
+
+
+# My batchnorm, supports standing stats    
+class myBN(nn.Module):
+  def __init__(self, num_channels, eps=1e-5, momentum=0.1):
+    super(myBN, self).__init__()
+    # momentum for updating running stats
+    self.momentum = momentum
+    # epsilon to avoid dividing by 0
+    self.eps = eps
+    # Momentum
+    self.momentum = momentum
+    # Register buffers
+    self.register_buffer('stored_mean', torch.zeros(num_channels))
+    self.register_buffer('stored_var',  torch.ones(num_channels))
+    self.register_buffer('accumulation_counter', torch.zeros(1))
+    # Accumulate running means and vars
+    self.accumulate_standing = False
+    
+  # reset standing stats
+  def reset_stats(self):
+    self.stored_mean[:] = 0
+    self.stored_var[:] = 0
+    self.accumulation_counter[:] = 0
+    
+  def forward(self, x, gain, bias):
+    if self.training:
+      out, mean, var = manual_bn(x, gain, bias, return_mean_var=True, eps=self.eps)
+      # If accumulating standing stats, increment them
+      if self.accumulate_standing:
+        self.stored_mean[:] = self.stored_mean + mean.data
+        self.stored_var[:] = self.stored_var + var.data
+        self.accumulation_counter += 1.0
+      # If not accumulating standing stats, take running averages
+      else:
+        self.stored_mean[:] = self.stored_mean * (1 - self.momentum) + mean * self.momentum
+        self.stored_var[:] = self.stored_var * (1 - self.momentum) + var * self.momentum
+      return out
+    # If not in training mode, use the stored statistics
+    else:         
+      mean = self.stored_mean.view(1, -1, 1, 1)
+      var = self.stored_var.view(1, -1, 1, 1)
+      # If using standing stats, divide them by the accumulation counter   
+      if self.accumulate_standing:
+        mean = mean / self.accumulation_counter
+        var = var / self.accumulation_counter
+      return fused_bn(x, mean, var, gain, bias, self.eps)
+
+
+# Simple function to handle groupnorm norm stylization                      
+def groupnorm(x, norm_style):
+  # If number of channels specified in norm_style:
+  if 'ch' in norm_style:
+    ch = int(norm_style.split('_')[-1])
+    groups = max(int(x.shape[1]) // ch, 1)
+  # If number of groups specified in norm style
+  elif 'grp' in norm_style:
+    groups = int(norm_style.split('_')[-1])
+  # If neither, default to groups = 16
+  else:
+    groups = 16
+  return F.group_norm(x, groups)
+
+
+# Class-conditional bn
+# output size is the number of channels, input size is for the linear layers
+# Andy's Note: this class feels messy but I'm not really sure how to clean it up
+# Suggestions welcome! (By which I mean, refactor this and make a pull request
+# if you want to make this more readable/usable). 
+class ccbn(nn.Module):
+  def __init__(self, output_size, input_size, which_linear, eps=1e-5, momentum=0.1,
+               cross_replica=False, mybn=False, norm_style='bn',):
+    super(ccbn, self).__init__()
+    self.output_size, self.input_size = output_size, input_size
+    # Prepare gain and bias layers
+    self.gain = which_linear(input_size, output_size)
+    self.bias = which_linear(input_size, output_size)
+    # epsilon to avoid dividing by 0
+    self.eps = eps
+    # Momentum
+    self.momentum = momentum
+    # Use cross-replica batchnorm?
+    self.cross_replica = cross_replica
+    # Use my batchnorm?
+    self.mybn = mybn
+    # Norm style?
+    self.norm_style = norm_style
+    
+    if self.cross_replica:
+      self.bn = SyncBN2d(output_size, eps=self.eps, momentum=self.momentum, affine=False)
+    elif self.mybn:
+      self.bn = myBN(output_size, self.eps, self.momentum)
+    elif self.norm_style in ['bn', 'in']:
+      self.register_buffer('stored_mean', torch.zeros(output_size))
+      self.register_buffer('stored_var',  torch.ones(output_size)) 
+    
+    
+  def forward(self, x, y):
+    # Calculate class-conditional gains and biases
+    gain = (1 + self.gain(y)).view(y.size(0), -1, 1, 1)
+    bias = self.bias(y).view(y.size(0), -1, 1, 1)
+    # If using my batchnorm
+    if self.mybn or self.cross_replica:
+      return self.bn(x, gain=gain, bias=bias)
+    # else:
+    else:
+      if self.norm_style == 'bn':
+        out = F.batch_norm(x, self.stored_mean, self.stored_var, None, None,
+                          self.training, 0.1, self.eps)
+      elif self.norm_style == 'in':
+        out = F.instance_norm(x, self.stored_mean, self.stored_var, None, None,
+                          self.training, 0.1, self.eps)
+      elif self.norm_style == 'gn':
+        out = groupnorm(x, self.normstyle)
+      elif self.norm_style == 'nonorm':
+        out = x
+      return out * gain + bias
+  def extra_repr(self):
+    s = 'out: {output_size}, in: {input_size},'
+    s +=' cross_replica={cross_replica}'
+    return s.format(**self.__dict__)
+
+
+# Normal, non-class-conditional BN
+class bn(nn.Module):
+  def __init__(self, output_size,  eps=1e-5, momentum=0.1,
+                cross_replica=False, mybn=False):
+    super(bn, self).__init__()
+    self.output_size= output_size
+    # Prepare gain and bias layers
+    self.gain = P(torch.ones(output_size), requires_grad=True)
+    self.bias = P(torch.zeros(output_size), requires_grad=True)
+    # epsilon to avoid dividing by 0
+    self.eps = eps
+    # Momentum
+    self.momentum = momentum
+    # Use cross-replica batchnorm?
+    self.cross_replica = cross_replica
+    # Use my batchnorm?
+    self.mybn = mybn
+    
+    if self.cross_replica:
+      self.bn = SyncBN2d(output_size, eps=self.eps, momentum=self.momentum, affine=False)    
+    elif mybn:
+      self.bn = myBN(output_size, self.eps, self.momentum)
+     # Register buffers if neither of the above
+    else:     
+      self.register_buffer('stored_mean', torch.zeros(output_size))
+      self.register_buffer('stored_var',  torch.ones(output_size))
+    
+  def forward(self, x, y=None):
+    if self.cross_replica or self.mybn:
+      gain = self.gain.view(1,-1,1,1)
+      bias = self.bias.view(1,-1,1,1)
+      return self.bn(x, gain=gain, bias=bias)
+    else:
+      return F.batch_norm(x, self.stored_mean, self.stored_var, self.gain,
+                          self.bias, self.training, self.momentum, self.eps)
+
+                          
+# Generator blocks
+# Note that this class assumes the kernel size and padding (and any other
+# settings) have been selected in the main generator module and passed in
+# through the which_conv arg. Similar rules apply with which_bn (the input
+# size [which is actually the number of channels of the conditional info] must 
+# be preselected)
+class GBlock(nn.Module):
+  def __init__(self, in_channels, out_channels,
+               which_conv=nn.Conv2d, which_bn=bn, activation=None, 
+               upsample=None):
+    super(GBlock, self).__init__()
+    
+    self.in_channels, self.out_channels = in_channels, out_channels
+    self.which_conv, self.which_bn = which_conv, which_bn
+    self.activation = activation
+    self.upsample = upsample
+    # Conv layers
+    self.conv1 = self.which_conv(self.in_channels, self.out_channels)
+    self.conv2 = self.which_conv(self.out_channels, self.out_channels)
+    self.learnable_sc = in_channels != out_channels or upsample
+    if self.learnable_sc:
+      self.conv_sc = self.which_conv(in_channels, out_channels, 
+                                     kernel_size=1, padding=0)
+    # Batchnorm layers
+    self.bn1 = self.which_bn(in_channels)
+    self.bn2 = self.which_bn(out_channels)
+    # upsample layers
+    self.upsample = upsample
+
+  def forward(self, x, y):
+    h = self.activation(self.bn1(x, y))
+    if self.upsample:
+      h = self.upsample(h)
+      x = self.upsample(x)
+    h = self.conv1(h)
+    h = self.activation(self.bn2(h, y))
+    h = self.conv2(h)
+    if self.learnable_sc:       
+      x = self.conv_sc(x)
+    return h + x
+    
+    
+# Residual block for the discriminator
+class DBlock(nn.Module):
+  def __init__(self, in_channels, out_channels, which_conv=SNConv2d, wide=True,
+               preactivation=False, activation=None, downsample=None,):
+    super(DBlock, self).__init__()
+    self.in_channels, self.out_channels = in_channels, out_channels
+    # If using wide D (as in SA-GAN and BigGAN), change the channel pattern
+    self.hidden_channels = self.out_channels if wide else self.in_channels
+    self.which_conv = which_conv
+    self.preactivation = preactivation
+    self.activation = activation
+    self.downsample = downsample
+        
+    # Conv layers
+    self.conv1 = self.which_conv(self.in_channels, self.hidden_channels)
+    self.conv2 = self.which_conv(self.hidden_channels, self.out_channels)
+    self.learnable_sc = True if (in_channels != out_channels) or downsample else False
+    if self.learnable_sc:
+      self.conv_sc = self.which_conv(in_channels, out_channels, 
+                                     kernel_size=1, padding=0)
+  def shortcut(self, x):
+    if self.preactivation:
+      if self.learnable_sc:
+        x = self.conv_sc(x)
+      if self.downsample:
+        x = self.downsample(x)
+    else:
+      if self.downsample:
+        x = self.downsample(x)
+      if self.learnable_sc:
+        x = self.conv_sc(x)
+    return x
+    
+  def forward(self, x):
+    if self.preactivation:
+      # h = self.activation(x) # NOT TODAY SATAN
+      # Andy's note: This line *must* be an out-of-place ReLU or it 
+      #              will negatively affect the shortcut connection.
+      h = F.relu(x)
+    else:
+      h = x    
+    h = self.conv1(h)
+    h = self.conv2(self.activation(h))
+    if self.downsample:
+      h = self.downsample(h)     
+        
+    return h + self.shortcut(x)
+    
+# dogball

src/models/big/sync_batchnorm/__init__.py

@@ -0,0 +1,12 @@
+# -*- coding: utf-8 -*-
+# File   : __init__.py
+# Author : Jiayuan Mao
+# Email  : maojiayuan@gmail.com
+# Date   : 27/01/2018
+# 
+# This file is part of Synchronized-BatchNorm-PyTorch.
+# https://github.com/vacancy/Synchronized-BatchNorm-PyTorch
+# Distributed under MIT License.
+
+from .batchnorm import SynchronizedBatchNorm1d, SynchronizedBatchNorm2d, SynchronizedBatchNorm3d
+from .replicate import DataParallelWithCallback, patch_replication_callback

src/models/big/sync_batchnorm/batchnorm.py

@@ -0,0 +1,349 @@
+# -*- coding: utf-8 -*-
+# File   : batchnorm.py
+# Author : Jiayuan Mao
+# Email  : maojiayuan@gmail.com
+# Date   : 27/01/2018
+#
+# This file is part of Synchronized-BatchNorm-PyTorch.
+# https://github.com/vacancy/Synchronized-BatchNorm-PyTorch
+# Distributed under MIT License.
+
+import collections
+
+import torch
+import torch.nn.functional as F
+
+from torch.nn.modules.batchnorm import _BatchNorm
+from torch.nn.parallel._functions import ReduceAddCoalesced, Broadcast
+
+from .comm import SyncMaster
+
+__all__ = ['SynchronizedBatchNorm1d', 'SynchronizedBatchNorm2d', 'SynchronizedBatchNorm3d']
+
+
+def _sum_ft(tensor):
+    """sum over the first and last dimention"""
+    return tensor.sum(dim=0).sum(dim=-1)
+
+
+def _unsqueeze_ft(tensor):
+    """add new dementions at the front and the tail"""
+    return tensor.unsqueeze(0).unsqueeze(-1)
+
+
+_ChildMessage = collections.namedtuple('_ChildMessage', ['sum', 'ssum', 'sum_size'])
+_MasterMessage = collections.namedtuple('_MasterMessage', ['sum', 'inv_std'])
+# _MasterMessage = collections.namedtuple('_MasterMessage', ['sum', 'ssum', 'sum_size'])
+
+class _SynchronizedBatchNorm(_BatchNorm):
+    def __init__(self, num_features, eps=1e-5, momentum=0.1, affine=True):
+        super(_SynchronizedBatchNorm, self).__init__(num_features, eps=eps, momentum=momentum, affine=affine)
+
+        self._sync_master = SyncMaster(self._data_parallel_master)
+
+        self._is_parallel = False
+        self._parallel_id = None
+        self._slave_pipe = None
+
+    def forward(self, input, gain=None, bias=None):
+        # If it is not parallel computation or is in evaluation mode, use PyTorch's implementation.
+        if not (self._is_parallel and self.training):
+            out = F.batch_norm(
+                input, self.running_mean, self.running_var, self.weight, self.bias,
+                self.training, self.momentum, self.eps)
+            if gain is not None:
+              out = out + gain
+            if bias is not None:
+              out = out + bias
+            return out
+
+        # Resize the input to (B, C, -1).
+        input_shape = input.size()
+        # print(input_shape)
+        input = input.view(input.size(0), input.size(1), -1)
+
+        # Compute the sum and square-sum.
+        sum_size = input.size(0) * input.size(2)
+        input_sum = _sum_ft(input)
+        input_ssum = _sum_ft(input ** 2)
+        # Reduce-and-broadcast the statistics.
+        # print('it begins')
+        if self._parallel_id == 0:
+            mean, inv_std = self._sync_master.run_master(_ChildMessage(input_sum, input_ssum, sum_size))
+        else:
+            mean, inv_std = self._slave_pipe.run_slave(_ChildMessage(input_sum, input_ssum, sum_size))
+        # if self._parallel_id == 0:
+            # # print('here')
+            # sum, ssum, num = self._sync_master.run_master(_ChildMessage(input_sum, input_ssum, sum_size))
+        # else:
+            # # print('there')
+            # sum, ssum, num = self._slave_pipe.run_slave(_ChildMessage(input_sum, input_ssum, sum_size))
+        
+        # print('how2')
+        # num = sum_size
+        # print('Sum: %f, ssum: %f, sumsize: %f, insum: %f' %(float(sum.sum().cpu()), float(ssum.sum().cpu()), float(sum_size), float(input_sum.sum().cpu()))) 
+        # Fix the graph
+        # sum = (sum.detach() - input_sum.detach()) + input_sum
+        # ssum = (ssum.detach() - input_ssum.detach()) + input_ssum
+        
+        # mean = sum / num
+        # var = ssum / num - mean ** 2
+        # # var = (ssum - mean * sum) / num
+        # inv_std = torch.rsqrt(var + self.eps)
+        
+        # Compute the output.
+        if gain is not None:
+          # print('gaining')
+          # scale = _unsqueeze_ft(inv_std) * gain.squeeze(-1)
+          # shift = _unsqueeze_ft(mean) * scale - bias.squeeze(-1)
+          # output = input * scale - shift
+          output = (input - _unsqueeze_ft(mean)) * (_unsqueeze_ft(inv_std) * gain.squeeze(-1)) + bias.squeeze(-1)
+        elif self.affine:
+            # MJY:: Fuse the multiplication for speed.
+            output = (input - _unsqueeze_ft(mean)) * _unsqueeze_ft(inv_std * self.weight) + _unsqueeze_ft(self.bias)        
+        else:
+            output = (input - _unsqueeze_ft(mean)) * _unsqueeze_ft(inv_std)
+
+        # Reshape it.
+        return output.view(input_shape)
+
+    def __data_parallel_replicate__(self, ctx, copy_id):
+        self._is_parallel = True
+        self._parallel_id = copy_id
+
+        # parallel_id == 0 means master device.
+        if self._parallel_id == 0:
+            ctx.sync_master = self._sync_master
+        else:
+            self._slave_pipe = ctx.sync_master.register_slave(copy_id)
+
+    def _data_parallel_master(self, intermediates):
+        """Reduce the sum and square-sum, compute the statistics, and broadcast it."""
+
+        # Always using same "device order" makes the ReduceAdd operation faster.
+        # Thanks to:: Tete Xiao (http://tetexiao.com/)
+        intermediates = sorted(intermediates, key=lambda i: i[1].sum.get_device())
+
+        to_reduce = [i[1][:2] for i in intermediates]
+        to_reduce = [j for i in to_reduce for j in i]  # flatten
+        target_gpus = [i[1].sum.get_device() for i in intermediates]
+
+        sum_size = sum([i[1].sum_size for i in intermediates])
+        sum_, ssum = ReduceAddCoalesced.apply(target_gpus[0], 2, *to_reduce)
+        mean, inv_std = self._compute_mean_std(sum_, ssum, sum_size)
+
+        broadcasted = Broadcast.apply(target_gpus, mean, inv_std)
+        # print('a')
+        # print(type(sum_), type(ssum), type(sum_size), sum_.shape, ssum.shape, sum_size)
+        # broadcasted = Broadcast.apply(target_gpus, sum_, ssum, torch.tensor(sum_size).float().to(sum_.device))
+        # print('b')
+        outputs = []
+        for i, rec in enumerate(intermediates):
+            outputs.append((rec[0], _MasterMessage(*broadcasted[i*2:i*2+2])))
+            # outputs.append((rec[0], _MasterMessage(*broadcasted[i*3:i*3+3])))
+
+        return outputs
+
+    def _compute_mean_std(self, sum_, ssum, size):
+        """Compute the mean and standard-deviation with sum and square-sum. This method
+        also maintains the moving average on the master device."""
+        assert size > 1, 'BatchNorm computes unbiased standard-deviation, which requires size > 1.'
+        mean = sum_ / size
+        sumvar = ssum - sum_ * mean
+        unbias_var = sumvar / (size - 1)
+        bias_var = sumvar / size
+
+        self.running_mean = (1 - self.momentum) * self.running_mean + self.momentum * mean.data
+        self.running_var = (1 - self.momentum) * self.running_var + self.momentum * unbias_var.data
+        return mean, torch.rsqrt(bias_var + self.eps)
+        # return mean, bias_var.clamp(self.eps) ** -0.5
+
+
+class SynchronizedBatchNorm1d(_SynchronizedBatchNorm):
+    r"""Applies Synchronized Batch Normalization over a 2d or 3d input that is seen as a
+    mini-batch.
+
+    .. math::
+
+        y = \frac{x - mean[x]}{ \sqrt{Var[x] + \epsilon}} * gamma + beta
+
+    This module differs from the built-in PyTorch BatchNorm1d as the mean and
+    standard-deviation are reduced across all devices during training.
+
+    For example, when one uses `nn.DataParallel` to wrap the network during
+    training, PyTorch's implementation normalize the tensor on each device using
+    the statistics only on that device, which accelerated the computation and
+    is also easy to implement, but the statistics might be inaccurate.
+    Instead, in this synchronized version, the statistics will be computed
+    over all training samples distributed on multiple devices.
+
+    Note that, for one-GPU or CPU-only case, this module behaves exactly same
+    as the built-in PyTorch implementation.
+
+    The mean and standard-deviation are calculated per-dimension over
+    the mini-batches and gamma and beta are learnable parameter vectors
+    of size C (where C is the input size).
+
+    During training, this layer keeps a running estimate of its computed mean
+    and variance. The running sum is kept with a default momentum of 0.1.
+
+    During evaluation, this running mean/variance is used for normalization.
+
+    Because the BatchNorm is done over the `C` dimension, computing statistics
+    on `(N, L)` slices, it's common terminology to call this Temporal BatchNorm
+
+    Args:
+        num_features: num_features from an expected input of size
+            `batch_size x num_features [x width]`
+        eps: a value added to the denominator for numerical stability.
+            Default: 1e-5
+        momentum: the value used for the running_mean and running_var
+            computation. Default: 0.1
+        affine: a boolean value that when set to ``True``, gives the layer learnable
+            affine parameters. Default: ``True``
+
+    Shape:
+        - Input: :math:`(N, C)` or :math:`(N, C, L)`
+        - Output: :math:`(N, C)` or :math:`(N, C, L)` (same shape as input)
+
+    Examples:
+        >>> # With Learnable Parameters
+        >>> m = SynchronizedBatchNorm1d(100)
+        >>> # Without Learnable Parameters
+        >>> m = SynchronizedBatchNorm1d(100, affine=False)
+        >>> input = torch.autograd.Variable(torch.randn(20, 100))
+        >>> output = m(input)
+    """
+
+    def _check_input_dim(self, input):
+        if input.dim() != 2 and input.dim() != 3:
+            raise ValueError('expected 2D or 3D input (got {}D input)'
+                             .format(input.dim()))
+        super(SynchronizedBatchNorm1d, self)._check_input_dim(input)
+
+
+class SynchronizedBatchNorm2d(_SynchronizedBatchNorm):
+    r"""Applies Batch Normalization over a 4d input that is seen as a mini-batch
+    of 3d inputs
+
+    .. math::
+
+        y = \frac{x - mean[x]}{ \sqrt{Var[x] + \epsilon}} * gamma + beta
+
+    This module differs from the built-in PyTorch BatchNorm2d as the mean and
+    standard-deviation are reduced across all devices during training.
+
+    For example, when one uses `nn.DataParallel` to wrap the network during
+    training, PyTorch's implementation normalize the tensor on each device using
+    the statistics only on that device, which accelerated the computation and
+    is also easy to implement, but the statistics might be inaccurate.
+    Instead, in this synchronized version, the statistics will be computed
+    over all training samples distributed on multiple devices.
+
+    Note that, for one-GPU or CPU-only case, this module behaves exactly same
+    as the built-in PyTorch implementation.
+
+    The mean and standard-deviation are calculated per-dimension over
+    the mini-batches and gamma and beta are learnable parameter vectors
+    of size C (where C is the input size).
+
+    During training, this layer keeps a running estimate of its computed mean
+    and variance. The running sum is kept with a default momentum of 0.1.
+
+    During evaluation, this running mean/variance is used for normalization.
+
+    Because the BatchNorm is done over the `C` dimension, computing statistics
+    on `(N, H, W)` slices, it's common terminology to call this Spatial BatchNorm
+
+    Args:
+        num_features: num_features from an expected input of
+            size batch_size x num_features x height x width
+        eps: a value added to the denominator for numerical stability.
+            Default: 1e-5
+        momentum: the value used for the running_mean and running_var
+            computation. Default: 0.1
+        affine: a boolean value that when set to ``True``, gives the layer learnable
+            affine parameters. Default: ``True``
+
+    Shape:
+        - Input: :math:`(N, C, H, W)`
+        - Output: :math:`(N, C, H, W)` (same shape as input)
+
+    Examples:
+        >>> # With Learnable Parameters
+        >>> m = SynchronizedBatchNorm2d(100)
+        >>> # Without Learnable Parameters
+        >>> m = SynchronizedBatchNorm2d(100, affine=False)
+        >>> input = torch.autograd.Variable(torch.randn(20, 100, 35, 45))
+        >>> output = m(input)
+    """
+
+    def _check_input_dim(self, input):
+        if input.dim() != 4:
+            raise ValueError('expected 4D input (got {}D input)'
+                             .format(input.dim()))
+        super(SynchronizedBatchNorm2d, self)._check_input_dim(input)
+
+
+class SynchronizedBatchNorm3d(_SynchronizedBatchNorm):
+    r"""Applies Batch Normalization over a 5d input that is seen as a mini-batch
+    of 4d inputs
+
+    .. math::
+
+        y = \frac{x - mean[x]}{ \sqrt{Var[x] + \epsilon}} * gamma + beta
+
+    This module differs from the built-in PyTorch BatchNorm3d as the mean and
+    standard-deviation are reduced across all devices during training.
+
+    For example, when one uses `nn.DataParallel` to wrap the network during
+    training, PyTorch's implementation normalize the tensor on each device using
+    the statistics only on that device, which accelerated the computation and
+    is also easy to implement, but the statistics might be inaccurate.
+    Instead, in this synchronized version, the statistics will be computed
+    over all training samples distributed on multiple devices.
+
+    Note that, for one-GPU or CPU-only case, this module behaves exactly same
+    as the built-in PyTorch implementation.
+
+    The mean and standard-deviation are calculated per-dimension over
+    the mini-batches and gamma and beta are learnable parameter vectors
+    of size C (where C is the input size).
+
+    During training, this layer keeps a running estimate of its computed mean
+    and variance. The running sum is kept with a default momentum of 0.1.
+
+    During evaluation, this running mean/variance is used for normalization.
+
+    Because the BatchNorm is done over the `C` dimension, computing statistics
+    on `(N, D, H, W)` slices, it's common terminology to call this Volumetric BatchNorm
+    or Spatio-temporal BatchNorm
+
+    Args:
+        num_features: num_features from an expected input of
+            size batch_size x num_features x depth x height x width
+        eps: a value added to the denominator for numerical stability.
+            Default: 1e-5
+        momentum: the value used for the running_mean and running_var
+            computation. Default: 0.1
+        affine: a boolean value that when set to ``True``, gives the layer learnable
+            affine parameters. Default: ``True``
+
+    Shape:
+        - Input: :math:`(N, C, D, H, W)`
+        - Output: :math:`(N, C, D, H, W)` (same shape as input)
+
+    Examples:
+        >>> # With Learnable Parameters
+        >>> m = SynchronizedBatchNorm3d(100)
+        >>> # Without Learnable Parameters
+        >>> m = SynchronizedBatchNorm3d(100, affine=False)
+        >>> input = torch.autograd.Variable(torch.randn(20, 100, 35, 45, 10))
+        >>> output = m(input)
+    """
+
+    def _check_input_dim(self, input):
+        if input.dim() != 5:
+            raise ValueError('expected 5D input (got {}D input)'
+                             .format(input.dim()))
+        super(SynchronizedBatchNorm3d, self)._check_input_dim(input)

src/models/big/sync_batchnorm/batchnorm_reimpl.py

@@ -0,0 +1,74 @@
+#! /usr/bin/env python3
+# -*- coding: utf-8 -*-
+# File   : batchnorm_reimpl.py
+# Author : acgtyrant
+# Date   : 11/01/2018
+#
+# This file is part of Synchronized-BatchNorm-PyTorch.
+# https://github.com/vacancy/Synchronized-BatchNorm-PyTorch
+# Distributed under MIT License.
+
+import torch
+import torch.nn as nn
+import torch.nn.init as init
+
+__all__ = ['BatchNormReimpl']
+
+
+class BatchNorm2dReimpl(nn.Module):
+    """
+    A re-implementation of batch normalization, used for testing the numerical
+    stability.
+
+    Author: acgtyrant
+    See also:
+    https://github.com/vacancy/Synchronized-BatchNorm-PyTorch/issues/14
+    """
+    def __init__(self, num_features, eps=1e-5, momentum=0.1):
+        super().__init__()
+
+        self.num_features = num_features
+        self.eps = eps
+        self.momentum = momentum
+        self.weight = nn.Parameter(torch.empty(num_features))
+        self.bias = nn.Parameter(torch.empty(num_features))
+        self.register_buffer('running_mean', torch.zeros(num_features))
+        self.register_buffer('running_var', torch.ones(num_features))
+        self.reset_parameters()
+
+    def reset_running_stats(self):
+        self.running_mean.zero_()
+        self.running_var.fill_(1)
+
+    def reset_parameters(self):
+        self.reset_running_stats()
+        init.uniform_(self.weight)
+        init.zeros_(self.bias)
+
+    def forward(self, input_):
+        batchsize, channels, height, width = input_.size()
+        numel = batchsize * height * width
+        input_ = input_.permute(1, 0, 2, 3).contiguous().view(channels, numel)
+        sum_ = input_.sum(1)
+        sum_of_square = input_.pow(2).sum(1)
+        mean = sum_ / numel
+        sumvar = sum_of_square - sum_ * mean
+
+        self.running_mean = (
+                (1 - self.momentum) * self.running_mean
+                + self.momentum * mean.detach()
+        )
+        unbias_var = sumvar / (numel - 1)
+        self.running_var = (
+                (1 - self.momentum) * self.running_var
+                + self.momentum * unbias_var.detach()
+        )
+
+        bias_var = sumvar / numel
+        inv_std = 1 / (bias_var + self.eps).pow(0.5)
+        output = (
+                (input_ - mean.unsqueeze(1)) * inv_std.unsqueeze(1) *
+                self.weight.unsqueeze(1) + self.bias.unsqueeze(1))
+
+        return output.view(channels, batchsize, height, width).permute(1, 0, 2, 3).contiguous()
+

src/models/big/sync_batchnorm/comm.py

@@ -0,0 +1,137 @@
+# -*- coding: utf-8 -*-
+# File   : comm.py
+# Author : Jiayuan Mao
+# Email  : maojiayuan@gmail.com
+# Date   : 27/01/2018
+# 
+# This file is part of Synchronized-BatchNorm-PyTorch.
+# https://github.com/vacancy/Synchronized-BatchNorm-PyTorch
+# Distributed under MIT License.
+
+import queue
+import collections
+import threading
+
+__all__ = ['FutureResult', 'SlavePipe', 'SyncMaster']
+
+
+class FutureResult(object):
+    """A thread-safe future implementation. Used only as one-to-one pipe."""
+
+    def __init__(self):
+        self._result = None
+        self._lock = threading.Lock()
+        self._cond = threading.Condition(self._lock)
+
+    def put(self, result):
+        with self._lock:
+            assert self._result is None, 'Previous result has\'t been fetched.'
+            self._result = result
+            self._cond.notify()
+
+    def get(self):
+        with self._lock:
+            if self._result is None:
+                self._cond.wait()
+
+            res = self._result
+            self._result = None
+            return res
+
+
+_MasterRegistry = collections.namedtuple('MasterRegistry', ['result'])
+_SlavePipeBase = collections.namedtuple('_SlavePipeBase', ['identifier', 'queue', 'result'])
+
+
+class SlavePipe(_SlavePipeBase):
+    """Pipe for master-slave communication."""
+
+    def run_slave(self, msg):
+        self.queue.put((self.identifier, msg))
+        ret = self.result.get()
+        self.queue.put(True)
+        return ret
+
+
+class SyncMaster(object):
+    """An abstract `SyncMaster` object.
+
+    - During the replication, as the data parallel will trigger an callback of each module, all slave devices should
+    call `register(id)` and obtain an `SlavePipe` to communicate with the master.
+    - During the forward pass, master device invokes `run_master`, all messages from slave devices will be collected,
+    and passed to a registered callback.
+    - After receiving the messages, the master device should gather the information and determine to message passed
+    back to each slave devices.
+    """
+
+    def __init__(self, master_callback):
+        """
+
+        Args:
+            master_callback: a callback to be invoked after having collected messages from slave devices.
+        """
+        self._master_callback = master_callback
+        self._queue = queue.Queue()
+        self._registry = collections.OrderedDict()
+        self._activated = False
+
+    def __getstate__(self):
+        return {'master_callback': self._master_callback}
+
+    def __setstate__(self, state):
+        self.__init__(state['master_callback'])
+
+    def register_slave(self, identifier):
+        """
+        Register an slave device.
+
+        Args:
+            identifier: an identifier, usually is the device id.
+
+        Returns: a `SlavePipe` object which can be used to communicate with the master device.
+
+        """
+        if self._activated:
+            assert self._queue.empty(), 'Queue is not clean before next initialization.'
+            self._activated = False
+            self._registry.clear()
+        future = FutureResult()
+        self._registry[identifier] = _MasterRegistry(future)
+        return SlavePipe(identifier, self._queue, future)
+
+    def run_master(self, master_msg):
+        """
+        Main entry for the master device in each forward pass.
+        The messages were first collected from each devices (including the master device), and then
+        an callback will be invoked to compute the message to be sent back to each devices
+        (including the master device).
+
+        Args:
+            master_msg: the message that the master want to send to itself. This will be placed as the first
+            message when calling `master_callback`. For detailed usage, see `_SynchronizedBatchNorm` for an example.
+
+        Returns: the message to be sent back to the master device.
+
+        """
+        self._activated = True
+
+        intermediates = [(0, master_msg)]
+        for i in range(self.nr_slaves):
+            intermediates.append(self._queue.get())
+
+        results = self._master_callback(intermediates)
+        assert results[0][0] == 0, 'The first result should belongs to the master.'
+
+        for i, res in results:
+            if i == 0:
+                continue
+            self._registry[i].result.put(res)
+
+        for i in range(self.nr_slaves):
+            assert self._queue.get() is True
+
+        return results[0][1]
+
+    @property
+    def nr_slaves(self):
+        return len(self._registry)

src/models/big/sync_batchnorm/replicate.py

@@ -0,0 +1,94 @@
+# -*- coding: utf-8 -*-
+# File   : replicate.py
+# Author : Jiayuan Mao
+# Email  : maojiayuan@gmail.com
+# Date   : 27/01/2018
+# 
+# This file is part of Synchronized-BatchNorm-PyTorch.
+# https://github.com/vacancy/Synchronized-BatchNorm-PyTorch
+# Distributed under MIT License.
+
+import functools
+
+from torch.nn.parallel.data_parallel import DataParallel
+
+__all__ = [
+    'CallbackContext',
+    'execute_replication_callbacks',
+    'DataParallelWithCallback',
+    'patch_replication_callback'
+]
+
+
+class CallbackContext(object):
+    pass
+
+
+def execute_replication_callbacks(modules):
+    """
+    Execute an replication callback `__data_parallel_replicate__` on each module created by original replication.
+
+    The callback will be invoked with arguments `__data_parallel_replicate__(ctx, copy_id)`
+
+    Note that, as all modules are isomorphism, we assign each sub-module with a context
+    (shared among multiple copies of this module on different devices).
+    Through this context, different copies can share some information.
+
+    We guarantee that the callback on the master copy (the first copy) will be called ahead of calling the callback
+    of any slave copies.
+    """
+    master_copy = modules[0]
+    nr_modules = len(list(master_copy.modules()))
+    ctxs = [CallbackContext() for _ in range(nr_modules)]
+
+    for i, module in enumerate(modules):
+        for j, m in enumerate(module.modules()):
+            if hasattr(m, '__data_parallel_replicate__'):
+                m.__data_parallel_replicate__(ctxs[j], i)
+
+
+class DataParallelWithCallback(DataParallel):
+    """
+    Data Parallel with a replication callback.
+
+    An replication callback `__data_parallel_replicate__` of each module will be invoked after being created by
+    original `replicate` function.
+    The callback will be invoked with arguments `__data_parallel_replicate__(ctx, copy_id)`
+
+    Examples:
+        > sync_bn = SynchronizedBatchNorm1d(10, eps=1e-5, affine=False)
+        > sync_bn = DataParallelWithCallback(sync_bn, device_ids=[0, 1])
+        # sync_bn.__data_parallel_replicate__ will be invoked.
+    """
+
+    def replicate(self, module, device_ids):
+        modules = super(DataParallelWithCallback, self).replicate(module, device_ids)
+        execute_replication_callbacks(modules)
+        return modules
+
+
+def patch_replication_callback(data_parallel):
+    """
+    Monkey-patch an existing `DataParallel` object. Add the replication callback.
+    Useful when you have customized `DataParallel` implementation.
+
+    Examples:
+        > sync_bn = SynchronizedBatchNorm1d(10, eps=1e-5, affine=False)
+        > sync_bn = DataParallel(sync_bn, device_ids=[0, 1])
+        > patch_replication_callback(sync_bn)
+        # this is equivalent to
+        > sync_bn = SynchronizedBatchNorm1d(10, eps=1e-5, affine=False)
+        > sync_bn = DataParallelWithCallback(sync_bn, device_ids=[0, 1])
+    """
+
+    assert isinstance(data_parallel, DataParallel)
+
+    old_replicate = data_parallel.replicate
+
+    @functools.wraps(old_replicate)
+    def new_replicate(module, device_ids):
+        modules = old_replicate(module, device_ids)
+        execute_replication_callbacks(modules)
+        return modules
+
+    data_parallel.replicate = new_replicate

src/models/big/sync_batchnorm/unittest.py

@@ -0,0 +1,29 @@
+# -*- coding: utf-8 -*-
+# File   : unittest.py
+# Author : Jiayuan Mao
+# Email  : maojiayuan@gmail.com
+# Date   : 27/01/2018
+#
+# This file is part of Synchronized-BatchNorm-PyTorch.
+# https://github.com/vacancy/Synchronized-BatchNorm-PyTorch
+# Distributed under MIT License.
+
+import unittest
+import torch
+
+
+class TorchTestCase(unittest.TestCase):
+    def assertTensorClose(self, x, y):
+        adiff = float((x - y).abs().max())
+        if (y == 0).all():
+            rdiff = 'NaN'
+        else:
+            rdiff = float((adiff / y).abs().max())
+
+        message = (
+            'Tensor close check failed\n'
+            'adiff={}\n'
+            'rdiff={}\n'
+        ).format(adiff, rdiff)
+        self.assertTrue(torch.allclose(x, y), message)
+

src/models/big/utils.py

@@ -0,0 +1,1193 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+
+''' Utilities file
+This file contains utility functions for bookkeeping, logging, and data loading.
+Methods which directly affect training should either go in layers, the model,
+or train_fns.py.
+'''
+
+from __future__ import print_function
+import sys
+import os
+import numpy as np
+import time
+import datetime
+import json
+import pickle
+from argparse import ArgumentParser
+import animal_hash
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torchvision
+import torchvision.transforms as transforms
+from torch.utils.data import DataLoader
+
+import datasets as dset
+
+def prepare_parser():
+  usage = 'Parser for all scripts.'
+  parser = ArgumentParser(description=usage)
+  
+  ### Dataset/Dataloader stuff ###
+  parser.add_argument(
+    '--dataset', type=str, default='I128_hdf5',
+    help='Which Dataset to train on, out of I128, I256, C10, C100;'
+         'Append "_hdf5" to use the hdf5 version for ISLVRC '
+         '(default: %(default)s)')
+  parser.add_argument(
+    '--augment', action='store_true', default=False,
+    help='Augment with random crops and flips (default: %(default)s)')
+  parser.add_argument(
+    '--num_workers', type=int, default=8,
+    help='Number of dataloader workers; consider using less for HDF5 '
+         '(default: %(default)s)')
+  parser.add_argument(
+    '--no_pin_memory', action='store_false', dest='pin_memory', default=True,
+    help='Pin data into memory through dataloader? (default: %(default)s)') 
+  parser.add_argument(
+    '--shuffle', action='store_true', default=False,
+    help='Shuffle the data (strongly recommended)? (default: %(default)s)')
+  parser.add_argument(
+    '--load_in_mem', action='store_true', default=False,
+    help='Load all data into memory? (default: %(default)s)')
+  parser.add_argument(
+    '--use_multiepoch_sampler', action='store_true', default=False,
+    help='Use the multi-epoch sampler for dataloader? (default: %(default)s)')
+  
+  
+  ### Model stuff ###
+  parser.add_argument(
+    '--model', type=str, default='BigGAN',
+    help='Name of the model module (default: %(default)s)')
+  parser.add_argument(
+    '--G_param', type=str, default='SN',
+    help='Parameterization style to use for G, spectral norm (SN) or SVD (SVD)'
+          ' or None (default: %(default)s)')
+  parser.add_argument(
+    '--D_param', type=str, default='SN',
+    help='Parameterization style to use for D, spectral norm (SN) or SVD (SVD)'
+         ' or None (default: %(default)s)')    
+  parser.add_argument(
+    '--G_ch', type=int, default=64,
+    help='Channel multiplier for G (default: %(default)s)')
+  parser.add_argument(
+    '--D_ch', type=int, default=64,
+    help='Channel multiplier for D (default: %(default)s)')
+  parser.add_argument(
+    '--G_depth', type=int, default=1,
+    help='Number of resblocks per stage in G? (default: %(default)s)')
+  parser.add_argument(
+    '--D_depth', type=int, default=1,
+    help='Number of resblocks per stage in D? (default: %(default)s)')
+  parser.add_argument(
+    '--D_thin', action='store_false', dest='D_wide', default=True,
+    help='Use the SN-GAN channel pattern for D? (default: %(default)s)')
+  parser.add_argument(
+    '--G_shared', action='store_true', default=False,
+    help='Use shared embeddings in G? (default: %(default)s)')
+  parser.add_argument(
+    '--shared_dim', type=int, default=0,
+    help='G''s shared embedding dimensionality; if 0, will be equal to dim_z. '
+         '(default: %(default)s)')
+  parser.add_argument(
+    '--dim_z', type=int, default=128,
+    help='Noise dimensionality: %(default)s)')
+  parser.add_argument(
+    '--z_var', type=float, default=1.0,
+    help='Noise variance: %(default)s)')    
+  parser.add_argument(
+    '--hier', action='store_true', default=False,
+    help='Use hierarchical z in G? (default: %(default)s)')
+  parser.add_argument(
+    '--cross_replica', action='store_true', default=False,
+    help='Cross_replica batchnorm in G?(default: %(default)s)')
+  parser.add_argument(
+    '--mybn', action='store_true', default=False,
+    help='Use my batchnorm (which supports standing stats?) %(default)s)')
+  parser.add_argument(
+    '--G_nl', type=str, default='relu',
+    help='Activation function for G (default: %(default)s)')
+  parser.add_argument(
+    '--D_nl', type=str, default='relu',
+    help='Activation function for D (default: %(default)s)')
+  parser.add_argument(
+    '--G_attn', type=str, default='64',
+    help='What resolutions to use attention on for G (underscore separated) '
+         '(default: %(default)s)')
+  parser.add_argument(
+    '--D_attn', type=str, default='64',
+    help='What resolutions to use attention on for D (underscore separated) '
+         '(default: %(default)s)')
+  parser.add_argument(
+    '--norm_style', type=str, default='bn',
+    help='Normalizer style for G, one of bn [batchnorm], in [instancenorm], '
+         'ln [layernorm], gn [groupnorm] (default: %(default)s)')
+         
+  ### Model init stuff ###
+  parser.add_argument(
+    '--seed', type=int, default=0,
+    help='Random seed to use; affects both initialization and '
+         ' dataloading. (default: %(default)s)')
+  parser.add_argument(
+    '--G_init', type=str, default='ortho',
+    help='Init style to use for G (default: %(default)s)')
+  parser.add_argument(
+    '--D_init', type=str, default='ortho',
+    help='Init style to use for D(default: %(default)s)')
+  parser.add_argument(
+    '--skip_init', action='store_true', default=False,
+    help='Skip initialization, ideal for testing when ortho init was used '
+          '(default: %(default)s)')
+  
+  ### Optimizer stuff ###
+  parser.add_argument(
+    '--G_lr', type=float, default=5e-5,
+    help='Learning rate to use for Generator (default: %(default)s)')
+  parser.add_argument(
+    '--D_lr', type=float, default=2e-4,
+    help='Learning rate to use for Discriminator (default: %(default)s)')
+  parser.add_argument(
+    '--G_B1', type=float, default=0.0,
+    help='Beta1 to use for Generator (default: %(default)s)')
+  parser.add_argument(
+    '--D_B1', type=float, default=0.0,
+    help='Beta1 to use for Discriminator (default: %(default)s)')
+  parser.add_argument(
+    '--G_B2', type=float, default=0.999,
+    help='Beta2 to use for Generator (default: %(default)s)')
+  parser.add_argument(
+    '--D_B2', type=float, default=0.999,
+    help='Beta2 to use for Discriminator (default: %(default)s)')
+    
+  ### Batch size, parallel, and precision stuff ###
+  parser.add_argument(
+    '--batch_size', type=int, default=64,
+    help='Default overall batchsize (default: %(default)s)')
+  parser.add_argument(
+    '--G_batch_size', type=int, default=0,
+    help='Batch size to use for G; if 0, same as D (default: %(default)s)')
+  parser.add_argument(
+    '--num_G_accumulations', type=int, default=1,
+    help='Number of passes to accumulate G''s gradients over '
+         '(default: %(default)s)')  
+  parser.add_argument(
+    '--num_D_steps', type=int, default=2,
+    help='Number of D steps per G step (default: %(default)s)')
+  parser.add_argument(
+    '--num_D_accumulations', type=int, default=1,
+    help='Number of passes to accumulate D''s gradients over '
+         '(default: %(default)s)')
+  parser.add_argument(
+    '--split_D', action='store_true', default=False,
+    help='Run D twice rather than concatenating inputs? (default: %(default)s)')
+  parser.add_argument(
+    '--num_epochs', type=int, default=100,
+    help='Number of epochs to train for (default: %(default)s)')
+  parser.add_argument(
+    '--parallel', action='store_true', default=False,
+    help='Train with multiple GPUs (default: %(default)s)')
+  parser.add_argument(
+    '--G_fp16', action='store_true', default=False,
+    help='Train with half-precision in G? (default: %(default)s)')
+  parser.add_argument(
+    '--D_fp16', action='store_true', default=False,
+    help='Train with half-precision in D? (default: %(default)s)')
+  parser.add_argument(
+    '--D_mixed_precision', action='store_true', default=False,
+    help='Train with half-precision activations but fp32 params in D? '
+         '(default: %(default)s)')
+  parser.add_argument(
+    '--G_mixed_precision', action='store_true', default=False,
+    help='Train with half-precision activations but fp32 params in G? '
+         '(default: %(default)s)')
+  parser.add_argument(
+    '--accumulate_stats', action='store_true', default=False,
+    help='Accumulate "standing" batchnorm stats? (default: %(default)s)')
+  parser.add_argument(
+    '--num_standing_accumulations', type=int, default=16,
+    help='Number of forward passes to use in accumulating standing stats? '
+         '(default: %(default)s)')        
+    
+  ### Bookkeping stuff ###  
+  parser.add_argument(
+    '--G_eval_mode', action='store_true', default=False,
+    help='Run G in eval mode (running/standing stats?) at sample/test time? '
+         '(default: %(default)s)')
+  parser.add_argument(
+    '--save_every', type=int, default=2000,
+    help='Save every X iterations (default: %(default)s)')
+  parser.add_argument(
+    '--num_save_copies', type=int, default=2,
+    help='How many copies to save (default: %(default)s)')
+  parser.add_argument(
+    '--num_best_copies', type=int, default=2,
+    help='How many previous best checkpoints to save (default: %(default)s)')
+  parser.add_argument(
+    '--which_best', type=str, default='IS',
+    help='Which metric to use to determine when to save new "best"'
+         'checkpoints, one of IS or FID (default: %(default)s)')
+  parser.add_argument(
+    '--no_fid', action='store_true', default=False,
+    help='Calculate IS only, not FID? (default: %(default)s)')
+  parser.add_argument(
+    '--test_every', type=int, default=5000,
+    help='Test every X iterations (default: %(default)s)')
+  parser.add_argument(
+    '--num_inception_images', type=int, default=50000,
+    help='Number of samples to compute inception metrics with '
+         '(default: %(default)s)')
+  parser.add_argument(
+    '--hashname', action='store_true', default=False,
+    help='Use a hash of the experiment name instead of the full config '
+         '(default: %(default)s)') 
+  parser.add_argument(
+    '--base_root', type=str, default='',
+    help='Default location to store all weights, samples, data, and logs '
+           ' (default: %(default)s)')
+  parser.add_argument(
+    '--data_root', type=str, default='data',
+    help='Default location where data is stored (default: %(default)s)')
+  parser.add_argument(
+    '--weights_root', type=str, default='weights',
+    help='Default location to store weights (default: %(default)s)')
+  parser.add_argument(
+    '--logs_root', type=str, default='logs',
+    help='Default location to store logs (default: %(default)s)')
+  parser.add_argument(
+    '--samples_root', type=str, default='samples',
+    help='Default location to store samples (default: %(default)s)')  
+  parser.add_argument(
+    '--pbar', type=str, default='mine',
+    help='Type of progressbar to use; one of "mine" or "tqdm" '
+         '(default: %(default)s)')
+  parser.add_argument(
+    '--name_suffix', type=str, default='',
+    help='Suffix for experiment name for loading weights for sampling '
+         '(consider "best0") (default: %(default)s)')
+  parser.add_argument(
+    '--experiment_name', type=str, default='',
+    help='Optionally override the automatic experiment naming with this arg. '
+         '(default: %(default)s)')
+  parser.add_argument(
+    '--config_from_name', action='store_true', default=False,
+    help='Use a hash of the experiment name instead of the full config '
+         '(default: %(default)s)')
+         
+  ### EMA Stuff ###
+  parser.add_argument(
+    '--ema', action='store_true', default=False,
+    help='Keep an ema of G''s weights? (default: %(default)s)')
+  parser.add_argument(
+    '--ema_decay', type=float, default=0.9999,
+    help='EMA decay rate (default: %(default)s)')
+  parser.add_argument(
+    '--use_ema', action='store_true', default=False,
+    help='Use the EMA parameters of G for evaluation? (default: %(default)s)')
+  parser.add_argument(
+    '--ema_start', type=int, default=0,
+    help='When to start updating the EMA weights (default: %(default)s)')
+  
+  ### Numerical precision and SV stuff ### 
+  parser.add_argument(
+    '--adam_eps', type=float, default=1e-8,
+    help='epsilon value to use for Adam (default: %(default)s)')
+  parser.add_argument(
+    '--BN_eps', type=float, default=1e-5,
+    help='epsilon value to use for BatchNorm (default: %(default)s)')
+  parser.add_argument(
+    '--SN_eps', type=float, default=1e-8,
+    help='epsilon value to use for Spectral Norm(default: %(default)s)')
+  parser.add_argument(
+    '--num_G_SVs', type=int, default=1,
+    help='Number of SVs to track in G (default: %(default)s)')
+  parser.add_argument(
+    '--num_D_SVs', type=int, default=1,
+    help='Number of SVs to track in D (default: %(default)s)')
+  parser.add_argument(
+    '--num_G_SV_itrs', type=int, default=1,
+    help='Number of SV itrs in G (default: %(default)s)')
+  parser.add_argument(
+    '--num_D_SV_itrs', type=int, default=1,
+    help='Number of SV itrs in D (default: %(default)s)')
+  
+  ### Ortho reg stuff ### 
+  parser.add_argument(
+    '--G_ortho', type=float, default=0.0, # 1e-4 is default for BigGAN
+    help='Modified ortho reg coefficient in G(default: %(default)s)')
+  parser.add_argument(
+    '--D_ortho', type=float, default=0.0,
+    help='Modified ortho reg coefficient in D (default: %(default)s)')
+  parser.add_argument(
+    '--toggle_grads', action='store_true', default=True,
+    help='Toggle D and G''s "requires_grad" settings when not training them? '
+         ' (default: %(default)s)')
+  
+  ### Which train function ###
+  parser.add_argument(
+    '--which_train_fn', type=str, default='GAN',
+    help='How2trainyourbois (default: %(default)s)')  
+  
+  ### Resume training stuff
+  parser.add_argument(
+    '--load_weights', type=str, default='',
+    help='Suffix for which weights to load (e.g. best0, copy0) '
+         '(default: %(default)s)')
+  parser.add_argument(
+    '--resume', action='store_true', default=False,
+    help='Resume training? (default: %(default)s)')
+  
+  ### Log stuff ###
+  parser.add_argument(
+    '--logstyle', type=str, default='%3.3e',
+    help='What style to use when logging training metrics?'
+         'One of: %#.#f/ %#.#e (float/exp, text),'
+         'pickle (python pickle),'
+         'npz (numpy zip),'
+         'mat (MATLAB .mat file) (default: %(default)s)')
+  parser.add_argument(
+    '--log_G_spectra', action='store_true', default=False,
+    help='Log the top 3 singular values in each SN layer in G? '
+         '(default: %(default)s)')
+  parser.add_argument(
+    '--log_D_spectra', action='store_true', default=False,
+    help='Log the top 3 singular values in each SN layer in D? '
+         '(default: %(default)s)')
+  parser.add_argument(
+    '--sv_log_interval', type=int, default=10,
+    help='Iteration interval for logging singular values '
+         ' (default: %(default)s)') 
+   
+  return parser
+
+# Arguments for sample.py; not presently used in train.py
+def add_sample_parser(parser):
+  parser.add_argument(
+    '--sample_npz', action='store_true', default=False,
+    help='Sample "sample_num_npz" images and save to npz? '
+         '(default: %(default)s)')
+  parser.add_argument(
+    '--sample_num_npz', type=int, default=50000,
+    help='Number of images to sample when sampling NPZs '
+         '(default: %(default)s)')
+  parser.add_argument(
+    '--sample_sheets', action='store_true', default=False,
+    help='Produce class-conditional sample sheets and stick them in '
+         'the samples root? (default: %(default)s)')
+  parser.add_argument(
+    '--sample_interps', action='store_true', default=False,
+    help='Produce interpolation sheets and stick them in '
+         'the samples root? (default: %(default)s)')         
+  parser.add_argument(
+    '--sample_sheet_folder_num', type=int, default=-1,
+    help='Number to use for the folder for these sample sheets '
+         '(default: %(default)s)')
+  parser.add_argument(
+    '--sample_random', action='store_true', default=False,
+    help='Produce a single random sheet? (default: %(default)s)')
+  parser.add_argument(
+    '--sample_trunc_curves', type=str, default='',
+    help='Get inception metrics with a range of variances?'
+         'To use this, specify a startpoint, step, and endpoint, e.g. '
+         '--sample_trunc_curves 0.2_0.1_1.0 for a startpoint of 0.2, '
+         'endpoint of 1.0, and stepsize of 1.0.  Note that this is '
+         'not exactly identical to using tf.truncated_normal, but should '
+         'have approximately the same effect. (default: %(default)s)')
+  parser.add_argument(
+    '--sample_inception_metrics', action='store_true', default=False,
+    help='Calculate Inception metrics with sample.py? (default: %(default)s)')  
+  return parser
+
+# Convenience dicts
+dset_dict = {'I32': dset.ImageFolder, 'I64': dset.ImageFolder, 
+             'I128': dset.ImageFolder, 'I256': dset.ImageFolder,
+             'I32_hdf5': dset.ILSVRC_HDF5, 'I64_hdf5': dset.ILSVRC_HDF5, 
+             'I128_hdf5': dset.ILSVRC_HDF5, 'I256_hdf5': dset.ILSVRC_HDF5,
+             'C10': dset.CIFAR10, 'C100': dset.CIFAR100}
+imsize_dict = {'I32': 32, 'I32_hdf5': 32,
+               'I64': 64, 'I64_hdf5': 64,
+               'I128': 128, 'I128_hdf5': 128,
+               'I256': 256, 'I256_hdf5': 256,
+               'C10': 32, 'C100': 32}
+root_dict = {'I32': 'ImageNet', 'I32_hdf5': 'ILSVRC32.hdf5',
+             'I64': 'ImageNet', 'I64_hdf5': 'ILSVRC64.hdf5',
+             'I128': 'ImageNet', 'I128_hdf5': 'ILSVRC128.hdf5',
+             'I256': 'ImageNet', 'I256_hdf5': 'ILSVRC256.hdf5',
+             'C10': 'cifar', 'C100': 'cifar'}
+nclass_dict = {'I32': 1000, 'I32_hdf5': 1000,
+               'I64': 1000, 'I64_hdf5': 1000,
+               'I128': 1000, 'I128_hdf5': 1000,
+               'I256': 1000, 'I256_hdf5': 1000,
+               'C10': 10, 'C100': 100}
+# Number of classes to put per sample sheet               
+classes_per_sheet_dict = {'I32': 50, 'I32_hdf5': 50,
+                          'I64': 50, 'I64_hdf5': 50,
+                          'I128': 20, 'I128_hdf5': 20,
+                          'I256': 20, 'I256_hdf5': 20,
+                          'C10': 10, 'C100': 100}
+activation_dict = {'inplace_relu': nn.ReLU(inplace=True),
+                   'relu': nn.ReLU(inplace=False),
+                   'ir': nn.ReLU(inplace=True),}
+
+class CenterCropLongEdge(object):
+  """Crops the given PIL Image on the long edge.
+  Args:
+      size (sequence or int): Desired output size of the crop. If size is an
+          int instead of sequence like (h, w), a square crop (size, size) is
+          made.
+  """
+  def __call__(self, img):
+    """
+    Args:
+        img (PIL Image): Image to be cropped.
+    Returns:
+        PIL Image: Cropped image.
+    """
+    return transforms.functional.center_crop(img, min(img.size))
+
+  def __repr__(self):
+    return self.__class__.__name__
+
+class RandomCropLongEdge(object):
+  """Crops the given PIL Image on the long edge with a random start point.
+  Args:
+      size (sequence or int): Desired output size of the crop. If size is an
+          int instead of sequence like (h, w), a square crop (size, size) is
+          made.
+  """
+  def __call__(self, img):
+    """
+    Args:
+        img (PIL Image): Image to be cropped.
+    Returns:
+        PIL Image: Cropped image.
+    """
+    size = (min(img.size), min(img.size))
+    # Only step forward along this edge if it's the long edge
+    i = (0 if size[0] == img.size[0] 
+          else np.random.randint(low=0,high=img.size[0] - size[0]))
+    j = (0 if size[1] == img.size[1]
+          else np.random.randint(low=0,high=img.size[1] - size[1]))
+    return transforms.functional.crop(img, i, j, size[0], size[1])
+
+  def __repr__(self):
+    return self.__class__.__name__
+
+    
+# multi-epoch Dataset sampler to avoid memory leakage and enable resumption of
+# training from the same sample regardless of if we stop mid-epoch
+class MultiEpochSampler(torch.utils.data.Sampler):
+  r"""Samples elements randomly over multiple epochs
+
+  Arguments:
+      data_source (Dataset): dataset to sample from
+      num_epochs (int) : Number of times to loop over the dataset
+      start_itr (int) : which iteration to begin from
+  """
+
+  def __init__(self, data_source, num_epochs, start_itr=0, batch_size=128):
+    self.data_source = data_source
+    self.num_samples = len(self.data_source)
+    self.num_epochs = num_epochs
+    self.start_itr = start_itr
+    self.batch_size = batch_size
+
+    if not isinstance(self.num_samples, int) or self.num_samples <= 0:
+      raise ValueError("num_samples should be a positive integeral "
+                       "value, but got num_samples={}".format(self.num_samples))
+
+  def __iter__(self):
+    n = len(self.data_source)
+    # Determine number of epochs
+    num_epochs = int(np.ceil((n * self.num_epochs 
+                              - (self.start_itr * self.batch_size)) / float(n)))
+    # Sample all the indices, and then grab the last num_epochs index sets;
+    # This ensures if we're starting at epoch 4, we're still grabbing epoch 4's
+    # indices
+    out = [torch.randperm(n) for epoch in range(self.num_epochs)][-num_epochs:]
+    # Ignore the first start_itr % n indices of the first epoch
+    out[0] = out[0][(self.start_itr * self.batch_size % n):]
+    # if self.replacement:
+      # return iter(torch.randint(high=n, size=(self.num_samples,), dtype=torch.int64).tolist())
+    # return iter(.tolist())
+    output = torch.cat(out).tolist()
+    print('Length dataset output is %d' % len(output))
+    return iter(output)
+
+  def __len__(self):
+    return len(self.data_source) * self.num_epochs - self.start_itr * self.batch_size
+
+
+# Convenience function to centralize all data loaders
+def get_data_loaders(dataset, data_root=None, augment=False, batch_size=64, 
+                     num_workers=8, shuffle=True, load_in_mem=False, hdf5=False,
+                     pin_memory=True, drop_last=True, start_itr=0,
+                     num_epochs=500, use_multiepoch_sampler=False,
+                     **kwargs):
+
+  # Append /FILENAME.hdf5 to root if using hdf5
+  data_root += '/%s' % root_dict[dataset]
+  print('Using dataset root location %s' % data_root)
+
+  which_dataset = dset_dict[dataset]
+  norm_mean = [0.5,0.5,0.5]
+  norm_std = [0.5,0.5,0.5]
+  image_size = imsize_dict[dataset]
+  # For image folder datasets, name of the file where we store the precomputed
+  # image locations to avoid having to walk the dirs every time we load.
+  dataset_kwargs = {'index_filename': '%s_imgs.npz' % dataset}
+  
+  # HDF5 datasets have their own inbuilt transform, no need to train_transform  
+  if 'hdf5' in dataset:
+    train_transform = None
+  else:
+    if augment:
+      print('Data will be augmented...')
+      if dataset in ['C10', 'C100']:
+        train_transform = [transforms.RandomCrop(32, padding=4),
+                           transforms.RandomHorizontalFlip()]
+      else:
+        train_transform = [RandomCropLongEdge(),
+                         transforms.Resize(image_size),
+                         transforms.RandomHorizontalFlip()]
+    else:
+      print('Data will not be augmented...')
+      if dataset in ['C10', 'C100']:
+        train_transform = []
+      else:
+        train_transform = [CenterCropLongEdge(), transforms.Resize(image_size)]
+      # train_transform = [transforms.Resize(image_size), transforms.CenterCrop]
+    train_transform = transforms.Compose(train_transform + [
+                     transforms.ToTensor(),
+                     transforms.Normalize(norm_mean, norm_std)])
+  train_set = which_dataset(root=data_root, transform=train_transform,
+                            load_in_mem=load_in_mem, **dataset_kwargs)
+
+  # Prepare loader; the loaders list is for forward compatibility with
+  # using validation / test splits.
+  loaders = []   
+  if use_multiepoch_sampler:
+    print('Using multiepoch sampler from start_itr %d...' % start_itr)
+    loader_kwargs = {'num_workers': num_workers, 'pin_memory': pin_memory}
+    sampler = MultiEpochSampler(train_set, num_epochs, start_itr, batch_size)
+    train_loader = DataLoader(train_set, batch_size=batch_size,
+                              sampler=sampler, **loader_kwargs)
+  else:
+    loader_kwargs = {'num_workers': num_workers, 'pin_memory': pin_memory,
+                     'drop_last': drop_last} # Default, drop last incomplete batch
+    train_loader = DataLoader(train_set, batch_size=batch_size,
+                              shuffle=shuffle, **loader_kwargs)
+  loaders.append(train_loader)
+  return loaders
+
+
+# Utility file to seed rngs
+def seed_rng(seed):
+  torch.manual_seed(seed)
+  torch.cuda.manual_seed(seed)
+  np.random.seed(seed)
+
+
+# Utility to peg all roots to a base root
+# If a base root folder is provided, peg all other root folders to it.
+def update_config_roots(config):
+  if config['base_root']:
+    print('Pegging all root folders to base root %s' % config['base_root'])
+    for key in ['data', 'weights', 'logs', 'samples']:
+      config['%s_root' % key] = '%s/%s' % (config['base_root'], key)
+  return config
+
+
+# Utility to prepare root folders if they don't exist; parent folder must exist
+def prepare_root(config):
+  for key in ['weights_root', 'logs_root', 'samples_root']:
+    if not os.path.exists(config[key]):
+      print('Making directory %s for %s...' % (config[key], key))
+      os.mkdir(config[key])
+
+
+# Simple wrapper that applies EMA to a model. COuld be better done in 1.0 using
+# the parameters() and buffers() module functions, but for now this works
+# with state_dicts using .copy_
+class ema(object):
+  def __init__(self, source, target, decay=0.9999, start_itr=0):
+    self.source = source
+    self.target = target
+    self.decay = decay
+    # Optional parameter indicating what iteration to start the decay at
+    self.start_itr = start_itr
+    # Initialize target's params to be source's
+    self.source_dict = self.source.state_dict()
+    self.target_dict = self.target.state_dict()
+    print('Initializing EMA parameters to be source parameters...')
+    with torch.no_grad():
+      for key in self.source_dict:
+        self.target_dict[key].data.copy_(self.source_dict[key].data)
+        # target_dict[key].data = source_dict[key].data # Doesn't work!
+
+  def update(self, itr=None):
+    # If an iteration counter is provided and itr is less than the start itr,
+    # peg the ema weights to the underlying weights.
+    if itr and itr < self.start_itr:
+      decay = 0.0
+    else:
+      decay = self.decay
+    with torch.no_grad():
+      for key in self.source_dict:
+        self.target_dict[key].data.copy_(self.target_dict[key].data * decay 
+                                     + self.source_dict[key].data * (1 - decay))
+
+
+# Apply modified ortho reg to a model
+# This function is an optimized version that directly computes the gradient,
+# instead of computing and then differentiating the loss.
+def ortho(model, strength=1e-4, blacklist=[]):
+  with torch.no_grad():
+    for param in model.parameters():
+      # Only apply this to parameters with at least 2 axes, and not in the blacklist
+      if len(param.shape) < 2 or any([param is item for item in blacklist]):
+        continue
+      w = param.view(param.shape[0], -1)
+      grad = (2 * torch.mm(torch.mm(w, w.t()) 
+              * (1. - torch.eye(w.shape[0], device=w.device)), w))
+      param.grad.data += strength * grad.view(param.shape)
+
+
+# Default ortho reg
+# This function is an optimized version that directly computes the gradient,
+# instead of computing and then differentiating the loss.
+def default_ortho(model, strength=1e-4, blacklist=[]):
+  with torch.no_grad():
+    for param in model.parameters():
+      # Only apply this to parameters with at least 2 axes & not in blacklist
+      if len(param.shape) < 2 or param in blacklist:
+        continue
+      w = param.view(param.shape[0], -1)
+      grad = (2 * torch.mm(torch.mm(w, w.t()) 
+               - torch.eye(w.shape[0], device=w.device), w))
+      param.grad.data += strength * grad.view(param.shape)
+
+
+# Convenience utility to switch off requires_grad
+def toggle_grad(model, on_or_off):
+  for param in model.parameters():
+    param.requires_grad = on_or_off
+
+
+# Function to join strings or ignore them
+# Base string is the string to link "strings," while strings
+# is a list of strings or Nones.
+def join_strings(base_string, strings):
+  return base_string.join([item for item in strings if item])
+
+
+# Save a model's weights, optimizer, and the state_dict
+def save_weights(G, D, state_dict, weights_root, experiment_name, 
+                 name_suffix=None, G_ema=None):
+  root = '/'.join([weights_root, experiment_name])
+  if not os.path.exists(root):
+    os.mkdir(root)
+  if name_suffix:
+    print('Saving weights to %s/%s...' % (root, name_suffix))
+  else:
+    print('Saving weights to %s...' % root)
+  torch.save(G.state_dict(), 
+              '%s/%s.pth' % (root, join_strings('_', ['G', name_suffix])))
+  torch.save(G.optim.state_dict(), 
+              '%s/%s.pth' % (root, join_strings('_', ['G_optim', name_suffix])))
+  torch.save(D.state_dict(), 
+              '%s/%s.pth' % (root, join_strings('_', ['D', name_suffix])))
+  torch.save(D.optim.state_dict(),
+              '%s/%s.pth' % (root, join_strings('_', ['D_optim', name_suffix])))
+  torch.save(state_dict,
+              '%s/%s.pth' % (root, join_strings('_', ['state_dict', name_suffix])))
+  if G_ema is not None:
+    torch.save(G_ema.state_dict(), 
+                '%s/%s.pth' % (root, join_strings('_', ['G_ema', name_suffix])))
+
+
+# Load a model's weights, optimizer, and the state_dict
+def load_weights(G, D, state_dict, weights_root, experiment_name, 
+                 name_suffix=None, G_ema=None, strict=True, load_optim=True):
+  root = '/'.join([weights_root, experiment_name])
+  if name_suffix:
+    print('Loading %s weights from %s...' % (name_suffix, root))
+  else:
+    print('Loading weights from %s...' % root)
+  if G is not None:
+    G.load_state_dict(
+      torch.load('%s/%s.pth' % (root, join_strings('_', ['G', name_suffix]))),
+      strict=strict)
+    if load_optim:
+      G.optim.load_state_dict(
+        torch.load('%s/%s.pth' % (root, join_strings('_', ['G_optim', name_suffix]))))
+  if D is not None:
+    D.load_state_dict(
+      torch.load('%s/%s.pth' % (root, join_strings('_', ['D', name_suffix]))),
+      strict=strict)
+    if load_optim:
+      D.optim.load_state_dict(
+        torch.load('%s/%s.pth' % (root, join_strings('_', ['D_optim', name_suffix]))))
+  # Load state dict
+  for item in state_dict:
+    state_dict[item] = torch.load('%s/%s.pth' % (root, join_strings('_', ['state_dict', name_suffix])))[item]
+  if G_ema is not None:
+    G_ema.load_state_dict(
+      torch.load('%s/%s.pth' % (root, join_strings('_', ['G_ema', name_suffix]))),
+      strict=strict)
+
+
+''' MetricsLogger originally stolen from VoxNet source code.
+    Used for logging inception metrics'''
+class MetricsLogger(object):
+  def __init__(self, fname, reinitialize=False):
+    self.fname = fname
+    self.reinitialize = reinitialize
+    if os.path.exists(self.fname):
+      if self.reinitialize:
+        print('{} exists, deleting...'.format(self.fname))
+        os.remove(self.fname)
+
+  def log(self, record=None, **kwargs):
+    """
+    Assumption: no newlines in the input.
+    """
+    if record is None:
+      record = {}
+    record.update(kwargs)
+    record['_stamp'] = time.time()
+    with open(self.fname, 'a') as f:
+      f.write(json.dumps(record, ensure_ascii=True) + '\n')
+
+
+# Logstyle is either:
+# '%#.#f' for floating point representation in text
+# '%#.#e' for exponent representation in text
+# 'npz' for output to npz # NOT YET SUPPORTED
+# 'pickle' for output to a python pickle # NOT YET SUPPORTED
+# 'mat' for output to a MATLAB .mat file # NOT YET SUPPORTED
+class MyLogger(object):
+  def __init__(self, fname, reinitialize=False, logstyle='%3.3f'):
+    self.root = fname
+    if not os.path.exists(self.root):
+      os.mkdir(self.root)
+    self.reinitialize = reinitialize
+    self.metrics = []
+    self.logstyle = logstyle # One of '%3.3f' or like '%3.3e'
+
+  # Delete log if re-starting and log already exists
+  def reinit(self, item):
+    if os.path.exists('%s/%s.log' % (self.root, item)):
+      if self.reinitialize:
+        # Only print the removal mess
+        if 'sv' in item :
+          if not any('sv' in item for item in self.metrics):
+            print('Deleting singular value logs...')
+        else:
+          print('{} exists, deleting...'.format('%s_%s.log' % (self.root, item)))
+        os.remove('%s/%s.log' % (self.root, item))
+  
+  # Log in plaintext; this is designed for being read in MATLAB(sorry not sorry)
+  def log(self, itr, **kwargs):
+    for arg in kwargs:
+      if arg not in self.metrics:
+        if self.reinitialize:
+          self.reinit(arg)
+        self.metrics += [arg]
+      if self.logstyle == 'pickle':
+        print('Pickle not currently supported...')
+         # with open('%s/%s.log' % (self.root, arg), 'a') as f:
+          # pickle.dump(kwargs[arg], f)
+      elif self.logstyle == 'mat':
+        print('.mat logstyle not currently supported...')
+      else:
+        with open('%s/%s.log' % (self.root, arg), 'a') as f:
+          f.write('%d: %s\n' % (itr, self.logstyle % kwargs[arg]))
+
+
+# Write some metadata to the logs directory
+def write_metadata(logs_root, experiment_name, config, state_dict):
+  with open(('%s/%s/metalog.txt' % 
+             (logs_root, experiment_name)), 'w') as writefile:
+    writefile.write('datetime: %s\n' % str(datetime.datetime.now()))
+    writefile.write('config: %s\n' % str(config))
+    writefile.write('state: %s\n' %str(state_dict))
+
+
+"""
+Very basic progress indicator to wrap an iterable in.
+
+Author: Jan Schlüter
+Andy's adds: time elapsed in addition to ETA, makes it possible to add
+estimated time to 1k iters instead of estimated time to completion.
+"""
+def progress(items, desc='', total=None, min_delay=0.1, displaytype='s1k'):
+  """
+  Returns a generator over `items`, printing the number and percentage of
+  items processed and the estimated remaining processing time before yielding
+  the next item. `total` gives the total number of items (required if `items`
+  has no length), and `min_delay` gives the minimum time in seconds between
+  subsequent prints. `desc` gives an optional prefix text (end with a space).
+  """
+  total = total or len(items)
+  t_start = time.time()
+  t_last = 0
+  for n, item in enumerate(items):
+    t_now = time.time()
+    if t_now - t_last > min_delay:
+      print("\r%s%d/%d (%6.2f%%)" % (
+              desc, n+1, total, n / float(total) * 100), end=" ")
+      if n > 0:
+        
+        if displaytype == 's1k': # minutes/seconds for 1000 iters
+          next_1000 = n + (1000 - n%1000)
+          t_done = t_now - t_start
+          t_1k = t_done / n * next_1000
+          outlist = list(divmod(t_done, 60)) + list(divmod(t_1k - t_done, 60))
+          print("(TE/ET1k: %d:%02d / %d:%02d)" % tuple(outlist), end=" ")
+        else:# displaytype == 'eta':
+          t_done = t_now - t_start
+          t_total = t_done / n * total
+          outlist = list(divmod(t_done, 60)) + list(divmod(t_total - t_done, 60))
+          print("(TE/ETA: %d:%02d / %d:%02d)" % tuple(outlist), end=" ")
+          
+      sys.stdout.flush()
+      t_last = t_now
+    yield item
+  t_total = time.time() - t_start
+  print("\r%s%d/%d (100.00%%) (took %d:%02d)" % ((desc, total, total) +
+                                                   divmod(t_total, 60)))
+
+
+# Sample function for use with inception metrics
+def sample(G, z_, y_, config):
+  with torch.no_grad():
+    z_.sample_()
+    y_.sample_()
+    if config['parallel']:
+      G_z =  nn.parallel.data_parallel(G, (z_, G.shared(y_)))
+    else:
+      G_z = G(z_, G.shared(y_))
+    return G_z, y_
+
+
+# Sample function for sample sheets
+def sample_sheet(G, classes_per_sheet, num_classes, samples_per_class, parallel,
+                 samples_root, experiment_name, folder_number, z_=None):
+  # Prepare sample directory
+  if not os.path.isdir('%s/%s' % (samples_root, experiment_name)):
+    os.mkdir('%s/%s' % (samples_root, experiment_name))
+  if not os.path.isdir('%s/%s/%d' % (samples_root, experiment_name, folder_number)):
+    os.mkdir('%s/%s/%d' % (samples_root, experiment_name, folder_number))
+  # loop over total number of sheets
+  for i in range(num_classes // classes_per_sheet):
+    ims = []
+    y = torch.arange(i * classes_per_sheet, (i + 1) * classes_per_sheet, device='cuda')
+    for j in range(samples_per_class):
+      if (z_ is not None) and hasattr(z_, 'sample_') and classes_per_sheet <= z_.size(0):
+        z_.sample_()
+      else:
+        z_ = torch.randn(classes_per_sheet, G.dim_z, device='cuda')        
+      with torch.no_grad():
+        if parallel:
+          o = nn.parallel.data_parallel(G, (z_[:classes_per_sheet], G.shared(y)))
+        else:
+          o = G(z_[:classes_per_sheet], G.shared(y))
+
+      ims += [o.data.cpu()]
+    # This line should properly unroll the images
+    out_ims = torch.stack(ims, 1).view(-1, ims[0].shape[1], ims[0].shape[2], 
+                                       ims[0].shape[3]).data.float().cpu()
+    # The path for the samples
+    image_filename = '%s/%s/%d/samples%d.jpg' % (samples_root, experiment_name, 
+                                                 folder_number, i)
+    torchvision.utils.save_image(out_ims, image_filename,
+                                 nrow=samples_per_class, normalize=True)
+
+
+# Interp function; expects x0 and x1 to be of shape (shape0, 1, rest_of_shape..)
+def interp(x0, x1, num_midpoints):
+  lerp = torch.linspace(0, 1.0, num_midpoints + 2, device='cuda').to(x0.dtype)
+  return ((x0 * (1 - lerp.view(1, -1, 1))) + (x1 * lerp.view(1, -1, 1)))
+
+
+# interp sheet function
+# Supports full, class-wise and intra-class interpolation
+def interp_sheet(G, num_per_sheet, num_midpoints, num_classes, parallel,
+                 samples_root, experiment_name, folder_number, sheet_number=0,
+                 fix_z=False, fix_y=False, device='cuda'):
+  # Prepare zs and ys
+  if fix_z: # If fix Z, only sample 1 z per row
+    zs = torch.randn(num_per_sheet, 1, G.dim_z, device=device)
+    zs = zs.repeat(1, num_midpoints + 2, 1).view(-1, G.dim_z)
+  else:
+    zs = interp(torch.randn(num_per_sheet, 1, G.dim_z, device=device),
+                torch.randn(num_per_sheet, 1, G.dim_z, device=device),
+                num_midpoints).view(-1, G.dim_z)
+  if fix_y: # If fix y, only sample 1 z per row
+    ys = sample_1hot(num_per_sheet, num_classes)
+    ys = G.shared(ys).view(num_per_sheet, 1, -1)
+    ys = ys.repeat(1, num_midpoints + 2, 1).view(num_per_sheet * (num_midpoints + 2), -1)
+  else:
+    ys = interp(G.shared(sample_1hot(num_per_sheet, num_classes)).view(num_per_sheet, 1, -1),
+                G.shared(sample_1hot(num_per_sheet, num_classes)).view(num_per_sheet, 1, -1),
+                num_midpoints).view(num_per_sheet * (num_midpoints + 2), -1)
+  # Run the net--note that we've already passed y through G.shared.
+  if G.fp16:
+    zs = zs.half()
+  with torch.no_grad():
+    if parallel:
+      out_ims = nn.parallel.data_parallel(G, (zs, ys)).data.cpu()
+    else:
+      out_ims = G(zs, ys).data.cpu()
+  interp_style = '' + ('Z' if not fix_z else '') + ('Y' if not fix_y else '')
+  image_filename = '%s/%s/%d/interp%s%d.jpg' % (samples_root, experiment_name,
+                                                folder_number, interp_style,
+                                                sheet_number)
+  torchvision.utils.save_image(out_ims, image_filename,
+                               nrow=num_midpoints + 2, normalize=True)
+
+
+# Convenience debugging function to print out gradnorms and shape from each layer
+# May need to rewrite this so we can actually see which parameter is which
+def print_grad_norms(net):
+    gradsums = [[float(torch.norm(param.grad).item()),
+                 float(torch.norm(param).item()), param.shape]
+                for param in net.parameters()]
+    order = np.argsort([item[0] for item in gradsums])
+    print(['%3.3e,%3.3e, %s' % (gradsums[item_index][0],
+                                gradsums[item_index][1],
+                                str(gradsums[item_index][2])) 
+                              for item_index in order])
+
+
+# Get singular values to log. This will use the state dict to find them
+# and substitute underscores for dots.
+def get_SVs(net, prefix):
+  d = net.state_dict()
+  return {('%s_%s' % (prefix, key)).replace('.', '_') :
+            float(d[key].item())
+            for key in d if 'sv' in key}
+
+
+# Name an experiment based on its config
+def name_from_config(config):
+  name = '_'.join([
+  item for item in [
+  'Big%s' % config['which_train_fn'],
+  config['dataset'],
+  config['model'] if config['model'] != 'BigGAN' else None,
+  'seed%d' % config['seed'],
+  'Gch%d' % config['G_ch'],
+  'Dch%d' % config['D_ch'],
+  'Gd%d' % config['G_depth'] if config['G_depth'] > 1 else None,
+  'Dd%d' % config['D_depth'] if config['D_depth'] > 1 else None,
+  'bs%d' % config['batch_size'],
+  'Gfp16' if config['G_fp16'] else None,
+  'Dfp16' if config['D_fp16'] else None,
+  'nDs%d' % config['num_D_steps'] if config['num_D_steps'] > 1 else None,
+  'nDa%d' % config['num_D_accumulations'] if config['num_D_accumulations'] > 1 else None,
+  'nGa%d' % config['num_G_accumulations'] if config['num_G_accumulations'] > 1 else None,
+  'Glr%2.1e' % config['G_lr'],
+  'Dlr%2.1e' % config['D_lr'],
+  'GB%3.3f' % config['G_B1'] if config['G_B1'] !=0.0 else None,
+  'GBB%3.3f' % config['G_B2'] if config['G_B2'] !=0.999 else None,
+  'DB%3.3f' % config['D_B1'] if config['D_B1'] !=0.0 else None,
+  'DBB%3.3f' % config['D_B2'] if config['D_B2'] !=0.999 else None,
+  'Gnl%s' % config['G_nl'],
+  'Dnl%s' % config['D_nl'],
+  'Ginit%s' % config['G_init'],
+  'Dinit%s' % config['D_init'],
+  'G%s' % config['G_param'] if config['G_param'] != 'SN' else None,
+  'D%s' % config['D_param'] if config['D_param'] != 'SN' else None,
+  'Gattn%s' % config['G_attn'] if config['G_attn'] != '0' else None,
+  'Dattn%s' % config['D_attn'] if config['D_attn'] != '0' else None,
+  'Gortho%2.1e' % config['G_ortho'] if config['G_ortho'] > 0.0 else None,
+  'Dortho%2.1e' % config['D_ortho'] if config['D_ortho'] > 0.0 else None,
+  config['norm_style'] if config['norm_style'] != 'bn' else None,
+  'cr' if config['cross_replica'] else None,
+  'Gshared' if config['G_shared'] else None,
+  'hier' if config['hier'] else None,
+  'ema' if config['ema'] else None,
+  config['name_suffix'] if config['name_suffix'] else None,
+  ]
+  if item is not None])
+  # dogball
+  if config['hashname']:
+    return hashname(name)
+  else:
+    return name
+
+
+# A simple function to produce a unique experiment name from the animal hashes.
+def hashname(name):
+  h = hash(name)
+  a = h % len(animal_hash.a)
+  h = h // len(animal_hash.a)
+  b = h % len(animal_hash.b)
+  h = h // len(animal_hash.c)
+  c = h % len(animal_hash.c)
+  return animal_hash.a[a] + animal_hash.b[b] + animal_hash.c[c]
+
+
+# Get GPU memory, -i is the index
+def query_gpu(indices):
+  os.system('nvidia-smi -i 0 --query-gpu=memory.free --format=csv')
+
+
+# Convenience function to count the number of parameters in a module
+def count_parameters(module):
+  print('Number of parameters: {}'.format(
+    sum([p.data.nelement() for p in module.parameters()])))
+
+   
+# Convenience function to sample an index, not actually a 1-hot
+def sample_1hot(batch_size, num_classes, device='cuda'):
+  return torch.randint(low=0, high=num_classes, size=(batch_size,),
+          device=device, dtype=torch.int64, requires_grad=False)
+
+
+# A highly simplified convenience class for sampling from distributions
+# One could also use PyTorch's inbuilt distributions package.
+# Note that this class requires initialization to proceed as
+# x = Distribution(torch.randn(size))
+# x.init_distribution(dist_type, **dist_kwargs)
+# x = x.to(device,dtype)
+# This is partially based on https://discuss.pytorch.org/t/subclassing-torch-tensor/23754/2
+class Distribution(torch.Tensor):
+  # Init the params of the distribution
+  def init_distribution(self, dist_type, **kwargs):    
+    self.dist_type = dist_type
+    self.dist_kwargs = kwargs
+    if self.dist_type == 'normal':
+      self.mean, self.var = kwargs['mean'], kwargs['var']
+    elif self.dist_type == 'categorical':
+      self.num_categories = kwargs['num_categories']
+
+  def sample_(self):
+    if self.dist_type == 'normal':
+      self.normal_(self.mean, self.var)
+    elif self.dist_type == 'categorical':
+      self.random_(0, self.num_categories)    
+    # return self.variable
+    
+  # Silly hack: overwrite the to() method to wrap the new object
+  # in a distribution as well
+  def to(self, *args, **kwargs):
+    new_obj = Distribution(self)
+    new_obj.init_distribution(self.dist_type, **self.dist_kwargs)
+    new_obj.data = super().to(*args, **kwargs)    
+    return new_obj
+
+
+# Convenience function to prepare a z and y vector
+def prepare_z_y(G_batch_size, dim_z, nclasses, device='cuda', 
+                fp16=False,z_var=1.0):
+  z_ = Distribution(torch.randn(G_batch_size, dim_z, requires_grad=False))
+  z_.init_distribution('normal', mean=0, var=z_var)
+  z_ = z_.to(device,torch.float16 if fp16 else torch.float32)   
+  
+  if fp16:
+    z_ = z_.half()
+
+  y_ = Distribution(torch.zeros(G_batch_size, requires_grad=False))
+  y_.init_distribution('categorical',num_categories=nclasses)
+  y_ = y_.to(device, torch.int64)
+  return z_, y_
+
+
+def initiate_standing_stats(net):
+  for module in net.modules():
+    if hasattr(module, 'accumulate_standing'):
+      module.reset_stats()
+      module.accumulate_standing = True
+
+
+def accumulate_standing_stats(net, z, y, nclasses, num_accumulations=16):
+  initiate_standing_stats(net)
+  net.train()
+  for i in range(num_accumulations):
+    with torch.no_grad():
+      z.normal_()
+      y.random_(0, nclasses)
+      x = net(z, net.shared(y)) # No need to parallelize here unless using syncbn
+  # Set to eval mode
+  net.eval() 
+
+
+# This version of Adam keeps an fp32 copy of the parameters and
+# does all of the parameter updates in fp32, while still doing the
+# forwards and backwards passes using fp16 (i.e. fp16 copies of the
+# parameters and fp16 activations).
+#
+# Note that this calls .float().cuda() on the params.
+import math
+from torch.optim.optimizer import Optimizer
+class Adam16(Optimizer):
+  def __init__(self, params, lr=1e-3, betas=(0.9, 0.999), eps=1e-8,weight_decay=0):
+    defaults = dict(lr=lr, betas=betas, eps=eps,
+            weight_decay=weight_decay)
+    params = list(params)
+    super(Adam16, self).__init__(params, defaults)
+      
+  # Safety modification to make sure we floatify our state
+  def load_state_dict(self, state_dict):
+    super(Adam16, self).load_state_dict(state_dict)
+    for group in self.param_groups:
+      for p in group['params']:
+        self.state[p]['exp_avg'] = self.state[p]['exp_avg'].float()
+        self.state[p]['exp_avg_sq'] = self.state[p]['exp_avg_sq'].float()
+        self.state[p]['fp32_p'] = self.state[p]['fp32_p'].float()
+
+  def step(self, closure=None):
+    """Performs a single optimization step.
+    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.float()
+        state = self.state[p]
+
+        # State initialization
+        if len(state) == 0:
+          state['step'] = 0
+          # Exponential moving average of gradient values
+          state['exp_avg'] = grad.new().resize_as_(grad).zero_()
+          # Exponential moving average of squared gradient values
+          state['exp_avg_sq'] = grad.new().resize_as_(grad).zero_()
+          # Fp32 copy of the weights
+          state['fp32_p'] = p.data.float()
+
+        exp_avg, exp_avg_sq = state['exp_avg'], state['exp_avg_sq']
+        beta1, beta2 = group['betas']
+
+        state['step'] += 1
+
+        if group['weight_decay'] != 0:
+          grad = grad.add(group['weight_decay'], state['fp32_p'])
+
+        # Decay the first and second moment running average coefficient
+        exp_avg.mul_(beta1).add_(1 - beta1, grad)
+        exp_avg_sq.mul_(beta2).addcmul_(1 - beta2, grad, grad)
+
+        denom = exp_avg_sq.sqrt().add_(group['eps'])
+
+        bias_correction1 = 1 - beta1 ** state['step']
+        bias_correction2 = 1 - beta2 ** state['step']
+        step_size = group['lr'] * math.sqrt(bias_correction2) / bias_correction1
+      
+        state['fp32_p'].addcdiv_(-step_size, exp_avg, denom)
+        p.data = state['fp32_p'].half()
+
+    return loss

src/models/cvae.py

@@ -0,0 +1,75 @@
+from torch import cat
+from torch.optim import Adam
+from torch.nn import Sequential, ModuleList, \
+                     Conv2d, Linear, \
+                     LeakyReLU, Tanh, \
+                     BatchNorm1d, BatchNorm2d, \
+                     ConvTranspose2d, UpsamplingBilinear2d
+
+from .neuralnetwork import NeuralNetwork
+
+
+# parameters for cVAE
+colors_dim = 3
+labels_dim = 37
+momentum = 0.99  # Batchnorm
+negative_slope = 0.2  # LeakyReLU
+optimizer = Adam
+betas = (0.5, 0.999)
+
+# hyperparameters
+learning_rate = 2e-4
+latent_dim = 128
+
+
+def genUpsample(input_channels, output_channels, stride, pad):
+   return Sequential(
+        ConvTranspose2d(input_channels, output_channels, 4, stride, pad, bias=False),
+        BatchNorm2d(output_channels),
+        LeakyReLU(negative_slope=negative_slope))
+
+
+def genUpsample2(input_channels, output_channels, kernel_size):
+   return Sequential(
+        Conv2d(input_channels, output_channels, kernel_size=kernel_size, stride=1, padding= (kernel_size-1) // 2),
+        BatchNorm2d(output_channels),
+        LeakyReLU(negative_slope=negative_slope),
+        Conv2d(output_channels, output_channels, kernel_size=kernel_size, stride=1, padding= (kernel_size-1) // 2),
+        BatchNorm2d(output_channels),
+        LeakyReLU(negative_slope=negative_slope),
+        UpsamplingBilinear2d(scale_factor=2))
+
+
+class ConditionalDecoder(NeuralNetwork):
+    def __init__(self, ll_scaling=1.0, dim_z=latent_dim):
+        super(ConditionalDecoder, self).__init__()
+        self.dim_z = dim_z
+        ngf = 32
+        self.init = genUpsample(self.dim_z, ngf * 16, 1, 0)
+        self.embedding = Sequential(
+            Linear(labels_dim, self.dim_z),
+            BatchNorm1d(self.dim_z, momentum=momentum),
+            LeakyReLU(negative_slope=negative_slope),
+        )
+        self.dense_init = Sequential(
+            Linear(self.dim_z*2, self.dim_z),
+            BatchNorm1d(self.dim_z, momentum=momentum),
+            LeakyReLU(negative_slope=negative_slope),
+        )
+        self.m_modules = ModuleList()  # to 4x4
+        self.c_modules = ModuleList()
+        for i in range(4):
+            self.m_modules.append(genUpsample2(ngf * 2**(4-i), ngf * 2**(3-i), 3))
+            self.c_modules.append(Sequential(Conv2d(ngf * 2**(3-i), colors_dim, 3, 1, 1, bias=False), Tanh()))
+        self.set_optimizer(optimizer, lr=learning_rate*ll_scaling, betas=betas)
+
+    def forward(self, latent, labels, step=3):
+        y = self.embedding(labels)
+        out = cat((latent, y), dim=1)
+        out = self.dense_init(out)
+        out = out.unsqueeze(2).unsqueeze(3)
+        out = self.init(out)
+        for i in range(step):
+            out = self.m_modules[i](out)
+        out = self.c_modules[step](self.m_modules[step](out))
+        return out

src/models/infoscc_gan.py

@@ -0,0 +1,272 @@
+from typing import Optional, Dict
+from functools import partial
+import math
+
+import torch
+import torch.nn as nn
+
+
+def get_activation(activation: str = "lrelu"):
+    actv_layers = {
+        "relu":  nn.ReLU,
+        "lrelu": partial(nn.LeakyReLU, 0.2),
+    }
+    assert activation in actv_layers, f"activation [{activation}] not implemented"
+    return actv_layers[activation]
+
+
+def get_normalization(normalization: str = "batch_norm"):
+    norm_layers = {
+        "instance_norm": nn.InstanceNorm2d,
+        "batch_norm":    nn.BatchNorm2d,
+        "group_norm":    partial(nn.GroupNorm, num_groups=8),
+        "layer_norm":    partial(nn.GroupNorm, num_groups=1),
+    }
+    assert normalization in norm_layers, f"normalization [{normalization}] not implemented"
+    return norm_layers[normalization]
+
+
+class ConvLayer(nn.Sequential):
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        kernel_size: int = 3,
+        stride: int = 1,
+        padding: Optional[int] = 1,
+        padding_mode: str = "zeros",
+        groups: int = 1,
+        bias: bool = True,
+        transposed: bool = False,
+        normalization: Optional[str] = None,
+        activation: Optional[str] = "lrelu",
+        pre_activate: bool = False,
+    ):
+        if transposed:
+            conv = partial(nn.ConvTranspose2d, output_padding=stride-1)
+            padding_mode = "zeros"
+        else:
+            conv = nn.Conv2d
+        layers = [
+            conv(
+                in_channels,
+                out_channels,
+                kernel_size=kernel_size,
+                stride=stride,
+                padding=padding,
+                padding_mode=padding_mode,
+                groups=groups,
+                bias=bias,
+            )
+        ]
+
+        norm_actv = []
+        if normalization is not None:
+            norm_actv.append(
+                get_normalization(normalization)(
+                    num_channels=in_channels if pre_activate else out_channels
+                )
+            )
+        if activation is not None:
+            norm_actv.append(
+                get_activation(activation)(inplace=True)
+            )
+
+        if pre_activate:
+            layers = norm_actv + layers
+        else:
+            layers = layers + norm_actv
+
+        super().__init__(
+            *layers
+        )
+
+
+class SubspaceLayer(nn.Module):
+    def __init__(
+        self,
+        dim: int,
+        n_basis: int,
+    ):
+        super().__init__()
+
+        self.U = nn.Parameter(torch.empty(n_basis, dim))
+        nn.init.orthogonal_(self.U)
+        self.L = nn.Parameter(torch.FloatTensor([3 * i for i in range(n_basis, 0, -1)]))
+        self.mu = nn.Parameter(torch.zeros(dim))
+
+    def forward(self, z):
+        return (self.L * z) @ self.U + self.mu
+
+
+class EigenBlock(nn.Module):
+    def __init__(
+        self,
+        width: int,
+        height: int,
+        in_channels: int,
+        out_channels: int,
+        n_basis: int,
+    ):
+        super().__init__()
+
+        self.projection = SubspaceLayer(dim=width*height*in_channels, n_basis=n_basis)
+        self.subspace_conv1 = ConvLayer(
+            in_channels,
+            in_channels,
+            kernel_size=1,
+            stride=1,
+            padding=0,
+            transposed=True,
+            activation=None,
+            normalization=None,
+        )
+        self.subspace_conv2 = ConvLayer(
+            in_channels,
+            out_channels,
+            kernel_size=3,
+            stride=2,
+            padding=1,
+            transposed=True,
+            activation=None,
+            normalization=None,
+        )
+
+        self.feature_conv1 = ConvLayer(
+            in_channels,
+            out_channels,
+            kernel_size=3,
+            stride=2,
+            transposed=True,
+            pre_activate=True,
+        )
+        self.feature_conv2 = ConvLayer(
+            out_channels,
+            out_channels,
+            kernel_size=3,
+            stride=1,
+            transposed=True,
+            pre_activate=True,
+        )
+
+    def forward(self, z, h):
+        phi = self.projection(z).view(h.shape)
+        h = self.feature_conv1(h + self.subspace_conv1(phi))
+        h = self.feature_conv2(h + self.subspace_conv2(phi))
+        return h
+
+
+class ConditionalGenerator(nn.Module):
+
+    """Conditional generator
+    It generates images from one hot label + noise sampled from N(0, 1) with explorable z injection space
+    Based on EigenGAN
+    """
+
+    def __init__(self,
+                 size: int,
+                 y_size: int,
+                 z_size: int,
+                 out_channels: int = 3,
+                 n_basis: int = 6,
+                 noise_dim: int = 512,
+                 base_channels: int = 16,
+                 max_channels: int = 512,
+                 y_type: str = 'one_hot'):
+
+        if y_type not in ['one_hot', 'multi_label', 'mixed', 'real']:
+            raise ValueError('Unsupported `y_type`')
+
+        super(ConditionalGenerator, self).__init__()
+
+        assert (size & (size - 1) == 0) and size != 0, "img size should be a power of 2"
+
+        self.y_type = y_type
+        self.y_size = y_size
+        self.eps_size = z_size
+
+        self.noise_dim = noise_dim
+        self.n_basis = n_basis
+        self.n_blocks = int(math.log(size, 2)) - 2
+
+        def get_channels(i_block):
+            return min(max_channels, base_channels * (2 ** (self.n_blocks - i_block)))
+
+        self.y_fc = nn.Linear(self.y_size, self.y_size)
+        self.concat_fc = nn.Linear(self.y_size + self.eps_size, self.noise_dim)
+
+        self.fc = nn.Linear(self.noise_dim, 4 * 4 * get_channels(0))
+
+        self.blocks = nn.ModuleList()
+        for i in range(self.n_blocks):
+            self.blocks.append(
+                EigenBlock(
+                    width=4 * (2 ** i),
+                    height=4 * (2 ** i),
+                    in_channels=get_channels(i),
+                    out_channels=get_channels(i + 1),
+                    n_basis=self.n_basis,
+                )
+            )
+
+        self.out = nn.Sequential(
+            ConvLayer(base_channels, out_channels, kernel_size=7, stride=1, padding=3, pre_activate=True),
+            nn.Tanh(),
+        )
+
+    def forward(self,
+                y: torch.Tensor,
+                eps: Optional[torch.Tensor] = None,
+                zs: Optional[torch.Tensor] = None,
+                return_eps: bool = False):
+
+        bs = y.size(0)
+
+        if eps is None:
+            eps = self.sample_eps(bs)
+
+        if zs is None:
+            zs = self.sample_zs(bs)
+
+        y_out = self.y_fc(y)
+        concat = torch.cat((y_out, eps), dim=1)
+        concat = self.concat_fc(concat)
+
+        out = self.fc(concat).view(len(eps), -1, 4, 4)
+        for block, z in zip(self.blocks, zs.permute(1, 0, 2)):
+            out = block(z, out)
+        out = self.out(out)
+
+        if return_eps:
+            return out, concat
+
+        return out
+
+    def sample_zs(self, batch: int, truncation: float = 1.):
+        device = self.get_device()
+        zs = torch.randn(batch, self.n_blocks, self.n_basis, device=device)
+
+        if truncation < 1.:
+            zs = torch.zeros_like(zs) * (1 - truncation) + zs * truncation
+        return zs
+
+    def sample_eps(self, batch: int, truncation: float = 1.):
+        device = self.get_device()
+        eps = torch.randn(batch, self.eps_size, device=device)
+
+        if truncation < 1.:
+            eps = torch.zeros_like(eps) * (1 - truncation) + eps * truncation
+        return eps
+
+    def get_device(self):
+        return self.fc.weight.device
+
+    def orthogonal_regularizer(self):
+        reg = []
+        for layer in self.modules():
+            if isinstance(layer, SubspaceLayer):
+                UUT = layer.U @ layer.U.t()
+                reg.append(
+                    ((UUT - torch.eye(UUT.shape[0], device=UUT.device)) ** 2).mean()
+                )
+        return sum(reg) / len(reg)

src/models/neuralnetwork.py

@@ -0,0 +1,46 @@
+import torch
+
+
+class NeuralNetwork(torch.nn.Module):
+    """ base class with convenient procedures used by all NN"""
+    def __init__(self):
+        super(NeuralNetwork, self).__init__()
+        self.parameter_file = f"parameter_state_dict_{self._get_name()}.pth"
+        # self.cuda() ## all NN shall run on cuda ### doesnt seem to work
+
+    def save(self) -> None:
+        """ save learned parameters to parameter_file """
+        torch.save(self.state_dict(), self.parameter_file)
+
+    def load(self) -> None:
+        """ load learned parameters from parameter_file """
+        self.load_state_dict(torch.load(self.parameter_file))
+        self.eval()
+
+    @staticmethod
+    def same_padding(kernel_size=1) -> float:
+        """ return padding required to mimic 'same' padding in tensorflow """
+        return (kernel_size-1) // 2
+
+    def set_optimizer(self, optimizer, **kwargs) -> None:
+        self.optimizer = optimizer(self.parameters(), **kwargs)
+
+    def get_total_number_parameters(self) -> float:
+        """ return total number of parameters """
+        return sum([p.numel() for p in classifier.parameters()])
+
+    def zero_grad(self):
+        """ faster implementation of zero_grad """
+        for p in self.parameters():
+            p.grad = None
+#        self.zero_grad(set_to_none=True)
+
+
+def update_networks_on_loss(loss: torch.Tensor, *networks) -> None:
+    if not loss:
+        return
+    for network in networks:
+        network.zero_grad()
+    loss.backward()
+    for network in networks:
+        network.optimizer.step()

src/models/parameter.py

@@ -0,0 +1,53 @@
+""" hardcoded parameter
+
+these can be changed in a jupyter notebook during runtime via
+
+>>> import parameter
+>>> parameter.parameter = new_value
+
+"""
+
+from torch.optim import Adam
+
+###############
+## hardcoded ##
+###############
+
+
+# Input
+image_dim = 64
+colors_dim = 3
+labels_dim = 37 #3
+input_size = (colors_dim,image_dim,image_dim)
+
+
+#############
+## mutable ##
+#############
+
+class Parameter:
+    """ container for hyperparameters"""
+
+    def __init__(self):
+        # Encoder/Decoder
+        self.latent_dim = 8
+        self.decoder_dim = self.latent_dim  # differs from latent_dim if PCA applied before decoder
+
+        # General
+        self.learning_rate = 0.0002
+        self.betas = (0.5,0.999)  ## 0.999 is default beta2 in tensorflow
+        self.optimizer = Adam
+        self.negative_slope = 0.2 # for LeakyReLU
+        self.momentum = 0.99 # for BatchNorm
+
+        # Loss weights
+        self.alpha = 1 # switch VAE (1) / AE (0)
+        self.beta = 1 # weight for KL-loss
+        self.gamma = 1024 # weight for learned-metric-loss (https://arxiv.org/pdf/1512.09300.pdf)
+        self.delta = 1 # weight for class-loss
+        self.zeta = 0.5 # weight for MSE-loss
+
+    def return_parameter_dict(self):
+        return(self.__dict__)
+
+parameter = Parameter()

src/utils/__init__.py

@@ -0,0 +1,2 @@
+from .utils import download_file
+from .utils import sample_labels

src/utils/utils.py

@@ -0,0 +1,14 @@
+import numpy as np
+import gdown
+
+import torch
+
+
+def download_file(file_id: str, output_path: str):
+    gdown.download(f'https://drive.google.com/uc?id={file_id}', output_path)
+
+
+def sample_labels(labels: torch.Tensor, n: int) -> torch.Tensor:
+    high = labels.shape[0]
+    idx = np.random.randint(0, high, size=n)
+    return labels[idx]