First commit

LICENSE.txt

@@ -0,0 +1,97 @@
+Copyright (c) 2020, NVIDIA Corporation. All rights reserved.
+
+
+NVIDIA Source Code License for StyleGAN2 with Adaptive Discriminator Augmentation (ADA)
+
+
+=======================================================================
+
+1. Definitions
+
+"Licensor" means any person or entity that distributes its Work.
+
+"Software" means the original work of authorship made available under
+this License.
+
+"Work" means the Software and any additions to or derivative works of
+the Software that are made available under this License.
+
+The terms "reproduce," "reproduction," "derivative works," and
+"distribution" have the meaning as provided under U.S. copyright law;
+provided, however, that for the purposes of this License, derivative
+works shall not include works that remain separable from, or merely
+link (or bind by name) to the interfaces of, the Work.
+
+Works, including the Software, are "made available" under this License
+by including in or with the Work either (a) a copyright notice
+referencing the applicability of this License to the Work, or (b) a
+copy of this License.
+
+2. License Grants
+
+    2.1 Copyright Grant. Subject to the terms and conditions of this
+    License, each Licensor grants to you a perpetual, worldwide,
+    non-exclusive, royalty-free, copyright license to reproduce,
+    prepare derivative works of, publicly display, publicly perform,
+    sublicense and distribute its Work and any resulting derivative
+    works in any form.
+
+3. Limitations
+
+    3.1 Redistribution. You may reproduce or distribute the Work only
+    if (a) you do so under this License, (b) you include a complete
+    copy of this License with your distribution, and (c) you retain
+    without modification any copyright, patent, trademark, or
+    attribution notices that are present in the Work.
+
+    3.2 Derivative Works. You may specify that additional or different
+    terms apply to the use, reproduction, and distribution of your
+    derivative works of the Work ("Your Terms") only if (a) Your Terms
+    provide that the use limitation in Section 3.3 applies to your
+    derivative works, and (b) you identify the specific derivative
+    works that are subject to Your Terms. Notwithstanding Your Terms,
+    this License (including the redistribution requirements in Section
+    3.1) will continue to apply to the Work itself.
+
+    3.3 Use Limitation. The Work and any derivative works thereof only
+    may be used or intended for use non-commercially. Notwithstanding
+    the foregoing, NVIDIA and its affiliates may use the Work and any
+    derivative works commercially. As used herein, "non-commercially"
+    means for research or evaluation purposes only.
+
+    3.4 Patent Claims. If you bring or threaten to bring a patent claim
+    against any Licensor (including any claim, cross-claim or
+    counterclaim in a lawsuit) to enforce any patents that you allege
+    are infringed by any Work, then your rights under this License from
+    such Licensor (including the grant in Section 2.1) will terminate
+    immediately.
+
+    3.5 Trademarks. This License does not grant any rights to use any
+    Licensor’s or its affiliates’ names, logos, or trademarks, except
+    as necessary to reproduce the notices described in this License.
+
+    3.6 Termination. If you violate any term of this License, then your
+    rights under this License (including the grant in Section 2.1) will
+    terminate immediately.
+
+4. Disclaimer of Warranty.
+
+THE WORK IS PROVIDED "AS IS" WITHOUT WARRANTIES OR CONDITIONS OF ANY
+KIND, EITHER EXPRESS OR IMPLIED, INCLUDING WARRANTIES OR CONDITIONS OF
+MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE, TITLE OR
+NON-INFRINGEMENT. YOU BEAR THE RISK OF UNDERTAKING ANY ACTIVITIES UNDER
+THIS LICENSE.
+
+5. Limitation of Liability.
+
+EXCEPT AS PROHIBITED BY APPLICABLE LAW, IN NO EVENT AND UNDER NO LEGAL
+THEORY, WHETHER IN TORT (INCLUDING NEGLIGENCE), CONTRACT, OR OTHERWISE
+SHALL ANY LICENSOR BE LIABLE TO YOU FOR DAMAGES, INCLUDING ANY DIRECT,
+INDIRECT, SPECIAL, INCIDENTAL, OR CONSEQUENTIAL DAMAGES ARISING OUT OF
+OR RELATED TO THIS LICENSE, THE USE OR INABILITY TO USE THE WORK
+(INCLUDING BUT NOT LIMITED TO LOSS OF GOODWILL, BUSINESS INTERRUPTION,
+LOST PROFITS OR DATA, COMPUTER FAILURE OR MALFUNCTION, OR ANY OTHER
+COMMERCIAL DAMAGES OR LOSSES), EVEN IF THE LICENSOR HAS BEEN ADVISED OF
+THE POSSIBILITY OF SUCH DAMAGES.
+
+=======================================================================

README.md

@@ -0,0 +1,63 @@
+## rasm
+Arabic art using GANs. We currently have two models for generating calligraphy and mosaics.  
+
+## Notebooks 
+
+<table class="tg">
+  <tr>
+    <th class="tg-yw4l"><b>Name</b></th>
+    <th class="tg-yw4l"><b>Notebook</b></th>
+  </tr>
+  <tr>
+    <td class="tg-yw4l">Visualization</td>
+    <td class="tg-yw4l"><a href="https://colab.research.google.com/github/ARBML/rasm/blob/master/demo.ipynb">
+  <img src="https://colab.research.google.com/assets/colab-badge.svg" width = '100px' >
+</a></td>
+  </tr>
+</table>
+
+## Visualization 
+A set of functions for vis, interpolation and animation. Mostly tested in colab notebooks. 
+
+### Load Model 
+```python 
+from rasm import Rasm
+model = Rasm(mode = 'calligraphy')
+model = Rasm(mode = 'mosaics')
+```
+
+### Generate random 
+```python 
+model.generate_randomly()
+```
+
+### Generate grid 
+```python 
+model.generate_grid()
+```
+
+### Generate animation 
+```python
+model.generate_animation(size = 2, steps = 20)
+```
+
+![alt text](video.gif)
+
+## Sample Models 
+
+### Mosaics 
+![alt text](imgs/mosaic.png)
+![alt text](imgs/mosaicsv2.png)
+![alt text](imgs/mosaicsv3.png)
+![alt text](imgs/mosaicsv4.png)
+### Calligraphy 
+![alt text](imgs/calligraphyv2.PNG)
+![alt text](imgs/calligraphyv3.png)
+![alt text](imgs/calligraphyv4.png)
+![alt text](imgs/calligraphyv5.png)
+
+
+## References 
+- Gan-surgery: https://github.com/aydao/stylegan2-surgery
+- WikiArt model: https://github.com/pbaylies/stylegan2 
+- Starter-Notebook: https://github.com/Hephyrius/Stylegan2-Ada-Google-Colab-Starter-Notebook/

demo.ipynb

@@ -0,0 +1,610 @@
+{
+  "nbformat": 4,
+  "nbformat_minor": 0,
+  "metadata": {
+    "colab": {
+      "name": "SGAN Vis.ipynb",
+      "provenance": [],
+      "machine_shape": "hm"
+    },
+    "kernelspec": {
+      "name": "python3",
+      "display_name": "Python 3"
+    },
+    "accelerator": "GPU",
+    "widgets": {
+      "application/vnd.jupyter.widget-state+json": {
+        "edfc5ae9a3924ee6a04811ab3dec1656": {
+          "model_module": "@jupyter-widgets/controls",
+          "model_name": "VBoxModel",
+          "state": {
+            "_view_name": "VBoxView",
+            "_dom_classes": [],
+            "_model_name": "VBoxModel",
+            "_view_module": "@jupyter-widgets/controls",
+            "_model_module_version": "1.5.0",
+            "_view_count": null,
+            "_view_module_version": "1.5.0",
+            "box_style": "",
+            "layout": "IPY_MODEL_86ea189295a54dd0aa2a771a95b12f00",
+            "_model_module": "@jupyter-widgets/controls",
+            "children": [
+              "IPY_MODEL_61e1782fb995412fbef8e59c570e8823",
+              "IPY_MODEL_a03932d2929c4225b7735d0831a18566"
+            ]
+          }
+        },
+        "86ea189295a54dd0aa2a771a95b12f00": {
+          "model_module": "@jupyter-widgets/base",
+          "model_name": "LayoutModel",
+          "state": {
+            "_view_name": "LayoutView",
+            "grid_template_rows": null,
+            "right": null,
+            "justify_content": null,
+            "_view_module": "@jupyter-widgets/base",
+            "overflow": null,
+            "_model_module_version": "1.2.0",
+            "_view_count": null,
+            "flex_flow": null,
+            "width": null,
+            "min_width": null,
+            "border": null,
+            "align_items": null,
+            "bottom": null,
+            "_model_module": "@jupyter-widgets/base",
+            "top": null,
+            "grid_column": null,
+            "overflow_y": null,
+            "overflow_x": null,
+            "grid_auto_flow": null,
+            "grid_area": null,
+            "grid_template_columns": null,
+            "flex": null,
+            "_model_name": "LayoutModel",
+            "justify_items": null,
+            "grid_row": null,
+            "max_height": null,
+            "align_content": null,
+            "visibility": null,
+            "align_self": null,
+            "height": null,
+            "min_height": null,
+            "padding": null,
+            "grid_auto_rows": null,
+            "grid_gap": null,
+            "max_width": null,
+            "order": null,
+            "_view_module_version": "1.2.0",
+            "grid_template_areas": null,
+            "object_position": null,
+            "object_fit": null,
+            "grid_auto_columns": null,
+            "margin": null,
+            "display": null,
+            "left": null
+          }
+        },
+        "61e1782fb995412fbef8e59c570e8823": {
+          "model_module": "@jupyter-widgets/controls",
+          "model_name": "HTMLModel",
+          "state": {
+            "_view_name": "HTMLView",
+            "style": "IPY_MODEL_c5b6fad0fd114c90ac0f23d91546846f",
+            "_dom_classes": [],
+            "description": "",
+            "_model_name": "HTMLModel",
+            "placeholder": "​",
+            "_view_module": "@jupyter-widgets/controls",
+            "_model_module_version": "1.5.0",
+            "value": "Generating images: 1",
+            "_view_count": null,
+            "_view_module_version": "1.5.0",
+            "description_tooltip": null,
+            "_model_module": "@jupyter-widgets/controls",
+            "layout": "IPY_MODEL_7dea244035a34d0f8658e75a12b5cd89"
+          }
+        },
+        "a03932d2929c4225b7735d0831a18566": {
+          "model_module": "@jupyter-widgets/controls",
+          "model_name": "IntProgressModel",
+          "state": {
+            "_view_name": "ProgressView",
+            "style": "IPY_MODEL_13d6a5cf5103499ebfb40e2e9c520d27",
+            "_dom_classes": [],
+            "description": "",
+            "_model_name": "IntProgressModel",
+            "bar_style": "success",
+            "max": 1,
+            "_view_module": "@jupyter-widgets/controls",
+            "_model_module_version": "1.5.0",
+            "value": 1,
+            "_view_count": null,
+            "_view_module_version": "1.5.0",
+            "orientation": "horizontal",
+            "min": 0,
+            "description_tooltip": null,
+            "_model_module": "@jupyter-widgets/controls",
+            "layout": "IPY_MODEL_653de804a191499ca46a8eb6b91e7a6b"
+          }
+        },
+        "c5b6fad0fd114c90ac0f23d91546846f": {
+          "model_module": "@jupyter-widgets/controls",
+          "model_name": "DescriptionStyleModel",
+          "state": {
+            "_view_name": "StyleView",
+            "_model_name": "DescriptionStyleModel",
+            "description_width": "",
+            "_view_module": "@jupyter-widgets/base",
+            "_model_module_version": "1.5.0",
+            "_view_count": null,
+            "_view_module_version": "1.2.0",
+            "_model_module": "@jupyter-widgets/controls"
+          }
+        },
+        "7dea244035a34d0f8658e75a12b5cd89": {
+          "model_module": "@jupyter-widgets/base",
+          "model_name": "LayoutModel",
+          "state": {
+            "_view_name": "LayoutView",
+            "grid_template_rows": null,
+            "right": null,
+            "justify_content": null,
+            "_view_module": "@jupyter-widgets/base",
+            "overflow": null,
+            "_model_module_version": "1.2.0",
+            "_view_count": null,
+            "flex_flow": null,
+            "width": null,
+            "min_width": null,
+            "border": null,
+            "align_items": null,
+            "bottom": null,
+            "_model_module": "@jupyter-widgets/base",
+            "top": null,
+            "grid_column": null,
+            "overflow_y": null,
+            "overflow_x": null,
+            "grid_auto_flow": null,
+            "grid_area": null,
+            "grid_template_columns": null,
+            "flex": null,
+            "_model_name": "LayoutModel",
+            "justify_items": null,
+            "grid_row": null,
+            "max_height": null,
+            "align_content": null,
+            "visibility": null,
+            "align_self": null,
+            "height": null,
+            "min_height": null,
+            "padding": null,
+            "grid_auto_rows": null,
+            "grid_gap": null,
+            "max_width": null,
+            "order": null,
+            "_view_module_version": "1.2.0",
+            "grid_template_areas": null,
+            "object_position": null,
+            "object_fit": null,
+            "grid_auto_columns": null,
+            "margin": null,
+            "display": null,
+            "left": null
+          }
+        },
+        "13d6a5cf5103499ebfb40e2e9c520d27": {
+          "model_module": "@jupyter-widgets/controls",
+          "model_name": "ProgressStyleModel",
+          "state": {
+            "_view_name": "StyleView",
+            "_model_name": "ProgressStyleModel",
+            "description_width": "",
+            "_view_module": "@jupyter-widgets/base",
+            "_model_module_version": "1.5.0",
+            "_view_count": null,
+            "_view_module_version": "1.2.0",
+            "bar_color": null,
+            "_model_module": "@jupyter-widgets/controls"
+          }
+        },
+        "653de804a191499ca46a8eb6b91e7a6b": {
+          "model_module": "@jupyter-widgets/base",
+          "model_name": "LayoutModel",
+          "state": {
+            "_view_name": "LayoutView",
+            "grid_template_rows": null,
+            "right": null,
+            "justify_content": null,
+            "_view_module": "@jupyter-widgets/base",
+            "overflow": null,
+            "_model_module_version": "1.2.0",
+            "_view_count": null,
+            "flex_flow": null,
+            "width": null,
+            "min_width": null,
+            "border": null,
+            "align_items": null,
+            "bottom": null,
+            "_model_module": "@jupyter-widgets/base",
+            "top": null,
+            "grid_column": null,
+            "overflow_y": null,
+            "overflow_x": null,
+            "grid_auto_flow": null,
+            "grid_area": null,
+            "grid_template_columns": null,
+            "flex": null,
+            "_model_name": "LayoutModel",
+            "justify_items": null,
+            "grid_row": null,
+            "max_height": null,
+            "align_content": null,
+            "visibility": null,
+            "align_self": null,
+            "height": null,
+            "min_height": null,
+            "padding": null,
+            "grid_auto_rows": null,
+            "grid_gap": null,
+            "max_width": null,
+            "order": null,
+            "_view_module_version": "1.2.0",
+            "grid_template_areas": null,
+            "object_position": null,
+            "object_fit": null,
+            "grid_auto_columns": null,
+            "margin": null,
+            "display": null,
+            "left": null
+          }
+        },
+        "ee14e110bec24331816d2243cdde5e37": {
+          "model_module": "@jupyter-widgets/controls",
+          "model_name": "VBoxModel",
+          "state": {
+            "_view_name": "VBoxView",
+            "_dom_classes": [],
+            "_model_name": "VBoxModel",
+            "_view_module": "@jupyter-widgets/controls",
+            "_model_module_version": "1.5.0",
+            "_view_count": null,
+            "_view_module_version": "1.5.0",
+            "box_style": "",
+            "layout": "IPY_MODEL_c3fbec3f54b54293a667dfe43c5aed56",
+            "_model_module": "@jupyter-widgets/controls",
+            "children": [
+              "IPY_MODEL_93630fd11e7f47fa9ef23b1057ec1dd9",
+              "IPY_MODEL_aa9410600ce04b5592f4d2a7b4531656"
+            ]
+          }
+        },
+        "c3fbec3f54b54293a667dfe43c5aed56": {
+          "model_module": "@jupyter-widgets/base",
+          "model_name": "LayoutModel",
+          "state": {
+            "_view_name": "LayoutView",
+            "grid_template_rows": null,
+            "right": null,
+            "justify_content": null,
+            "_view_module": "@jupyter-widgets/base",
+            "overflow": null,
+            "_model_module_version": "1.2.0",
+            "_view_count": null,
+            "flex_flow": null,
+            "width": null,
+            "min_width": null,
+            "border": null,
+            "align_items": null,
+            "bottom": null,
+            "_model_module": "@jupyter-widgets/base",
+            "top": null,
+            "grid_column": null,
+            "overflow_y": null,
+            "overflow_x": null,
+            "grid_auto_flow": null,
+            "grid_area": null,
+            "grid_template_columns": null,
+            "flex": null,
+            "_model_name": "LayoutModel",
+            "justify_items": null,
+            "grid_row": null,
+            "max_height": null,
+            "align_content": null,
+            "visibility": null,
+            "align_self": null,
+            "height": null,
+            "min_height": null,
+            "padding": null,
+            "grid_auto_rows": null,
+            "grid_gap": null,
+            "max_width": null,
+            "order": null,
+            "_view_module_version": "1.2.0",
+            "grid_template_areas": null,
+            "object_position": null,
+            "object_fit": null,
+            "grid_auto_columns": null,
+            "margin": null,
+            "display": null,
+            "left": null
+          }
+        },
+        "93630fd11e7f47fa9ef23b1057ec1dd9": {
+          "model_module": "@jupyter-widgets/controls",
+          "model_name": "HTMLModel",
+          "state": {
+            "_view_name": "HTMLView",
+            "style": "IPY_MODEL_af65853e39c04163a6173f438cc0ae71",
+            "_dom_classes": [],
+            "description": "",
+            "_model_name": "HTMLModel",
+            "placeholder": "​",
+            "_view_module": "@jupyter-widgets/controls",
+            "_model_module_version": "1.5.0",
+            "value": "Generating images: 9",
+            "_view_count": null,
+            "_view_module_version": "1.5.0",
+            "description_tooltip": null,
+            "_model_module": "@jupyter-widgets/controls",
+            "layout": "IPY_MODEL_2bb901e39be5406bb1f495030180732b"
+          }
+        },
+        "aa9410600ce04b5592f4d2a7b4531656": {
+          "model_module": "@jupyter-widgets/controls",
+          "model_name": "IntProgressModel",
+          "state": {
+            "_view_name": "ProgressView",
+            "style": "IPY_MODEL_f2b4a35ed75343aab26390a8c26b6e0f",
+            "_dom_classes": [],
+            "description": "",
+            "_model_name": "IntProgressModel",
+            "bar_style": "success",
+            "max": 9,
+            "_view_module": "@jupyter-widgets/controls",
+            "_model_module_version": "1.5.0",
+            "value": 9,
+            "_view_count": null,
+            "_view_module_version": "1.5.0",
+            "orientation": "horizontal",
+            "min": 0,
+            "description_tooltip": null,
+            "_model_module": "@jupyter-widgets/controls",
+            "layout": "IPY_MODEL_df1e1c500f9c42e09fd96d7cb6cca316"
+          }
+        },
+        "af65853e39c04163a6173f438cc0ae71": {
+          "model_module": "@jupyter-widgets/controls",
+          "model_name": "DescriptionStyleModel",
+          "state": {
+            "_view_name": "StyleView",
+            "_model_name": "DescriptionStyleModel",
+            "description_width": "",
+            "_view_module": "@jupyter-widgets/base",
+            "_model_module_version": "1.5.0",
+            "_view_count": null,
+            "_view_module_version": "1.2.0",
+            "_model_module": "@jupyter-widgets/controls"
+          }
+        },
+        "2bb901e39be5406bb1f495030180732b": {
+          "model_module": "@jupyter-widgets/base",
+          "model_name": "LayoutModel",
+          "state": {
+            "_view_name": "LayoutView",
+            "grid_template_rows": null,
+            "right": null,
+            "justify_content": null,
+            "_view_module": "@jupyter-widgets/base",
+            "overflow": null,
+            "_model_module_version": "1.2.0",
+            "_view_count": null,
+            "flex_flow": null,
+            "width": null,
+            "min_width": null,
+            "border": null,
+            "align_items": null,
+            "bottom": null,
+            "_model_module": "@jupyter-widgets/base",
+            "top": null,
+            "grid_column": null,
+            "overflow_y": null,
+            "overflow_x": null,
+            "grid_auto_flow": null,
+            "grid_area": null,
+            "grid_template_columns": null,
+            "flex": null,
+            "_model_name": "LayoutModel",
+            "justify_items": null,
+            "grid_row": null,
+            "max_height": null,
+            "align_content": null,
+            "visibility": null,
+            "align_self": null,
+            "height": null,
+            "min_height": null,
+            "padding": null,
+            "grid_auto_rows": null,
+            "grid_gap": null,
+            "max_width": null,
+            "order": null,
+            "_view_module_version": "1.2.0",
+            "grid_template_areas": null,
+            "object_position": null,
+            "object_fit": null,
+            "grid_auto_columns": null,
+            "margin": null,
+            "display": null,
+            "left": null
+          }
+        },
+        "f2b4a35ed75343aab26390a8c26b6e0f": {
+          "model_module": "@jupyter-widgets/controls",
+          "model_name": "ProgressStyleModel",
+          "state": {
+            "_view_name": "StyleView",
+            "_model_name": "ProgressStyleModel",
+            "description_width": "",
+            "_view_module": "@jupyter-widgets/base",
+            "_model_module_version": "1.5.0",
+            "_view_count": null,
+            "_view_module_version": "1.2.0",
+            "bar_color": null,
+            "_model_module": "@jupyter-widgets/controls"
+          }
+        },
+        "df1e1c500f9c42e09fd96d7cb6cca316": {
+          "model_module": "@jupyter-widgets/base",
+          "model_name": "LayoutModel",
+          "state": {
+            "_view_name": "LayoutView",
+            "grid_template_rows": null,
+            "right": null,
+            "justify_content": null,
+            "_view_module": "@jupyter-widgets/base",
+            "overflow": null,
+            "_model_module_version": "1.2.0",
+            "_view_count": null,
+            "flex_flow": null,
+            "width": null,
+            "min_width": null,
+            "border": null,
+            "align_items": null,
+            "bottom": null,
+            "_model_module": "@jupyter-widgets/base",
+            "top": null,
+            "grid_column": null,
+            "overflow_y": null,
+            "overflow_x": null,
+            "grid_auto_flow": null,
+            "grid_area": null,
+            "grid_template_columns": null,
+            "flex": null,
+            "_model_name": "LayoutModel",
+            "justify_items": null,
+            "grid_row": null,
+            "max_height": null,
+            "align_content": null,
+            "visibility": null,
+            "align_self": null,
+            "height": null,
+            "min_height": null,
+            "padding": null,
+            "grid_auto_rows": null,
+            "grid_gap": null,
+            "max_width": null,
+            "order": null,
+            "_view_module_version": "1.2.0",
+            "grid_template_areas": null,
+            "object_position": null,
+            "object_fit": null,
+            "grid_auto_columns": null,
+            "margin": null,
+            "display": null,
+            "left": null
+          }
+        }
+      }
+    }
+  },
+  "cells": [
+    {
+      "cell_type": "code",
+      "metadata": {
+        "id": "QvYDzQccgMg_"
+      },
+      "source": [
+        "%tensorflow_version 1.x\r\n",
+        "import tensorflow as tf"
+      ],
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "code",
+      "metadata": {
+        "colab": {
+          "base_uri": "https://localhost:8080/"
+        },
+        "id": "y8VaukPJgclY",
+        "outputId": "56ac601b-2cba-427e-c9bb-860d583c1cf6"
+      },
+      "source": [
+        "%cd /content\n",
+        "!rm -rf /content/rasm\n",
+        "!git clone https://github.com/ARBML/rasm\n",
+        "%cd rasm"
+      ],
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "code",
+      "metadata": {
+        "colab": {
+          "base_uri": "https://localhost:8080/"
+        },
+        "id": "BPyug4mhnEEz",
+        "outputId": "ad0b269c-23f8-4376-8830-9f9d0541b6c8"
+      },
+      "source": [
+        "from rasm import Rasm\n",
+        "model = Rasm(mode = 'calligraphy')"
+      ],
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "code",
+      "metadata": {
+        "colab": {
+          "base_uri": "https://localhost:8080/",
+          "height": 919,
+          "referenced_widgets": [
+            "edfc5ae9a3924ee6a04811ab3dec1656",
+            "86ea189295a54dd0aa2a771a95b12f00",
+            "61e1782fb995412fbef8e59c570e8823",
+            "a03932d2929c4225b7735d0831a18566",
+            "c5b6fad0fd114c90ac0f23d91546846f",
+            "7dea244035a34d0f8658e75a12b5cd89",
+            "13d6a5cf5103499ebfb40e2e9c520d27",
+            "653de804a191499ca46a8eb6b91e7a6b"
+          ]
+        },
+        "id": "e5mniebmwJiy",
+        "outputId": "cf409647-73e5-479e-d711-c3f914721ee5"
+      },
+      "source": [
+        "model.generate_randomly()"
+      ],
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "code",
+      "metadata": {
+        "colab": {
+          "base_uri": "https://localhost:8080/",
+          "height": 919,
+          "referenced_widgets": [
+            "ee14e110bec24331816d2243cdde5e37",
+            "c3fbec3f54b54293a667dfe43c5aed56",
+            "93630fd11e7f47fa9ef23b1057ec1dd9",
+            "aa9410600ce04b5592f4d2a7b4531656",
+            "af65853e39c04163a6173f438cc0ae71",
+            "2bb901e39be5406bb1f495030180732b",
+            "f2b4a35ed75343aab26390a8c26b6e0f",
+            "df1e1c500f9c42e09fd96d7cb6cca316"
+          ]
+        },
+        "id": "F2RGfy_9wRFS",
+        "outputId": "2b5fa5af-401f-44b0-f2a6-13da08507396"
+      },
+      "source": [
+        "model.generate_grid()"
+      ],
+      "execution_count": null,
+      "outputs": []
+    }
+  ]
+}

dnnlib/__init__.py

@@ -0,0 +1,24 @@
+# Copyright (c) 2020, NVIDIA CORPORATION.  All rights reserved.
+#
+# NVIDIA CORPORATION and its licensors retain all intellectual property
+# and proprietary rights in and to this software, related documentation
+# and any modifications thereto.  Any use, reproduction, disclosure or
+# distribution of this software and related documentation without an express
+# license agreement from NVIDIA CORPORATION is strictly prohibited.
+
+from . import submission
+
+from .submission.run_context import RunContext
+
+from .submission.submit import SubmitTarget
+from .submission.submit import PathType
+from .submission.submit import SubmitConfig
+from .submission.submit import submit_run
+from .submission.submit import submit_diagnostic
+from .submission.submit import get_path_from_template
+from .submission.submit import convert_path
+from .submission.submit import make_run_dir_path
+
+from .util import EasyDict
+
+submit_config: SubmitConfig = None # Package level variable for SubmitConfig which is only valid when inside the run function.

dnnlib/submission/__init__.py

@@ -0,0 +1,8 @@
+# Copyright (c) 2019, NVIDIA Corporation. All rights reserved.
+#
+# This work is made available under the Nvidia Source Code License-NC.
+# To view a copy of this license, visit
+# https://nvlabs.github.io/stylegan2/license.html
+
+from . import run_context
+from . import submit

dnnlib/submission/internal/__init__.py

@@ -0,0 +1,7 @@
+# Copyright (c) 2019, NVIDIA Corporation. All rights reserved.
+#
+# This work is made available under the Nvidia Source Code License-NC.
+# To view a copy of this license, visit
+# https://nvlabs.github.io/stylegan2/license.html
+
+from . import local

dnnlib/submission/internal/local.py

@@ -0,0 +1,22 @@
+# Copyright (c) 2019, NVIDIA Corporation. All rights reserved.
+#
+# This work is made available under the Nvidia Source Code License-NC.
+# To view a copy of this license, visit
+# https://nvlabs.github.io/stylegan2/license.html
+
+class TargetOptions():
+    def __init__(self):
+        self.do_not_copy_source_files = False
+
+class Target():
+    def __init__(self):
+        pass
+
+    def finalize_submit_config(self, submit_config, host_run_dir):
+        # print ('Local submit ', end='', flush=True)
+        submit_config.run_dir = host_run_dir
+
+    def submit(self, submit_config, host_run_dir):
+        from ..submit import run_wrapper, convert_path
+        # print('- run_dir: %s' % convert_path(submit_config.run_dir), flush=True)
+        return run_wrapper(submit_config)

dnnlib/submission/run_context.py

@@ -0,0 +1,110 @@
+# Copyright (c) 2019, NVIDIA Corporation. All rights reserved.
+#
+# This work is made available under the Nvidia Source Code License-NC.
+# To view a copy of this license, visit
+# https://nvlabs.github.io/stylegan2/license.html
+
+"""Helpers for managing the run/training loop."""
+
+import datetime
+import json
+import os
+import pprint
+import time
+import types
+
+from typing import Any
+
+from . import submit
+
+# Singleton RunContext
+_run_context = None
+
+class RunContext(object):
+    """Helper class for managing the run/training loop.
+
+    The context will hide the implementation details of a basic run/training loop.
+    It will set things up properly, tell if run should be stopped, and then cleans up.
+    User should call update periodically and use should_stop to determine if run should be stopped.
+
+    Args:
+        submit_config: The SubmitConfig that is used for the current run.
+        config_module: (deprecated) The whole config module that is used for the current run.
+    """
+
+    def __init__(self, submit_config: submit.SubmitConfig, config_module: types.ModuleType = None):
+        global _run_context
+        # Only a single RunContext can be alive
+        assert _run_context is None
+        _run_context = self
+        self.submit_config = submit_config
+        self.should_stop_flag = False
+        self.has_closed = False
+        self.start_time = time.time()
+        self.last_update_time = time.time()
+        self.last_update_interval = 0.0
+        self.progress_monitor_file_path = None
+
+        # vestigial config_module support just prints a warning
+        if config_module is not None:
+            print("RunContext.config_module parameter support has been removed.")
+
+        # write out details about the run to a text file
+        self.run_txt_data = {"task_name": submit_config.task_name, "host_name": submit_config.host_name, "start_time": datetime.datetime.now().isoformat(sep=" ")}
+        with open(os.path.join(submit_config.run_dir, "run.txt"), "w") as f:
+            pprint.pprint(self.run_txt_data, stream=f, indent=4, width=200, compact=False)
+
+    def __enter__(self) -> "RunContext":
+        return self
+
+    def __exit__(self, exc_type: Any, exc_value: Any, traceback: Any) -> None:
+        self.close()
+
+    def update(self, loss: Any = 0, cur_epoch: Any = 0, max_epoch: Any = None) -> None:
+        """Do general housekeeping and keep the state of the context up-to-date.
+        Should be called often enough but not in a tight loop."""
+        assert not self.has_closed
+
+        self.last_update_interval = time.time() - self.last_update_time
+        self.last_update_time = time.time()
+
+        if os.path.exists(os.path.join(self.submit_config.run_dir, "abort.txt")):
+            self.should_stop_flag = True
+
+    def should_stop(self) -> bool:
+        """Tell whether a stopping condition has been triggered one way or another."""
+        return self.should_stop_flag
+
+    def get_time_since_start(self) -> float:
+        """How much time has passed since the creation of the context."""
+        return time.time() - self.start_time
+
+    def get_time_since_last_update(self) -> float:
+        """How much time has passed since the last call to update."""
+        return time.time() - self.last_update_time
+
+    def get_last_update_interval(self) -> float:
+        """How much time passed between the previous two calls to update."""
+        return self.last_update_interval
+
+    def close(self) -> None:
+        """Close the context and clean up.
+        Should only be called once."""
+        if not self.has_closed:
+            # update the run.txt with stopping time
+            self.run_txt_data["stop_time"] = datetime.datetime.now().isoformat(sep=" ")
+            with open(os.path.join(self.submit_config.run_dir, "run.txt"), "w") as f:
+                pprint.pprint(self.run_txt_data, stream=f, indent=4, width=200, compact=False)
+            self.has_closed = True
+
+            # detach the global singleton
+            global _run_context
+            if _run_context is self:
+                _run_context = None
+
+    @staticmethod
+    def get():
+        import dnnlib
+        if _run_context is not None:
+            return _run_context
+        return RunContext(dnnlib.submit_config)

dnnlib/submission/submit.py

@@ -0,0 +1,369 @@
+# Copyright (c) 2019, NVIDIA Corporation. All rights reserved.
+#
+# This work is made available under the Nvidia Source Code License-NC.
+# To view a copy of this license, visit
+# https://nvlabs.github.io/stylegan2/license.html
+
+"""Submit a function to be run either locally or in a computing cluster."""
+
+import copy
+import inspect
+import os
+import pathlib
+import pickle
+import platform
+import pprint
+import re
+import shutil
+import sys
+import time
+import traceback
+
+from enum import Enum
+
+from .. import util
+from ..util import EasyDict
+
+from . import internal
+
+class SubmitTarget(Enum):
+    """The target where the function should be run.
+
+    LOCAL: Run it locally.
+    """
+    LOCAL = 1
+    DIAGNOSTIC = 17
+
+
+class PathType(Enum):
+    """Determines in which format should a path be formatted.
+
+    WINDOWS: Format with Windows style.
+    LINUX: Format with Linux/Posix style.
+    AUTO: Use current OS type to select either WINDOWS or LINUX.
+    """
+    WINDOWS = 1
+    LINUX = 2
+    AUTO = 3
+
+
+class PlatformExtras:
+    """A mixed bag of values used by dnnlib heuristics.
+
+    Attributes:
+
+        data_reader_buffer_size: Used by DataReader to size internal shared memory buffers.
+        data_reader_process_count: Number of worker processes to spawn (zero for single thread operation)
+    """
+    def __init__(self):
+        self.data_reader_buffer_size = 1<<30    # 1 GB
+        self.data_reader_process_count = 0      # single threaded default
+
+
+_user_name_override = None
+
+class SubmitConfig(util.EasyDict):
+    """Strongly typed config dict needed to submit runs.
+
+    Attributes:
+        run_dir_root: Path to the run dir root. Can be optionally templated with tags. Needs to always be run through get_path_from_template.
+        run_desc: Description of the run. Will be used in the run dir and task name.
+        run_dir_ignore: List of file patterns used to ignore files when copying files to the run dir.
+        run_dir_extra_files: List of (abs_path, rel_path) tuples of file paths. rel_path root will be the src directory inside the run dir.
+        submit_target: Submit target enum value. Used to select where the run is actually launched.
+        num_gpus: Number of GPUs used/requested for the run.
+        print_info: Whether to print debug information when submitting.
+        local.do_not_copy_source_files: Do not copy source files from the working directory to the run dir.
+        run_id: Automatically populated value during submit.
+        run_name: Automatically populated value during submit.
+        run_dir: Automatically populated value during submit.
+        run_func_name: Automatically populated value during submit.
+        run_func_kwargs: Automatically populated value during submit.
+        user_name: Automatically populated value during submit. Can be set by the user which will then override the automatic value.
+        task_name: Automatically populated value during submit.
+        host_name: Automatically populated value during submit.
+        platform_extras: Automatically populated values during submit.  Used by various dnnlib libraries such as the DataReader class.
+    """
+
+    def __init__(self):
+        super().__init__()
+
+        # run (set these)
+        self.run_dir_root = ""  # should always be passed through get_path_from_template
+        self.run_desc = ""
+        self.run_dir_ignore = ["__pycache__", "*.pyproj", "*.sln", "*.suo", ".cache", ".idea", ".vs", ".vscode", "_cudacache"]
+        self.run_dir_extra_files = []
+
+        # submit (set these)
+        self.submit_target = SubmitTarget.LOCAL
+        self.num_gpus = 1
+        self.print_info = False
+        self.nvprof = False
+        self.local = internal.local.TargetOptions()
+        self.datasets = []
+
+        # (automatically populated)
+        self.run_id = None
+        self.run_name = None
+        self.run_dir = None
+        self.run_func_name = None
+        self.run_func_kwargs = None
+        self.user_name = None
+        self.task_name = None
+        self.host_name = "localhost"
+        self.platform_extras = PlatformExtras()
+
+
+def get_path_from_template(path_template: str, path_type: PathType = PathType.AUTO) -> str:
+    """Replace tags in the given path template and return either Windows or Linux formatted path."""
+    # automatically select path type depending on running OS
+    if path_type == PathType.AUTO:
+        if platform.system() == "Windows":
+            path_type = PathType.WINDOWS
+        elif platform.system() == "Linux":
+            path_type = PathType.LINUX
+        else:
+            raise RuntimeError("Unknown platform")
+
+    path_template = path_template.replace("<USERNAME>", get_user_name())
+
+    # return correctly formatted path
+    if path_type == PathType.WINDOWS:
+        return str(pathlib.PureWindowsPath(path_template))
+    elif path_type == PathType.LINUX:
+        return str(pathlib.PurePosixPath(path_template))
+    else:
+        raise RuntimeError("Unknown platform")
+
+
+def get_template_from_path(path: str) -> str:
+    """Convert a normal path back to its template representation."""
+    path = path.replace("\\", "/")
+    return path
+
+
+def convert_path(path: str, path_type: PathType = PathType.AUTO) -> str:
+    """Convert a normal path to template and the convert it back to a normal path with given path type."""
+    path_template = get_template_from_path(path)
+    path = get_path_from_template(path_template, path_type)
+    return path
+
+
+def set_user_name_override(name: str) -> None:
+    """Set the global username override value."""
+    global _user_name_override
+    _user_name_override = name
+
+
+def get_user_name():
+    """Get the current user name."""
+    if _user_name_override is not None:
+        return _user_name_override
+    elif platform.system() == "Windows":
+        return os.getlogin()
+    elif platform.system() == "Linux":
+        try:
+            import pwd
+            return pwd.getpwuid(os.geteuid()).pw_name
+        except:
+            return "unknown"
+    else:
+        raise RuntimeError("Unknown platform")
+
+
+def make_run_dir_path(*paths):
+    """Make a path/filename that resides under the current submit run_dir.
+
+    Args:
+        *paths: Path components to be passed to os.path.join
+
+    Returns:
+        A file/dirname rooted at submit_config.run_dir.  If there's no
+        submit_config or run_dir, the base directory is the current
+        working directory.
+
+    E.g., `os.path.join(dnnlib.submit_config.run_dir, "output.txt"))`
+    """
+    import dnnlib
+    if (dnnlib.submit_config is None) or (dnnlib.submit_config.run_dir is None):
+        return os.path.join(os.getcwd(), *paths)
+    return os.path.join(dnnlib.submit_config.run_dir, *paths)
+
+
+def _create_run_dir_local(submit_config: SubmitConfig) -> str:
+    """Create a new run dir with increasing ID number at the start."""
+    run_dir_root = get_path_from_template(submit_config.run_dir_root, PathType.AUTO)
+
+    if not os.path.exists(run_dir_root):
+        os.makedirs(run_dir_root)
+
+    submit_config.run_id = _get_next_run_id_local(run_dir_root)
+    submit_config.run_name = "{0:05d}-{1}".format(submit_config.run_id, submit_config.run_desc)
+    run_dir = os.path.join(run_dir_root, submit_config.run_name)
+
+    if os.path.exists(run_dir):
+        raise RuntimeError("The run dir already exists! ({0})".format(run_dir))
+
+    os.makedirs(run_dir)
+
+    return run_dir
+
+
+def _get_next_run_id_local(run_dir_root: str) -> int:
+    """Reads all directory names in a given directory (non-recursive) and returns the next (increasing) run id. Assumes IDs are numbers at the start of the directory names."""
+    dir_names = [d for d in os.listdir(run_dir_root) if os.path.isdir(os.path.join(run_dir_root, d))]
+    r = re.compile("^\\d+")  # match one or more digits at the start of the string
+    run_id = 0
+
+    for dir_name in dir_names:
+        m = r.match(dir_name)
+
+        if m is not None:
+            i = int(m.group())
+            run_id = max(run_id, i + 1)
+
+    return run_id
+
+
+def _populate_run_dir(submit_config: SubmitConfig, run_dir: str) -> None:
+    """Copy all necessary files into the run dir. Assumes that the dir exists, is local, and is writable."""
+    pickle.dump(submit_config, open(os.path.join(run_dir, "submit_config.pkl"), "wb"))
+    with open(os.path.join(run_dir, "submit_config.txt"), "w") as f:
+        pprint.pprint(submit_config, stream=f, indent=4, width=200, compact=False)
+
+    if (submit_config.submit_target == SubmitTarget.LOCAL) and submit_config.local.do_not_copy_source_files:
+        return
+
+    files = []
+
+    run_func_module_dir_path = util.get_module_dir_by_obj_name(submit_config.run_func_name)
+    assert '.' in submit_config.run_func_name
+    for _idx in range(submit_config.run_func_name.count('.') - 1):
+        run_func_module_dir_path = os.path.dirname(run_func_module_dir_path)
+    files += util.list_dir_recursively_with_ignore(run_func_module_dir_path, ignores=submit_config.run_dir_ignore, add_base_to_relative=False)
+
+    dnnlib_module_dir_path = util.get_module_dir_by_obj_name("dnnlib")
+    files += util.list_dir_recursively_with_ignore(dnnlib_module_dir_path, ignores=submit_config.run_dir_ignore, add_base_to_relative=True)
+
+    files += submit_config.run_dir_extra_files
+
+    files = [(f[0], os.path.join(run_dir, "src", f[1])) for f in files]
+    files += [(os.path.join(dnnlib_module_dir_path, "submission", "internal", "run.py"), os.path.join(run_dir, "run.py"))]
+
+    util.copy_files_and_create_dirs(files)
+
+
+
+def run_wrapper(submit_config: SubmitConfig) -> None:
+    """Wrap the actual run function call for handling logging, exceptions, typing, etc."""
+    is_local = submit_config.submit_target == SubmitTarget.LOCAL
+
+    # when running locally, redirect stderr to stdout, log stdout to a file, and force flushing
+    if is_local:
+        logger = util.Logger(file_name=os.path.join(submit_config.run_dir, "log.txt"), file_mode="w", should_flush=True)
+    else:  # when running in a cluster, redirect stderr to stdout, and just force flushing (log writing is handled by run.sh)
+        logger = util.Logger(file_name=None, should_flush=True)
+
+    import dnnlib
+    dnnlib.submit_config = submit_config
+
+    exit_with_errcode = False
+    try:
+        # print("dnnlib: Running {0}() on {1}...".format(submit_config.run_func_name, submit_config.host_name))
+        start_time = time.time()
+
+        run_func_obj = util.get_obj_by_name(submit_config.run_func_name)
+        assert callable(run_func_obj)
+        sig = inspect.signature(run_func_obj)
+        if 'submit_config' in sig.parameters:
+            run_func_obj(submit_config=submit_config, **submit_config.run_func_kwargs)
+        else:
+            run_func_obj(**submit_config.run_func_kwargs)
+
+        # print("dnnlib: Finished {0}() in {1}.".format(submit_config.run_func_name, util.format_time(time.time() - start_time)))
+    except:
+        if is_local:
+            raise
+        else:
+            traceback.print_exc()
+
+            try:
+                log_src = os.path.join(submit_config.run_dir, "log.txt")
+                log_dst = os.path.join(get_path_from_template(submit_config.run_dir_root), "{0}-error.txt".format(submit_config.run_name))
+                shutil.copyfile(log_src, log_dst)
+            except:
+                print("Failing hard, check stack trace")
+
+            # Defer sys.exit(1) to happen after we close the logs and create a _finished.txt
+            exit_with_errcode = True
+    finally:
+        if submit_config.submit_target != SubmitTarget.DIAGNOSTIC:
+            open(os.path.join(submit_config.run_dir, "_finished.txt"), "w").close()
+
+    dnnlib.RunContext.get().close()
+    dnnlib.submit_config = None
+    logger.close()
+
+    # If we hit an error, get out of the script now and signal the error
+    # to whatever process that started this script.
+    if exit_with_errcode:
+        sys.exit(1)
+
+    return submit_config
+
+
+def submit_run(submit_config: SubmitConfig, run_func_name: str, **run_func_kwargs) -> None:
+    """Create a run dir, gather files related to the run, copy files to the run dir, and launch the run in appropriate place."""
+    submit_config = copy.deepcopy(submit_config)
+
+    submit_target = submit_config.submit_target
+    farm = None
+    if submit_target == SubmitTarget.LOCAL:
+        farm = internal.local.Target()
+    assert farm is not None # unknown target
+
+    # Disallow submitting jobs with zero num_gpus.
+    if (submit_config.num_gpus is None) or (submit_config.num_gpus == 0):
+        raise RuntimeError("submit_config.num_gpus must be set to a non-zero value")
+
+    if submit_config.user_name is None:
+        submit_config.user_name = get_user_name()
+
+    submit_config.run_func_name = run_func_name
+    submit_config.run_func_kwargs = run_func_kwargs
+
+    #--------------------------------------------------------------------
+    # Prepare submission by populating the run dir
+    #--------------------------------------------------------------------
+    host_run_dir = _create_run_dir_local(submit_config)
+
+    submit_config.task_name = "{0}-{1:05d}-{2}".format(submit_config.user_name, submit_config.run_id, submit_config.run_desc)
+    docker_valid_name_regex = "^[a-zA-Z0-9][a-zA-Z0-9_.-]+$"
+    if not re.match(docker_valid_name_regex, submit_config.task_name):
+        raise RuntimeError("Invalid task name.  Probable reason: unacceptable characters in your submit_config.run_desc.  Task name must be accepted by the following regex: " + docker_valid_name_regex + ", got " + submit_config.task_name)
+
+    # Farm specific preparations for a submit
+    farm.finalize_submit_config(submit_config, host_run_dir)
+    _populate_run_dir(submit_config, host_run_dir)
+    return farm.submit(submit_config, host_run_dir)
+
+def submit_diagnostic(submit_config: SubmitConfig, run_func_name: str, **run_func_kwargs) -> None:
+    """Launch a run without creating a run directory."""
+    submit_config = copy.deepcopy(submit_config)
+
+    submit_target = submit_config.submit_target
+    farm = None
+    if submit_target == SubmitTarget.LOCAL or submit_target == SubmitTarget.DIAGNOSTIC:
+        farm = internal.local.Target()
+    assert farm is not None # unknown target
+
+    if submit_config.user_name is None:
+        submit_config.user_name = get_user_name()
+
+    submit_config.run_func_name = run_func_name
+    submit_config.run_func_kwargs = run_func_kwargs
+
+    host_run_dir = ""
+    # Farm specific preparations for a submit
+    farm.finalize_submit_config(submit_config, host_run_dir)
+    return farm.submit(submit_config, host_run_dir)

dnnlib/tflib/__init__.py

@@ -0,0 +1,20 @@
+# Copyright (c) 2020, NVIDIA CORPORATION.  All rights reserved.
+#
+# NVIDIA CORPORATION and its licensors retain all intellectual property
+# and proprietary rights in and to this software, related documentation
+# and any modifications thereto.  Any use, reproduction, disclosure or
+# distribution of this software and related documentation without an express
+# license agreement from NVIDIA CORPORATION is strictly prohibited.
+
+from . import autosummary
+from . import network
+from . import optimizer
+from . import tfutil
+from . import custom_ops
+
+from .tfutil import *
+from .network import Network
+
+from .optimizer import Optimizer
+
+from .custom_ops import get_plugin

dnnlib/tflib/autosummary.py

@@ -0,0 +1,193 @@
+# Copyright (c) 2020, NVIDIA CORPORATION.  All rights reserved.
+#
+# NVIDIA CORPORATION and its licensors retain all intellectual property
+# and proprietary rights in and to this software, related documentation
+# and any modifications thereto.  Any use, reproduction, disclosure or
+# distribution of this software and related documentation without an express
+# license agreement from NVIDIA CORPORATION is strictly prohibited.
+
+"""Helper for adding automatically tracked values to Tensorboard.
+
+Autosummary creates an identity op that internally keeps track of the input
+values and automatically shows up in TensorBoard. The reported value
+represents an average over input components. The average is accumulated
+constantly over time and flushed when save_summaries() is called.
+
+Notes:
+- The output tensor must be used as an input for something else in the
+  graph. Otherwise, the autosummary op will not get executed, and the average
+  value will not get accumulated.
+- It is perfectly fine to include autosummaries with the same name in
+  several places throughout the graph, even if they are executed concurrently.
+- It is ok to also pass in a python scalar or numpy array. In this case, it
+  is added to the average immediately.
+"""
+
+from collections import OrderedDict
+import numpy as np
+import tensorflow as tf
+from tensorboard import summary as summary_lib
+from tensorboard.plugins.custom_scalar import layout_pb2
+
+from . import tfutil
+from .tfutil import TfExpression
+from .tfutil import TfExpressionEx
+
+# Enable "Custom scalars" tab in TensorBoard for advanced formatting.
+# Disabled by default to reduce tfevents file size.
+enable_custom_scalars = False
+
+_dtype = tf.float64
+_vars = OrderedDict()  # name => [var, ...]
+_immediate = OrderedDict()  # name => update_op, update_value
+_finalized = False
+_merge_op = None
+
+
+def _create_var(name: str, value_expr: TfExpression) -> TfExpression:
+    """Internal helper for creating autosummary accumulators."""
+    assert not _finalized
+    name_id = name.replace("/", "_")
+    v = tf.cast(value_expr, _dtype)
+
+    if v.shape.is_fully_defined():
+        size = np.prod(v.shape.as_list())
+        size_expr = tf.constant(size, dtype=_dtype)
+    else:
+        size = None
+        size_expr = tf.reduce_prod(tf.cast(tf.shape(v), _dtype))
+
+    if size == 1:
+        if v.shape.ndims != 0:
+            v = tf.reshape(v, [])
+        v = [size_expr, v, tf.square(v)]
+    else:
+        v = [size_expr, tf.reduce_sum(v), tf.reduce_sum(tf.square(v))]
+    v = tf.cond(tf.is_finite(v[1]), lambda: tf.stack(v), lambda: tf.zeros(3, dtype=_dtype))
+
+    with tfutil.absolute_name_scope("Autosummary/" + name_id), tf.control_dependencies(None):
+        var = tf.Variable(tf.zeros(3, dtype=_dtype), trainable=False)  # [sum(1), sum(x), sum(x**2)]
+    update_op = tf.cond(tf.is_variable_initialized(var), lambda: tf.assign_add(var, v), lambda: tf.assign(var, v))
+
+    if name in _vars:
+        _vars[name].append(var)
+    else:
+        _vars[name] = [var]
+    return update_op
+
+
+def autosummary(name: str, value: TfExpressionEx, passthru: TfExpressionEx = None, condition: TfExpressionEx = True) -> TfExpressionEx:
+    """Create a new autosummary.
+
+    Args:
+        name:     Name to use in TensorBoard
+        value:    TensorFlow expression or python value to track
+        passthru: Optionally return this TF node without modifications but tack an autosummary update side-effect to this node.
+
+    Example use of the passthru mechanism:
+
+    n = autosummary('l2loss', loss, passthru=n)
+
+    This is a shorthand for the following code:
+
+    with tf.control_dependencies([autosummary('l2loss', loss)]):
+        n = tf.identity(n)
+    """
+    tfutil.assert_tf_initialized()
+    name_id = name.replace("/", "_")
+
+    if tfutil.is_tf_expression(value):
+        with tf.name_scope("summary_" + name_id), tf.device(value.device):
+            condition = tf.convert_to_tensor(condition, name='condition')
+            update_op = tf.cond(condition, lambda: tf.group(_create_var(name, value)), tf.no_op)
+            with tf.control_dependencies([update_op]):
+                return tf.identity(value if passthru is None else passthru)
+
+    else:  # python scalar or numpy array
+        assert not tfutil.is_tf_expression(passthru)
+        assert not tfutil.is_tf_expression(condition)
+        if condition:
+            if name not in _immediate:
+                with tfutil.absolute_name_scope("Autosummary/" + name_id), tf.device(None), tf.control_dependencies(None):
+                    update_value = tf.placeholder(_dtype)
+                    update_op = _create_var(name, update_value)
+                    _immediate[name] = update_op, update_value
+            update_op, update_value = _immediate[name]
+            tfutil.run(update_op, {update_value: value})
+        return value if passthru is None else passthru
+
+
+def finalize_autosummaries() -> None:
+    """Create the necessary ops to include autosummaries in TensorBoard report.
+    Note: This should be done only once per graph.
+    """
+    global _finalized
+    tfutil.assert_tf_initialized()
+
+    if _finalized:
+        return None
+
+    _finalized = True
+    tfutil.init_uninitialized_vars([var for vars_list in _vars.values() for var in vars_list])
+
+    # Create summary ops.
+    with tf.device(None), tf.control_dependencies(None):
+        for name, vars_list in _vars.items():
+            name_id = name.replace("/", "_")
+            with tfutil.absolute_name_scope("Autosummary/" + name_id):
+                moments = tf.add_n(vars_list)
+                moments /= moments[0]
+                with tf.control_dependencies([moments]):  # read before resetting
+                    reset_ops = [tf.assign(var, tf.zeros(3, dtype=_dtype)) for var in vars_list]
+                    with tf.name_scope(None), tf.control_dependencies(reset_ops):  # reset before reporting
+                        mean = moments[1]
+                        std = tf.sqrt(moments[2] - tf.square(moments[1]))
+                        tf.summary.scalar(name, mean)
+                        if enable_custom_scalars:
+                            tf.summary.scalar("xCustomScalars/" + name + "/margin_lo", mean - std)
+                            tf.summary.scalar("xCustomScalars/" + name + "/margin_hi", mean + std)
+
+    # Setup layout for custom scalars.
+    layout = None
+    if enable_custom_scalars:
+        cat_dict = OrderedDict()
+        for series_name in sorted(_vars.keys()):
+            p = series_name.split("/")
+            cat = p[0] if len(p) >= 2 else ""
+            chart = "/".join(p[1:-1]) if len(p) >= 3 else p[-1]
+            if cat not in cat_dict:
+                cat_dict[cat] = OrderedDict()
+            if chart not in cat_dict[cat]:
+                cat_dict[cat][chart] = []
+            cat_dict[cat][chart].append(series_name)
+        categories = []
+        for cat_name, chart_dict in cat_dict.items():
+            charts = []
+            for chart_name, series_names in chart_dict.items():
+                series = []
+                for series_name in series_names:
+                    series.append(layout_pb2.MarginChartContent.Series(
+                        value=series_name,
+                        lower="xCustomScalars/" + series_name + "/margin_lo",
+                        upper="xCustomScalars/" + series_name + "/margin_hi"))
+                margin = layout_pb2.MarginChartContent(series=series)
+                charts.append(layout_pb2.Chart(title=chart_name, margin=margin))
+            categories.append(layout_pb2.Category(title=cat_name, chart=charts))
+        layout = summary_lib.custom_scalar_pb(layout_pb2.Layout(category=categories))
+    return layout
+
+def save_summaries(file_writer, global_step=None):
+    """Call FileWriter.add_summary() with all summaries in the default graph,
+    automatically finalizing and merging them on the first call.
+    """
+    global _merge_op
+    tfutil.assert_tf_initialized()
+
+    if _merge_op is None:
+        layout = finalize_autosummaries()
+        if layout is not None:
+            file_writer.add_summary(layout)
+        with tf.device(None), tf.control_dependencies(None):
+            _merge_op = tf.summary.merge_all()
+
+    file_writer.add_summary(_merge_op.eval(), global_step)

dnnlib/tflib/custom_ops.py

@@ -0,0 +1,181 @@
+# Copyright (c) 2020, NVIDIA CORPORATION.  All rights reserved.
+#
+# NVIDIA CORPORATION and its licensors retain all intellectual property
+# and proprietary rights in and to this software, related documentation
+# and any modifications thereto.  Any use, reproduction, disclosure or
+# distribution of this software and related documentation without an express
+# license agreement from NVIDIA CORPORATION is strictly prohibited.
+
+"""TensorFlow custom ops builder.
+"""
+
+import glob
+import os
+import re
+import uuid
+import hashlib
+import tempfile
+import shutil
+import tensorflow as tf
+from tensorflow.python.client import device_lib # pylint: disable=no-name-in-module
+
+from .. import util
+
+#----------------------------------------------------------------------------
+# Global options.
+
+cuda_cache_path = None
+cuda_cache_version_tag = 'v1'
+do_not_hash_included_headers = True # Speed up compilation by assuming that headers included by the CUDA code never change.
+verbose = False # Print status messages to stdout.
+
+#----------------------------------------------------------------------------
+# Internal helper funcs.
+
+def _find_compiler_bindir():
+    hostx64_paths = sorted(glob.glob('C:/Program Files (x86)/Microsoft Visual Studio/*/Professional/VC/Tools/MSVC/*/bin/Hostx64/x64'), reverse=True)
+    if hostx64_paths != []:
+        return hostx64_paths[0]
+    hostx64_paths = sorted(glob.glob('C:/Program Files (x86)/Microsoft Visual Studio/*/BuildTools/VC/Tools/MSVC/*/bin/Hostx64/x64'), reverse=True)
+    if hostx64_paths != []:
+        return hostx64_paths[0]
+    hostx64_paths = sorted(glob.glob('C:/Program Files (x86)/Microsoft Visual Studio/*/Community/VC/Tools/MSVC/*/bin/Hostx64/x64'), reverse=True)
+    if hostx64_paths != []:
+        return hostx64_paths[0]
+    vc_bin_dir = 'C:/Program Files (x86)/Microsoft Visual Studio 14.0/vc/bin'
+    if os.path.isdir(vc_bin_dir):
+        return vc_bin_dir
+    return None
+
+def _get_compute_cap(device):
+    caps_str = device.physical_device_desc
+    m = re.search('compute capability: (\\d+).(\\d+)', caps_str)
+    major = m.group(1)
+    minor = m.group(2)
+    return (major, minor)
+
+def _get_cuda_gpu_arch_string():
+    gpus = [x for x in device_lib.list_local_devices() if x.device_type == 'GPU']
+    if len(gpus) == 0:
+        raise RuntimeError('No GPU devices found')
+    (major, minor) = _get_compute_cap(gpus[0])
+    return 'sm_%s%s' % (major, minor)
+
+def _run_cmd(cmd):
+    with os.popen(cmd) as pipe:
+        output = pipe.read()
+        status = pipe.close()
+    if status is not None:
+        raise RuntimeError('NVCC returned an error. See below for full command line and output log:\n\n%s\n\n%s' % (cmd, output))
+
+def _prepare_nvcc_cli(opts):
+    cmd = 'nvcc  --std=c++11 -DNDEBUG ' + opts.strip()
+    cmd += ' --disable-warnings'
+    cmd += ' --include-path "%s"' % tf.sysconfig.get_include()
+    cmd += ' --include-path "%s"' % os.path.join(tf.sysconfig.get_include(), 'external', 'protobuf_archive', 'src')
+    cmd += ' --include-path "%s"' % os.path.join(tf.sysconfig.get_include(), 'external', 'com_google_absl')
+    cmd += ' --include-path "%s"' % os.path.join(tf.sysconfig.get_include(), 'external', 'eigen_archive')
+
+    compiler_bindir = _find_compiler_bindir()
+    if compiler_bindir is None:
+        # Require that _find_compiler_bindir succeeds on Windows.  Allow
+        # nvcc to use whatever is the default on Linux.
+        if os.name == 'nt':
+            raise RuntimeError('Could not find MSVC/GCC/CLANG installation on this computer. Check compiler_bindir_search_path list in "%s".' % __file__)
+    else:
+        cmd += ' --compiler-bindir "%s"' % compiler_bindir
+    cmd += ' 2>&1'
+    return cmd
+
+#----------------------------------------------------------------------------
+# Main entry point.
+
+_plugin_cache = dict()
+
+def get_plugin(cuda_file, extra_nvcc_options=[]):
+    cuda_file_base = os.path.basename(cuda_file)
+    cuda_file_name, cuda_file_ext = os.path.splitext(cuda_file_base)
+
+    # Already in cache?
+    if cuda_file in _plugin_cache:
+        return _plugin_cache[cuda_file]
+
+    # Setup plugin.
+    if verbose:
+        print('Setting up TensorFlow plugin "%s": ' % cuda_file_base, end='', flush=True)
+    try:
+        # Hash CUDA source.
+        md5 = hashlib.md5()
+        with open(cuda_file, 'rb') as f:
+            md5.update(f.read())
+        md5.update(b'\n')
+
+        # Hash headers included by the CUDA code by running it through the preprocessor.
+        if not do_not_hash_included_headers:
+            if verbose:
+                print('Preprocessing... ', end='', flush=True)
+            with tempfile.TemporaryDirectory() as tmp_dir:
+                tmp_file = os.path.join(tmp_dir, cuda_file_name + '_tmp' + cuda_file_ext)
+                _run_cmd(_prepare_nvcc_cli('"%s" --preprocess -o "%s" --keep --keep-dir "%s"' % (cuda_file, tmp_file, tmp_dir)))
+                with open(tmp_file, 'rb') as f:
+                    bad_file_str = ('"' + cuda_file.replace('\\', '/') + '"').encode('utf-8') # __FILE__ in error check macros
+                    good_file_str = ('"' + cuda_file_base + '"').encode('utf-8')
+                    for ln in f:
+                        if not ln.startswith(b'# ') and not ln.startswith(b'#line '): # ignore line number pragmas
+                            ln = ln.replace(bad_file_str, good_file_str)
+                            md5.update(ln)
+                    md5.update(b'\n')
+
+        # Select compiler options.
+        compile_opts = ''
+        if os.name == 'nt':
+            compile_opts += '"%s"' % os.path.join(tf.sysconfig.get_lib(), 'python', '_pywrap_tensorflow_internal.lib')
+        elif os.name == 'posix':
+            compile_opts += f' --compiler-options \'-fPIC\''
+            compile_opts += f' --compiler-options \'{" ".join(tf.sysconfig.get_compile_flags())}\''
+            compile_opts += f' --linker-options \'{" ".join(tf.sysconfig.get_link_flags())}\''
+        else:
+            assert False # not Windows or Linux, w00t?
+        compile_opts += f' --gpu-architecture={_get_cuda_gpu_arch_string()}'
+        compile_opts += ' --use_fast_math'
+        for opt in extra_nvcc_options:
+            compile_opts += ' ' + opt
+        nvcc_cmd = _prepare_nvcc_cli(compile_opts)
+
+        # Hash build configuration.
+        md5.update(('nvcc_cmd: ' + nvcc_cmd).encode('utf-8') + b'\n')
+        md5.update(('tf.VERSION: ' + tf.VERSION).encode('utf-8') + b'\n')
+        md5.update(('cuda_cache_version_tag: ' + cuda_cache_version_tag).encode('utf-8') + b'\n')
+
+        # Compile if not already compiled.
+        cache_dir = util.make_cache_dir_path('tflib-cudacache') if cuda_cache_path is None else cuda_cache_path
+        bin_file_ext = '.dll' if os.name == 'nt' else '.so'
+        bin_file = os.path.join(cache_dir, cuda_file_name + '_' + md5.hexdigest() + bin_file_ext)
+        if not os.path.isfile(bin_file):
+            if verbose:
+                print('Compiling... ', end='', flush=True)
+            with tempfile.TemporaryDirectory() as tmp_dir:
+                tmp_file = os.path.join(tmp_dir, cuda_file_name + '_tmp' + bin_file_ext)
+                _run_cmd(nvcc_cmd + ' "%s" --shared -o "%s" --keep --keep-dir "%s"' % (cuda_file, tmp_file, tmp_dir))
+                os.makedirs(cache_dir, exist_ok=True)
+                intermediate_file = os.path.join(cache_dir, cuda_file_name + '_' + uuid.uuid4().hex + '_tmp' + bin_file_ext)
+                shutil.copyfile(tmp_file, intermediate_file)
+                os.rename(intermediate_file, bin_file) # atomic
+
+        # Load.
+        if verbose:
+            print('Loading... ', end='', flush=True)
+        plugin = tf.load_op_library(bin_file)
+
+        # Add to cache.
+        _plugin_cache[cuda_file] = plugin
+        if verbose:
+            print('Done.', flush=True)
+        return plugin
+
+    except:
+        if verbose:
+            print('Failed!', flush=True)
+        raise
+
+#----------------------------------------------------------------------------

dnnlib/tflib/network.py

@@ -0,0 +1,825 @@
+# Copyright (c) 2020, NVIDIA CORPORATION.  All rights reserved.
+#
+# NVIDIA CORPORATION and its licensors retain all intellectual property
+# and proprietary rights in and to this software, related documentation
+# and any modifications thereto.  Any use, reproduction, disclosure or
+# distribution of this software and related documentation without an express
+# license agreement from NVIDIA CORPORATION is strictly prohibited.
+
+"""Helper for managing networks."""
+
+import types
+import inspect
+import re
+import uuid
+import sys
+import copy
+import numpy as np
+import tensorflow as tf
+
+from collections import OrderedDict
+from typing import Any, List, Tuple, Union, Callable
+
+from . import tfutil
+from .. import util
+
+from .tfutil import TfExpression, TfExpressionEx
+
+# pylint: disable=protected-access
+# pylint: disable=attribute-defined-outside-init
+# pylint: disable=too-many-public-methods
+
+_import_handlers = []  # Custom import handlers for dealing with legacy data in pickle import.
+_import_module_src = dict()  # Source code for temporary modules created during pickle import.
+
+
+def import_handler(handler_func):
+    """Function decorator for declaring custom import handlers."""
+    _import_handlers.append(handler_func)
+    return handler_func
+
+
+class Network:
+    """Generic network abstraction.
+
+    Acts as a convenience wrapper for a parameterized network construction
+    function, providing several utility methods and convenient access to
+    the inputs/outputs/weights.
+
+    Network objects can be safely pickled and unpickled for long-term
+    archival purposes. The pickling works reliably as long as the underlying
+    network construction function is defined in a standalone Python module
+    that has no side effects or application-specific imports.
+
+    Args:
+        name: Network name. Used to select TensorFlow name and variable scopes. Defaults to build func name if None.
+        func_name: Fully qualified name of the underlying network construction function, or a top-level function object.
+        static_kwargs: Keyword arguments to be passed in to the network construction function.
+    """
+
+    def __init__(self, name: str = None, func_name: Any = None, **static_kwargs):
+        # Locate the user-specified build function.
+        assert isinstance(func_name, str) or util.is_top_level_function(func_name)
+        if util.is_top_level_function(func_name):
+            func_name = util.get_top_level_function_name(func_name)
+        module, func_name = util.get_module_from_obj_name(func_name)
+        func = util.get_obj_from_module(module, func_name)
+
+        # Dig up source code for the module containing the build function.
+        module_src = _import_module_src.get(module, None)
+        if module_src is None:
+            module_src = inspect.getsource(module)
+
+        # Initialize fields.
+        self._init_fields(name=(name or func_name), static_kwargs=static_kwargs, build_func=func, build_func_name=func_name, build_module_src=module_src)
+
+    def _init_fields(self, name: str, static_kwargs: dict, build_func: Callable, build_func_name: str, build_module_src: str) -> None:
+        tfutil.assert_tf_initialized()
+        assert isinstance(name, str)
+        assert len(name) >= 1
+        assert re.fullmatch(r"[A-Za-z0-9_.\\-]*", name)
+        assert isinstance(static_kwargs, dict)
+        assert util.is_pickleable(static_kwargs)
+        assert callable(build_func)
+        assert isinstance(build_func_name, str)
+        assert isinstance(build_module_src, str)
+
+        # Choose TensorFlow name scope.
+        with tf.name_scope(None):
+            scope = tf.get_default_graph().unique_name(name, mark_as_used=True)
+
+        # Query current TensorFlow device.
+        with tfutil.absolute_name_scope(scope), tf.control_dependencies(None):
+            device = tf.no_op(name="_QueryDevice").device
+
+        # Immutable state.
+        self._name                  = name
+        self._scope                 = scope
+        self._device                = device
+        self._static_kwargs         = util.EasyDict(copy.deepcopy(static_kwargs))
+        self._build_func            = build_func
+        self._build_func_name       = build_func_name
+        self._build_module_src      = build_module_src
+
+        # State before _init_graph().
+        self._var_inits             = dict()    # var_name => initial_value, set to None by _init_graph()
+        self._all_inits_known       = False     # Do we know for sure that _var_inits covers all the variables?
+        self._components            = None      # subnet_name => Network, None if the components are not known yet
+
+        # Initialized by _init_graph().
+        self._input_templates       = None
+        self._output_templates      = None
+        self._own_vars              = None
+
+        # Cached values initialized the respective methods.
+        self._input_shapes          = None
+        self._output_shapes         = None
+        self._input_names           = None
+        self._output_names          = None
+        self._vars                  = None
+        self._trainables            = None
+        self._var_global_to_local   = None
+        self._run_cache             = dict()
+        self.epochs = tf.Variable(0., dtype=tf.float32, name='epochs')
+
+    def _init_graph(self) -> None:
+        assert self._var_inits is not None
+        assert self._input_templates is None
+        assert self._output_templates is None
+        assert self._own_vars is None
+
+        # Initialize components.
+        if self._components is None:
+            self._components = util.EasyDict()
+
+        # Choose build func kwargs.
+        build_kwargs = dict(self.static_kwargs)
+        build_kwargs["is_template_graph"] = True
+        build_kwargs["components"] = self._components
+
+        # Override scope and device, and ignore surrounding control dependencies.
+        with tfutil.absolute_variable_scope(self.scope, reuse=False), tfutil.absolute_name_scope(self.scope), tf.device(self.device), tf.control_dependencies(None):
+            assert tf.get_variable_scope().name == self.scope
+            assert tf.get_default_graph().get_name_scope() == self.scope
+
+            # Create input templates.
+            self._input_templates = []
+            for param in inspect.signature(self._build_func).parameters.values():
+                if param.kind == param.POSITIONAL_OR_KEYWORD and param.default is param.empty:
+                    self._input_templates.append(tf.placeholder(tf.float32, name=param.name))
+
+            # Call build func.
+            out_expr = self._build_func(*self._input_templates, **build_kwargs)
+
+        # Collect output templates and variables.
+        assert tfutil.is_tf_expression(out_expr) or isinstance(out_expr, tuple)
+        self._output_templates = [out_expr] if tfutil.is_tf_expression(out_expr) else list(out_expr)
+        self._own_vars = OrderedDict((var.name[len(self.scope) + 1:].split(":")[0], var) for var in tf.global_variables(self.scope + "/"))
+
+        # Check for errors.
+        if len(self._input_templates) == 0:
+            raise ValueError("Network build func did not list any inputs.")
+        if len(self._output_templates) == 0:
+            raise ValueError("Network build func did not return any outputs.")
+        if any(not tfutil.is_tf_expression(t) for t in self._output_templates):
+            raise ValueError("Network outputs must be TensorFlow expressions.")
+        if any(t.shape.ndims is None for t in self._input_templates):
+            raise ValueError("Network input shapes not defined. Please call x.set_shape() for each input.")
+        if any(t.shape.ndims is None for t in self._output_templates):
+            raise ValueError("Network output shapes not defined. Please call x.set_shape() where applicable.")
+        if any(not isinstance(comp, Network) for comp in self._components.values()):
+            raise ValueError("Components of a Network must be Networks themselves.")
+        if len(self._components) != len(set(comp.name for comp in self._components.values())):
+            raise ValueError("Components of a Network must have unique names.")
+
+        # Initialize variables.
+        if len(self._var_inits):
+            tfutil.set_vars({self._get_vars()[name]: value for name, value in self._var_inits.items() if name in self._get_vars()})
+        remaining_inits = [var.initializer for name, var in self._own_vars.items() if name not in self._var_inits]
+        if self._all_inits_known:
+            assert len(remaining_inits) == 0
+        else:
+            tfutil.run(remaining_inits)
+        self._var_inits = None
+
+    @property
+    def name(self):
+        """User-specified name string."""
+        return self._name
+
+    @property
+    def scope(self):
+        """Unique TensorFlow scope containing template graph and variables, derived from the user-specified name."""
+        return self._scope
+
+    @property
+    def device(self):
+        """Name of the TensorFlow device that the weights of this network reside on. Determined by the current device at construction time."""
+        return self._device
+
+    @property
+    def static_kwargs(self):
+        """EasyDict of arguments passed to the user-supplied build func."""
+        return copy.deepcopy(self._static_kwargs)
+
+    @property
+    def components(self):
+        """EasyDict of sub-networks created by the build func."""
+        return copy.copy(self._get_components())
+
+    def _get_components(self):
+        if self._components is None:
+            self._init_graph()
+            assert self._components is not None
+        return self._components
+
+    @property
+    def input_shapes(self):
+        """List of input tensor shapes, including minibatch dimension."""
+        if self._input_shapes is None:
+            self._input_shapes = [t.shape.as_list() for t in self.input_templates]
+        return copy.deepcopy(self._input_shapes)
+
+    @property
+    def output_shapes(self):
+        """List of output tensor shapes, including minibatch dimension."""
+        if self._output_shapes is None:
+            self._output_shapes = [t.shape.as_list() for t in self.output_templates]
+        return copy.deepcopy(self._output_shapes)
+
+    @property
+    def input_shape(self):
+        """Short-hand for input_shapes[0]."""
+        return self.input_shapes[0]
+
+    @property
+    def output_shape(self):
+        """Short-hand for output_shapes[0]."""
+        return self.output_shapes[0]
+
+    @property
+    def num_inputs(self):
+        """Number of input tensors."""
+        return len(self.input_shapes)
+
+    @property
+    def num_outputs(self):
+        """Number of output tensors."""
+        return len(self.output_shapes)
+
+    @property
+    def input_names(self):
+        """Name string for each input."""
+        if self._input_names is None:
+            self._input_names = [t.name.split("/")[-1].split(":")[0] for t in self.input_templates]
+        return copy.copy(self._input_names)
+
+    @property
+    def output_names(self):
+        """Name string for each output."""
+        if self._output_names is None:
+            self._output_names = [t.name.split("/")[-1].split(":")[0] for t in self.output_templates]
+        return copy.copy(self._output_names)
+
+    @property
+    def input_templates(self):
+        """Input placeholders in the template graph."""
+        if self._input_templates is None:
+            self._init_graph()
+            assert self._input_templates is not None
+        return copy.copy(self._input_templates)
+
+    @property
+    def output_templates(self):
+        """Output tensors in the template graph."""
+        if self._output_templates is None:
+            self._init_graph()
+            assert self._output_templates is not None
+        return copy.copy(self._output_templates)
+
+    @property
+    def own_vars(self):
+        """Variables defined by this network (local_name => var), excluding sub-networks."""
+        return copy.copy(self._get_own_vars())
+
+    def _get_own_vars(self):
+        if self._own_vars is None:
+            self._init_graph()
+            assert self._own_vars is not None
+        return self._own_vars
+
+    @property
+    def vars(self):
+        """All variables (local_name => var)."""
+        return copy.copy(self._get_vars())
+
+    def _get_vars(self):
+        if self._vars is None:
+            self._vars = OrderedDict(self._get_own_vars())
+            for comp in self._get_components().values():
+                self._vars.update((comp.name + "/" + name, var) for name, var in comp._get_vars().items())
+        return self._vars
+
+    @property
+    def trainables(self):
+        """All trainable variables (local_name => var)."""
+        return copy.copy(self._get_trainables())
+
+    def _get_trainables(self):
+        if self._trainables is None:
+            self._trainables = OrderedDict((name, var) for name, var in self.vars.items() if var.trainable)
+        return self._trainables
+
+    @property
+    def var_global_to_local(self):
+        """Mapping from variable global names to local names."""
+        return copy.copy(self._get_var_global_to_local())
+
+    def _get_var_global_to_local(self):
+        if self._var_global_to_local is None:
+            self._var_global_to_local = OrderedDict((var.name.split(":")[0], name) for name, var in self.vars.items())
+        return self._var_global_to_local
+
+    def reset_own_vars(self) -> None:
+        """Re-initialize all variables of this network, excluding sub-networks."""
+        if self._var_inits is None or self._components is None:
+            tfutil.run([var.initializer for var in self._get_own_vars().values()])
+        else:
+            self._var_inits.clear()
+            self._all_inits_known = False
+
+    def reset_vars(self) -> None:
+        """Re-initialize all variables of this network, including sub-networks."""
+        if self._var_inits is None:
+            tfutil.run([var.initializer for var in self._get_vars().values()])
+        else:
+            self._var_inits.clear()
+            self._all_inits_known = False
+            if self._components is not None:
+                for comp in self._components.values():
+                    comp.reset_vars()
+
+    def reset_trainables(self) -> None:
+        """Re-initialize all trainable variables of this network, including sub-networks."""
+        tfutil.run([var.initializer for var in self._get_trainables().values()])
+
+    def get_output_for(self, *in_expr: TfExpression, return_as_list: bool = False, **dynamic_kwargs) -> Union[TfExpression, List[TfExpression]]:
+        """Construct TensorFlow expression(s) for the output(s) of this network, given the input expression(s).
+        The graph is placed on the current TensorFlow device."""
+        assert len(in_expr) == self.num_inputs
+        assert not all(expr is None for expr in in_expr)
+        self._get_vars()  # ensure that all variables have been created
+
+        # Choose build func kwargs.
+        build_kwargs = dict(self.static_kwargs)
+        build_kwargs.update(dynamic_kwargs)
+        build_kwargs["is_template_graph"] = False
+        build_kwargs["components"] = self._components
+
+        # Build TensorFlow graph to evaluate the network.
+        with tfutil.absolute_variable_scope(self.scope, reuse=True), tf.name_scope(self.name):
+            assert tf.get_variable_scope().name == self.scope
+            valid_inputs = [expr for expr in in_expr if expr is not None]
+            final_inputs = []
+            for expr, name, shape in zip(in_expr, self.input_names, self.input_shapes):
+                if expr is not None:
+                    expr = tf.identity(expr, name=name)
+                else:
+                    expr = tf.zeros([tf.shape(valid_inputs[0])[0]] + shape[1:], name=name)
+                final_inputs.append(expr)
+            out_expr = self._build_func(*final_inputs, **build_kwargs)
+
+        # Propagate input shapes back to the user-specified expressions.
+        for expr, final in zip(in_expr, final_inputs):
+            if isinstance(expr, tf.Tensor):
+                expr.set_shape(final.shape)
+
+        # Express outputs in the desired format.
+        assert tfutil.is_tf_expression(out_expr) or isinstance(out_expr, tuple)
+        if return_as_list:
+            out_expr = [out_expr] if tfutil.is_tf_expression(out_expr) else list(out_expr)
+        return out_expr
+
+    def get_var_local_name(self, var_or_global_name: Union[TfExpression, str]) -> str:
+        """Get the local name of a given variable, without any surrounding name scopes."""
+        assert tfutil.is_tf_expression(var_or_global_name) or isinstance(var_or_global_name, str)
+        global_name = var_or_global_name if isinstance(var_or_global_name, str) else var_or_global_name.name
+        return self._get_var_global_to_local()[global_name]
+
+    def find_var(self, var_or_local_name: Union[TfExpression, str]) -> TfExpression:
+        """Find variable by local or global name."""
+        assert tfutil.is_tf_expression(var_or_local_name) or isinstance(var_or_local_name, str)
+        return self._get_vars()[var_or_local_name] if isinstance(var_or_local_name, str) else var_or_local_name
+
+    def get_var(self, var_or_local_name: Union[TfExpression, str]) -> np.ndarray:
+        """Get the value of a given variable as NumPy array.
+        Note: This method is very inefficient -- prefer to use tflib.run(list_of_vars) whenever possible."""
+        return self.find_var(var_or_local_name).eval()
+
+    def set_var(self, var_or_local_name: Union[TfExpression, str], new_value: Union[int, float, np.ndarray]) -> None:
+        """Set the value of a given variable based on the given NumPy array.
+        Note: This method is very inefficient -- prefer to use tflib.set_vars() whenever possible."""
+        tfutil.set_vars({self.find_var(var_or_local_name): new_value})
+
+    def __getstate__(self) -> dict:
+        """Pickle export."""
+        state = dict()
+        state["version"]            = 5
+        state["name"]               = self.name
+        state["static_kwargs"]      = dict(self.static_kwargs)
+        state["components"]         = dict(self.components)
+        state["build_module_src"]   = self._build_module_src
+        state["build_func_name"]    = self._build_func_name
+        state["variables"]          = list(zip(self._get_own_vars().keys(), tfutil.run(list(self._get_own_vars().values()))))
+        state["input_shapes"]       = self.input_shapes
+        state["output_shapes"]      = self.output_shapes
+        state["input_names"]        = self.input_names
+        state["output_names"]       = self.output_names
+        return state
+
+    def __setstate__(self, state: dict) -> None:
+        """Pickle import."""
+
+        # Execute custom import handlers.
+        for handler in _import_handlers:
+            state = handler(state)
+
+        # Get basic fields.
+        assert state["version"] in [2, 3, 4, 5]
+        name = state["name"]
+        static_kwargs = state["static_kwargs"]
+        build_module_src = state["build_module_src"]
+        build_func_name = state["build_func_name"]
+
+        # Create temporary module from the imported source code.
+        module_name = "_tflib_network_import_" + uuid.uuid4().hex
+        module = types.ModuleType(module_name)
+        sys.modules[module_name] = module
+        _import_module_src[module] = build_module_src
+        exec(build_module_src, module.__dict__) # pylint: disable=exec-used
+        build_func = util.get_obj_from_module(module, build_func_name)
+
+        # Initialize fields.
+        self._init_fields(name=name, static_kwargs=static_kwargs, build_func=build_func, build_func_name=build_func_name, build_module_src=build_module_src)
+        self._var_inits.update(copy.deepcopy(state["variables"]))
+        self._all_inits_known   = True
+        self._components        = util.EasyDict(state.get("components", {}))
+        self._input_shapes      = copy.deepcopy(state.get("input_shapes", None))
+        self._output_shapes     = copy.deepcopy(state.get("output_shapes", None))
+        self._input_names       = copy.deepcopy(state.get("input_names", None))
+        self._output_names      = copy.deepcopy(state.get("output_names", None))
+
+    def clone(self, name: str = None, **new_static_kwargs) -> "Network":
+        """Create a clone of this network with its own copy of the variables."""
+        static_kwargs = dict(self.static_kwargs)
+        static_kwargs.update(new_static_kwargs)
+        net = object.__new__(Network)
+        net._init_fields(name=(name or self.name), static_kwargs=static_kwargs, build_func=self._build_func, build_func_name=self._build_func_name, build_module_src=self._build_module_src)
+        net.copy_vars_from(self)
+        return net
+
+    def copy_own_vars_from(self, src_net: "Network") -> None:
+        """Copy the values of all variables from the given network, excluding sub-networks."""
+
+        # Source has unknown variables or unknown components => init now.
+        if (src_net._var_inits is not None and not src_net._all_inits_known) or src_net._components is None:
+            src_net._get_vars()
+
+       # Both networks are inited => copy directly.
+        if src_net._var_inits is None and self._var_inits is None:
+            names = [name for name in self._get_own_vars().keys() if name in src_net._get_own_vars()]
+            tfutil.set_vars(tfutil.run({self._get_vars()[name]: src_net._get_vars()[name] for name in names}))
+            return
+
+        # Read from source.
+        if src_net._var_inits is None:
+            value_dict = tfutil.run(src_net._get_own_vars())
+        else:
+            value_dict = src_net._var_inits
+
+        # Write to destination.
+        if self._var_inits is None:
+            tfutil.set_vars({self._get_vars()[name]: value for name, value in value_dict.items() if name in self._get_vars()})
+        else:
+            self._var_inits.update(value_dict)
+
+    def copy_vars_from(self, src_net: "Network") -> None:
+        """Copy the values of all variables from the given network, including sub-networks."""
+
+        # Source has unknown variables or unknown components => init now.
+        if (src_net._var_inits is not None and not src_net._all_inits_known) or src_net._components is None:
+            src_net._get_vars()
+
+        # Source is inited, but destination components have not been created yet => set as initial values.
+        if src_net._var_inits is None and self._components is None:
+            self._var_inits.update(tfutil.run(src_net._get_vars()))
+            return
+
+        # Destination has unknown components => init now.
+        if self._components is None:
+            self._get_vars()
+
+        # Both networks are inited => copy directly.
+        if src_net._var_inits is None and self._var_inits is None:
+            names = [name for name in self._get_vars().keys() if name in src_net._get_vars()]
+            tfutil.set_vars(tfutil.run({self._get_vars()[name]: src_net._get_vars()[name] for name in names}))
+            return
+
+        # Copy recursively, component by component.
+        self.copy_own_vars_from(src_net)
+        for name, src_comp in src_net._components.items():
+            if name in self._components:
+                self._components[name].copy_vars_from(src_comp)
+
+    def copy_trainables_from(self, src_net: "Network") -> None:
+        """Copy the values of all trainable variables from the given network, including sub-networks."""
+        names = [name for name in self._get_trainables().keys() if name in src_net._get_trainables()]
+        tfutil.set_vars(tfutil.run({self._get_vars()[name]: src_net._get_vars()[name] for name in names}))
+
+    def copy_compatible_trainables_from(self, src_net: "Network") -> None:
+        """Copy the compatible values of all trainable variables from the given network, including sub-networks"""
+        names = []
+        for name in self.trainables.keys():
+            if name not in src_net.trainables:
+                print("Not restoring (not present):     {}".format(name))
+            elif self.trainables[name].shape != src_net.trainables[name].shape:
+                print("Not restoring (different shape): {}".format(name))
+
+            if name in src_net.trainables and self.trainables[name].shape == src_net.trainables[name].shape:
+                names.append(name)
+
+        tfutil.set_vars(tfutil.run({self.vars[name]: src_net.vars[name] for name in names}))
+
+    def apply_swa(self, src_net, epoch):
+        """Perform stochastic weight averaging on the compatible values of all trainable variables from the given network, including sub-networks"""
+        names = []
+        for name in self.trainables.keys():
+            if name not in src_net.trainables:
+                print("Not restoring (not present):     {}".format(name))
+            elif self.trainables[name].shape != src_net.trainables[name].shape:
+                print("Not restoring (different shape): {}".format(name))
+
+            if name in src_net.trainables and self.trainables[name].shape == src_net.trainables[name].shape:
+                names.append(name)
+
+        scale_new_data = 1.0 - 1.0 / (epoch + 1)
+        scale_moving_average = (1.0 - scale_new_data)
+        tfutil.set_vars(tfutil.run({self.vars[name]: (src_net.vars[name] * scale_new_data + self.vars[name] * scale_moving_average) for name in names}))
+
+    def convert(self, new_func_name: str, new_name: str = None, **new_static_kwargs) -> "Network":
+        """Create new network with the given parameters, and copy all variables from this network."""
+        if new_name is None:
+            new_name = self.name
+        static_kwargs = dict(self.static_kwargs)
+        static_kwargs.update(new_static_kwargs)
+        net = Network(name=new_name, func_name=new_func_name, **static_kwargs)
+        net.copy_vars_from(self)
+        return net
+
+    def setup_as_moving_average_of(self, src_net: "Network", beta: TfExpressionEx = 0.99, beta_nontrainable: TfExpressionEx = 0.0) -> tf.Operation:
+        """Construct a TensorFlow op that updates the variables of this network
+        to be slightly closer to those of the given network."""
+        with tfutil.absolute_name_scope(self.scope + "/_MovingAvg"):
+            ops = []
+            for name, var in self._get_vars().items():
+                if name in src_net._get_vars():
+                    cur_beta = beta if var.trainable else beta_nontrainable
+                    new_value = tfutil.lerp(src_net._get_vars()[name], var, cur_beta)
+                    ops.append(var.assign(new_value))
+            return tf.group(*ops)
+
+    def update_epochs(self, epochs: TfExpressionEx = 0) -> tf.Operation:
+        """Construct a TensorFlow op that updates the epoch counter of this network."""
+        with tfutil.absolute_name_scope(self.scope + "/_Epochs"):
+            op = self.epochs.assign(epochs)
+            return op
+
+    def run(self,
+            *in_arrays: Tuple[Union[np.ndarray, None], ...],
+            input_transform: dict = None,
+            output_transform: dict = None,
+            return_as_list: bool = False,
+            print_progress: bool = False,
+            minibatch_size: int = None,
+            num_gpus: int = 1,
+            assume_frozen: bool = False,
+            custom_inputs: Any = None,
+            **dynamic_kwargs) -> Union[np.ndarray, Tuple[np.ndarray, ...], List[np.ndarray]]:
+        """Run this network for the given NumPy array(s), and return the output(s) as NumPy array(s).
+
+        Args:
+            input_transform:    A dict specifying a custom transformation to be applied to the input tensor(s) before evaluating the network.
+                                The dict must contain a 'func' field that points to a top-level function. The function is called with the input
+                                TensorFlow expression(s) as positional arguments. Any remaining fields of the dict will be passed in as kwargs.
+            output_transform:   A dict specifying a custom transformation to be applied to the output tensor(s) after evaluating the network.
+                                The dict must contain a 'func' field that points to a top-level function. The function is called with the output
+                                TensorFlow expression(s) as positional arguments. Any remaining fields of the dict will be passed in as kwargs.
+            return_as_list:     True = return a list of NumPy arrays, False = return a single NumPy array, or a tuple if there are multiple outputs.
+            print_progress:     Print progress to the console? Useful for very large input arrays.
+            minibatch_size:     Maximum minibatch size to use, None = disable batching.
+            num_gpus:           Number of GPUs to use.
+            assume_frozen:      Improve multi-GPU performance by assuming that the trainable parameters will remain changed between calls.
+            custom_inputs:      Allow to use another tensor as input instead of default placeholders.
+            dynamic_kwargs:     Additional keyword arguments to be passed into the network build function.
+        """
+        assert len(in_arrays) == self.num_inputs
+        assert not all(arr is None for arr in in_arrays)
+        assert input_transform is None or util.is_top_level_function(input_transform["func"])
+        assert output_transform is None or util.is_top_level_function(output_transform["func"])
+        output_transform, dynamic_kwargs = _handle_legacy_output_transforms(output_transform, dynamic_kwargs)
+        num_items = in_arrays[0].shape[0]
+        if minibatch_size is None:
+            minibatch_size = num_items
+
+        # Construct unique hash key from all arguments that affect the TensorFlow graph.
+        key = dict(input_transform=input_transform, output_transform=output_transform, num_gpus=num_gpus, assume_frozen=assume_frozen, dynamic_kwargs=dynamic_kwargs)
+        def unwind_key(obj):
+            if isinstance(obj, dict):
+                return [(key, unwind_key(value)) for key, value in sorted(obj.items())]
+            if callable(obj):
+                return util.get_top_level_function_name(obj)
+            return obj
+        key = repr(unwind_key(key))
+
+        # Build graph.
+        if key not in self._run_cache:
+            with tfutil.absolute_name_scope(self.scope + "/_Run"), tf.control_dependencies(None):
+                if custom_inputs is not None:
+                    with tf.device("/gpu:0"):
+                        in_expr = [input_builder(name) for input_builder, name in zip(custom_inputs, self.input_names)]
+                        in_split = list(zip(*[tf.split(x, num_gpus) for x in in_expr]))
+                else:
+                    with tf.device("/cpu:0"):
+                        in_expr = [tf.placeholder(tf.float32, name=name) for name in self.input_names]
+                        in_split = list(zip(*[tf.split(x, num_gpus) for x in in_expr]))
+
+                out_split = []
+                for gpu in range(num_gpus):
+                    with tf.device(self.device if num_gpus == 1 else "/gpu:%d" % gpu):
+                        net_gpu = self.clone() if assume_frozen else self
+                        in_gpu = in_split[gpu]
+
+                        if input_transform is not None:
+                            in_kwargs = dict(input_transform)
+                            in_gpu = in_kwargs.pop("func")(*in_gpu, **in_kwargs)
+                            in_gpu = [in_gpu] if tfutil.is_tf_expression(in_gpu) else list(in_gpu)
+
+                        assert len(in_gpu) == self.num_inputs
+                        out_gpu = net_gpu.get_output_for(*in_gpu, return_as_list=True, **dynamic_kwargs)
+
+                        if output_transform is not None:
+                            out_kwargs = dict(output_transform)
+                            out_gpu = out_kwargs.pop("func")(*out_gpu, **out_kwargs)
+                            out_gpu = [out_gpu] if tfutil.is_tf_expression(out_gpu) else list(out_gpu)
+
+                        assert len(out_gpu) == self.num_outputs
+                        out_split.append(out_gpu)
+
+                with tf.device("/cpu:0"):
+                    out_expr = [tf.concat(outputs, axis=0) for outputs in zip(*out_split)]
+                    self._run_cache[key] = in_expr, out_expr
+
+        # Run minibatches.
+        in_expr, out_expr = self._run_cache[key]
+        out_arrays = [np.empty([num_items] + expr.shape.as_list()[1:], expr.dtype.name) for expr in out_expr]
+
+        for mb_begin in range(0, num_items, minibatch_size):
+            if print_progress:
+                print("\r%d / %d" % (mb_begin, num_items), end="")
+
+            mb_end = min(mb_begin + minibatch_size, num_items)
+            mb_num = mb_end - mb_begin
+            mb_in = [src[mb_begin : mb_end] if src is not None else np.zeros([mb_num] + shape[1:]) for src, shape in zip(in_arrays, self.input_shapes)]
+            mb_out = tf.get_default_session().run(out_expr, dict(zip(in_expr, mb_in)))
+
+            for dst, src in zip(out_arrays, mb_out):
+                dst[mb_begin: mb_end] = src
+
+        # Done.
+        if print_progress:
+            print("\r%d / %d" % (num_items, num_items))
+
+        if not return_as_list:
+            out_arrays = out_arrays[0] if len(out_arrays) == 1 else tuple(out_arrays)
+        return out_arrays
+
+    def list_ops(self) -> List[TfExpression]:
+        _ = self.output_templates  # ensure that the template graph has been created
+        include_prefix = self.scope + "/"
+        exclude_prefix = include_prefix + "_"
+        ops = tf.get_default_graph().get_operations()
+        ops = [op for op in ops if op.name.startswith(include_prefix)]
+        ops = [op for op in ops if not op.name.startswith(exclude_prefix)]
+        return ops
+
+    def list_layers(self) -> List[Tuple[str, TfExpression, List[TfExpression]]]:
+        """Returns a list of (layer_name, output_expr, trainable_vars) tuples corresponding to
+        individual layers of the network. Mainly intended to be used for reporting."""
+        layers = []
+
+        def recurse(scope, parent_ops, parent_vars, level):
+            if len(parent_ops) == 0 and len(parent_vars) == 0:
+                return
+
+            # Ignore specific patterns.
+            if any(p in scope for p in ["/Shape", "/strided_slice", "/Cast", "/concat", "/Assign"]):
+                return
+
+            # Filter ops and vars by scope.
+            global_prefix = scope + "/"
+            local_prefix = global_prefix[len(self.scope) + 1:]
+            cur_ops = [op for op in parent_ops if op.name.startswith(global_prefix) or op.name == global_prefix[:-1]]
+            cur_vars = [(name, var) for name, var in parent_vars if name.startswith(local_prefix) or name == local_prefix[:-1]]
+            if not cur_ops and not cur_vars:
+                return
+
+            # Filter out all ops related to variables.
+            for var in [op for op in cur_ops if op.type.startswith("Variable")]:
+                var_prefix = var.name + "/"
+                cur_ops = [op for op in cur_ops if not op.name.startswith(var_prefix)]
+
+            # Scope does not contain ops as immediate children => recurse deeper.
+            contains_direct_ops = any("/" not in op.name[len(global_prefix):] and op.type not in ["Identity", "Cast", "Transpose"] for op in cur_ops)
+            if (level == 0 or not contains_direct_ops) and (len(cur_ops) != 0 or len(cur_vars) != 0):
+                visited = set()
+                for rel_name in [op.name[len(global_prefix):] for op in cur_ops] + [name[len(local_prefix):] for name, _var in cur_vars]:
+                    token = rel_name.split("/")[0]
+                    if token not in visited:
+                        recurse(global_prefix + token, cur_ops, cur_vars, level + 1)
+                        visited.add(token)
+                return
+
+            # Report layer.
+            layer_name = scope[len(self.scope) + 1:]
+            layer_output = cur_ops[-1].outputs[0] if cur_ops else cur_vars[-1][1]
+            layer_trainables = [var for _name, var in cur_vars if var.trainable]
+            layers.append((layer_name, layer_output, layer_trainables))
+
+        recurse(self.scope, self.list_ops(), list(self._get_vars().items()), 0)
+        return layers
+    
+    def print_layers(self, title: str = None, hide_layers_with_no_params: bool = False) -> None:
+        """Print a summary table of the network structure."""
+        rows = [[title if title is not None else self.name, "Params", "OutputShape", "WeightShape"]]
+        rows += [["---"] * 4]
+        total_params = 0
+
+        for layer_name, layer_output, layer_trainables in self.list_layers():
+            num_params = sum(int(np.prod(var.shape.as_list())) for var in layer_trainables)
+            weights = [var for var in layer_trainables if var.name.endswith("/weight:0")]
+            weights.sort(key=lambda x: len(x.name))
+            if len(weights) == 0 and len(layer_trainables) == 1:
+                weights = layer_trainables
+            total_params += num_params
+
+            if not hide_layers_with_no_params or num_params != 0:
+                num_params_str = str(num_params) if num_params > 0 else "-"
+                output_shape_str = str(layer_output.shape)
+                weight_shape_str = str(weights[0].shape) if len(weights) >= 1 else "-"
+                rows += [[layer_name, num_params_str, output_shape_str, weight_shape_str]]
+
+        rows += [["---"] * 4]
+        rows += [["Total", str(total_params), "", ""]]
+
+        widths = [max(len(cell) for cell in column) for column in zip(*rows)]
+        print()
+        for row in rows:
+            print("  ".join(cell + " " * (width - len(cell)) for cell, width in zip(row, widths)))
+        print()
+
+    def setup_weight_histograms(self, title: str = None) -> None:
+        """Construct summary ops to include histograms of all trainable parameters in TensorBoard."""
+        if title is None:
+            title = self.name
+
+        with tf.name_scope(None), tf.device(None), tf.control_dependencies(None):
+            for local_name, var in self._get_trainables().items():
+                if "/" in local_name:
+                    p = local_name.split("/")
+                    name = title + "_" + p[-1] + "/" + "_".join(p[:-1])
+                else:
+                    name = title + "_toplevel/" + local_name
+
+                tf.summary.histogram(name, var)
+
+#----------------------------------------------------------------------------
+# Backwards-compatible emulation of legacy output transformation in Network.run().
+
+_print_legacy_warning = True
+
+def _handle_legacy_output_transforms(output_transform, dynamic_kwargs):
+    global _print_legacy_warning
+    legacy_kwargs = ["out_mul", "out_add", "out_shrink", "out_dtype"]
+    if not any(kwarg in dynamic_kwargs for kwarg in legacy_kwargs):
+        return output_transform, dynamic_kwargs
+
+    if _print_legacy_warning:
+        _print_legacy_warning = False
+        print()
+        print("WARNING: Old-style output transformations in Network.run() are deprecated.")
+        print("Consider using 'output_transform=dict(func=tflib.convert_images_to_uint8)'")
+        print("instead of 'out_mul=127.5, out_add=127.5, out_dtype=np.uint8'.")
+        print()
+    assert output_transform is None
+
+    new_kwargs = dict(dynamic_kwargs)
+    new_transform = {kwarg: new_kwargs.pop(kwarg) for kwarg in legacy_kwargs if kwarg in dynamic_kwargs}
+    new_transform["func"] = _legacy_output_transform_func
+    return new_transform, new_kwargs
+
+def _legacy_output_transform_func(*expr, out_mul=1.0, out_add=0.0, out_shrink=1, out_dtype=None):
+    if out_mul != 1.0:
+        expr = [x * out_mul for x in expr]
+
+    if out_add != 0.0:
+        expr = [x + out_add for x in expr]
+
+    if out_shrink > 1:
+        ksize = [1, 1, out_shrink, out_shrink]
+        expr = [tf.nn.avg_pool(x, ksize=ksize, strides=ksize, padding="VALID", data_format="NCHW") for x in expr]
+
+    if out_dtype is not None:
+        if tf.as_dtype(out_dtype).is_integer:
+            expr = [tf.round(x) for x in expr]
+        expr = [tf.saturate_cast(x, out_dtype) for x in expr]
+    return expr

dnnlib/tflib/ops/__init__.py

@@ -0,0 +1,9 @@
+# Copyright (c) 2020, NVIDIA CORPORATION.  All rights reserved.
+#
+# NVIDIA CORPORATION and its licensors retain all intellectual property
+# and proprietary rights in and to this software, related documentation
+# and any modifications thereto.  Any use, reproduction, disclosure or
+# distribution of this software and related documentation without an express
+# license agreement from NVIDIA CORPORATION is strictly prohibited.
+
+# empty

dnnlib/tflib/ops/fused_bias_act.cu

@@ -0,0 +1,220 @@
+// Copyright (c) 2020, NVIDIA CORPORATION.  All rights reserved.
+//
+// NVIDIA CORPORATION and its licensors retain all intellectual property
+// and proprietary rights in and to this software, related documentation
+// and any modifications thereto.  Any use, reproduction, disclosure or
+// distribution of this software and related documentation without an express
+// license agreement from NVIDIA CORPORATION is strictly prohibited.
+
+#define EIGEN_USE_GPU
+#define __CUDA_INCLUDE_COMPILER_INTERNAL_HEADERS__
+#include "tensorflow/core/framework/op.h"
+#include "tensorflow/core/framework/op_kernel.h"
+#include "tensorflow/core/framework/shape_inference.h"
+#include <stdio.h>
+
+using namespace tensorflow;
+using namespace tensorflow::shape_inference;
+
+#define OP_CHECK_CUDA_ERROR(CTX, CUDA_CALL) do { cudaError_t err = CUDA_CALL; OP_REQUIRES(CTX, err == cudaSuccess, errors::Internal(cudaGetErrorName(err))); } while (false)
+
+//------------------------------------------------------------------------
+// CUDA kernel.
+
+template <class T>
+struct FusedBiasActKernelParams
+{
+    const T*    x;      // [sizeX]
+    const T*    b;      // [sizeB] or NULL
+    const T*    xref;   // [sizeX] or NULL
+    const T*    yref;   // [sizeX] or NULL
+    T*          y;      // [sizeX]
+
+    int         grad;
+    int         axis;
+    int         act;
+    float       alpha;
+    float       gain;
+    float       clamp;
+
+    int         sizeX;
+    int         sizeB;
+    int         stepB;
+    int         loopX;
+};
+
+template <class T>
+static __global__ void FusedBiasActKernel(const FusedBiasActKernelParams<T> p)
+{
+    const float expRange        = 80.0f;
+    const float halfExpRange    = 40.0f;
+    const float seluScale       = 1.0507009873554804934193349852946f;
+    const float seluAlpha       = 1.6732632423543772848170429916717f;
+
+    // Loop over elements.
+    int xi = blockIdx.x * p.loopX * blockDim.x + threadIdx.x;
+    for (int loopIdx = 0; loopIdx < p.loopX && xi < p.sizeX; loopIdx++, xi += blockDim.x)
+    {
+        // Load and apply bias.
+        float x = (float)p.x[xi];
+        if (p.b)
+            x += (float)p.b[(xi / p.stepB) % p.sizeB];
+        float xref = (p.xref) ? (float)p.xref[xi] : 0.0f;
+        float yref = (p.yref) ? (float)p.yref[xi] : 0.0f;
+        float yy = (p.gain != 0.0f) ? yref / p.gain : 0.0f;
+
+        // Evaluate activation func.
+        float y;
+        switch (p.act * 10 + p.grad)
+        {
+            // linear
+            default:
+            case 10: y = x; break;
+            case 11: y = x; break;
+            case 12: y = 0.0f; break;
+
+            // relu
+            case 20: y = (x > 0.0f) ? x : 0.0f; break;
+            case 21: y = (yy > 0.0f) ? x : 0.0f; break;
+            case 22: y = 0.0f; break;
+
+            // lrelu
+            case 30: y = (x > 0.0f) ? x : x * p.alpha; break;
+            case 31: y = (yy > 0.0f) ? x : x * p.alpha; break;
+            case 32: y = 0.0f; break;
+
+            // tanh
+            case 40: { float c = expf(x); float d = 1.0f / c; y = (x < -expRange) ? -1.0f : (x > expRange) ? 1.0f : (c - d) / (c + d); } break;
+            case 41: y = x * (1.0f - yy * yy); break;
+            case 42: y = x * (1.0f - yy * yy) * (-2.0f * yy); break;
+
+            // sigmoid
+            case 50: y = (x < -expRange) ? 0.0f : 1.0f / (expf(-x) + 1.0f); break;
+            case 51: y = x * yy * (1.0f - yy); break;
+            case 52: y = x * yy * (1.0f - yy) * (1.0f - 2.0f * yy); break;
+
+            // elu
+            case 60: y = (x >= 0.0f) ? x : expf(x) - 1.0f; break;
+            case 61: y = (yy >= 0.0f) ? x : x * (yy + 1.0f); break;
+            case 62: y = (yy >= 0.0f) ? 0.0f : x * (yy + 1.0f); break;
+
+            // selu
+            case 70: y = (x >= 0.0f) ? seluScale * x : (seluScale * seluAlpha) * (expf(x) - 1.0f); break;
+            case 71: y = (yy >= 0.0f) ? x * seluScale : x * (yy + seluScale * seluAlpha); break;
+            case 72: y = (yy >= 0.0f) ? 0.0f : x * (yy + seluScale * seluAlpha); break;
+
+            // softplus
+            case 80: y = (x > expRange) ? x : logf(expf(x) + 1.0f); break;
+            case 81: y = x * (1.0f - expf(-yy)); break;
+            case 82: { float c = expf(-yy); y = x * c * (1.0f - c); } break;
+
+            // swish
+            case 90: y = (x < -expRange) ? 0.0f : x / (expf(-x) + 1.0f); break;
+            case 91:
+            case 92:
+                {
+                    float c = expf(xref);
+                    float d = c + 1.0f;
+                    if (p.grad == 1)
+                        y = (xref > halfExpRange) ? x : x * c * (xref + d) / (d * d);
+                    else
+                        y = (xref > halfExpRange) ? 0.0f : x * c * (xref * (2.0f - d) + 2.0f * d) / (d * d * d);
+                    yref = (xref < -expRange) ? 0.0f : xref / (expf(-xref) + 1.0f) * p.gain;
+                }
+                break;
+        }
+
+        // Apply gain.
+        y *= p.gain;
+
+        // Clamp.
+        if (p.clamp >= 0.0f)
+        {
+            if (p.grad == 0)
+                y = (fabsf(y) < p.clamp) ? y : (y >= 0.0f) ? p.clamp : -p.clamp;
+            else
+                y = (fabsf(yref) < p.clamp) ? y : 0.0f;
+        }
+
+        // Store.
+        p.y[xi] = (T)y;
+    }
+}
+
+//------------------------------------------------------------------------
+// TensorFlow op.
+
+template <class T>
+struct FusedBiasActOp : public OpKernel
+{
+    FusedBiasActKernelParams<T> m_attribs;
+
+    FusedBiasActOp(OpKernelConstruction* ctx) : OpKernel(ctx)
+    {
+        memset(&m_attribs, 0, sizeof(m_attribs));
+        OP_REQUIRES_OK(ctx, ctx->GetAttr("grad",    &m_attribs.grad));
+        OP_REQUIRES_OK(ctx, ctx->GetAttr("axis",    &m_attribs.axis));
+        OP_REQUIRES_OK(ctx, ctx->GetAttr("act",     &m_attribs.act));
+        OP_REQUIRES_OK(ctx, ctx->GetAttr("alpha",   &m_attribs.alpha));
+        OP_REQUIRES_OK(ctx, ctx->GetAttr("gain",    &m_attribs.gain));
+        OP_REQUIRES_OK(ctx, ctx->GetAttr("clamp",   &m_attribs.clamp));
+        OP_REQUIRES(ctx, m_attribs.grad >= 0, errors::InvalidArgument("grad must be non-negative"));
+        OP_REQUIRES(ctx, m_attribs.axis >= 0, errors::InvalidArgument("axis must be non-negative"));
+        OP_REQUIRES(ctx, m_attribs.act >= 0, errors::InvalidArgument("act must be non-negative"));
+    }
+
+    void Compute(OpKernelContext* ctx)
+    {
+        FusedBiasActKernelParams<T> p = m_attribs;
+        cudaStream_t stream = ctx->eigen_device<Eigen::GpuDevice>().stream();
+
+        const Tensor& x     = ctx->input(0); // [...]
+        const Tensor& b     = ctx->input(1); // [sizeB] or [0]
+        const Tensor& xref  = ctx->input(2); // x.shape or [0]
+        const Tensor& yref  = ctx->input(3); // x.shape or [0]
+        p.x = x.flat<T>().data();
+        p.b = (b.NumElements()) ? b.flat<T>().data() : NULL;
+        p.xref = (xref.NumElements()) ? xref.flat<T>().data() : NULL;
+        p.yref = (yref.NumElements()) ? yref.flat<T>().data() : NULL;
+        OP_REQUIRES(ctx, b.NumElements() == 0 || m_attribs.axis < x.dims(), errors::InvalidArgument("axis out of bounds"));
+        OP_REQUIRES(ctx, b.dims() == 1, errors::InvalidArgument("b must have rank 1"));
+        OP_REQUIRES(ctx, b.NumElements() == 0 || b.NumElements() == x.dim_size(m_attribs.axis), errors::InvalidArgument("b has wrong number of elements"));
+        OP_REQUIRES(ctx, xref.NumElements() == 0 || xref.NumElements() == x.NumElements(), errors::InvalidArgument("xref has wrong number of elements"));
+        OP_REQUIRES(ctx, yref.NumElements() == 0 || yref.NumElements() == x.NumElements(), errors::InvalidArgument("yref has wrong number of elements"));
+        OP_REQUIRES(ctx, x.NumElements() <= kint32max, errors::InvalidArgument("x is too large"));
+
+        p.sizeX = (int)x.NumElements();
+        p.sizeB = (int)b.NumElements();
+        p.stepB = 1;
+        for (int i = m_attribs.axis + 1; i < x.dims(); i++)
+            p.stepB *= (int)x.dim_size(i);
+
+        Tensor* y = NULL; // x.shape
+        OP_REQUIRES_OK(ctx, ctx->allocate_output(0, x.shape(), &y));
+        p.y = y->flat<T>().data();
+
+        p.loopX = 4;
+        int blockSize = 4 * 32;
+        int gridSize = (p.sizeX - 1) / (p.loopX * blockSize) + 1;
+        void* args[] = {&p};
+        OP_CHECK_CUDA_ERROR(ctx, cudaLaunchKernel((void*)FusedBiasActKernel<T>, gridSize, blockSize, args, 0, stream));
+    }
+};
+
+REGISTER_OP("FusedBiasAct")
+    .Input      ("x: T")
+    .Input      ("b: T")
+    .Input      ("xref: T")
+    .Input      ("yref: T")
+    .Output     ("y: T")
+    .Attr       ("T: {float, half}")
+    .Attr       ("grad: int = 0")
+    .Attr       ("axis: int = 1")
+    .Attr       ("act: int = 0")
+    .Attr       ("alpha: float = 0.0")
+    .Attr       ("gain: float = 1.0")
+    .Attr       ("clamp: float = -1.0");
+REGISTER_KERNEL_BUILDER(Name("FusedBiasAct").Device(DEVICE_GPU).TypeConstraint<float>("T"), FusedBiasActOp<float>);
+REGISTER_KERNEL_BUILDER(Name("FusedBiasAct").Device(DEVICE_GPU).TypeConstraint<Eigen::half>("T"), FusedBiasActOp<Eigen::half>);
+
+//------------------------------------------------------------------------

dnnlib/tflib/ops/fused_bias_act.py

@@ -0,0 +1,211 @@
+# Copyright (c) 2020, NVIDIA CORPORATION.  All rights reserved.
+#
+# NVIDIA CORPORATION and its licensors retain all intellectual property
+# and proprietary rights in and to this software, related documentation
+# and any modifications thereto.  Any use, reproduction, disclosure or
+# distribution of this software and related documentation without an express
+# license agreement from NVIDIA CORPORATION is strictly prohibited.
+
+"""Custom TensorFlow ops for efficient bias and activation."""
+
+import os
+import numpy as np
+import tensorflow as tf
+from .. import custom_ops
+from ...util import EasyDict
+
+def _get_plugin():
+    return custom_ops.get_plugin(os.path.splitext(__file__)[0] + '.cu')
+
+#----------------------------------------------------------------------------
+
+activation_funcs = {
+    'linear':   EasyDict(func=lambda x, **_:        x,                          def_alpha=None, def_gain=1.0,           cuda_idx=1, ref='y', zero_2nd_grad=True),
+    'relu':     EasyDict(func=lambda x, **_:        tf.nn.relu(x),              def_alpha=None, def_gain=np.sqrt(2),    cuda_idx=2, ref='y', zero_2nd_grad=True),
+    'lrelu':    EasyDict(func=lambda x, alpha, **_: tf.nn.leaky_relu(x, alpha), def_alpha=0.2,  def_gain=np.sqrt(2),    cuda_idx=3, ref='y', zero_2nd_grad=True),
+    'tanh':     EasyDict(func=lambda x, **_:        tf.nn.tanh(x),              def_alpha=None, def_gain=1.0,           cuda_idx=4, ref='y', zero_2nd_grad=False),
+    'sigmoid':  EasyDict(func=lambda x, **_:        tf.nn.sigmoid(x),           def_alpha=None, def_gain=1.0,           cuda_idx=5, ref='y', zero_2nd_grad=False),
+    'elu':      EasyDict(func=lambda x, **_:        tf.nn.elu(x),               def_alpha=None, def_gain=1.0,           cuda_idx=6, ref='y', zero_2nd_grad=False),
+    'selu':     EasyDict(func=lambda x, **_:        tf.nn.selu(x),              def_alpha=None, def_gain=1.0,           cuda_idx=7, ref='y', zero_2nd_grad=False),
+    'softplus': EasyDict(func=lambda x, **_:        tf.nn.softplus(x),          def_alpha=None, def_gain=1.0,           cuda_idx=8, ref='y', zero_2nd_grad=False),
+    'swish':    EasyDict(func=lambda x, **_:        tf.nn.sigmoid(x) * x,       def_alpha=None, def_gain=np.sqrt(2),    cuda_idx=9, ref='x', zero_2nd_grad=False),
+}
+
+#----------------------------------------------------------------------------
+
+def fused_bias_act(x, b=None, axis=1, act='linear', alpha=None, gain=None, clamp=None, impl='cuda'):
+    r"""Fused bias and activation function.
+
+    Adds bias `b` to activation tensor `x`, evaluates activation function `act`,
+    and scales the result by `gain`. Each of the steps is optional. In most cases,
+    the fused op is considerably more efficient than performing the same calculation
+    using standard TensorFlow ops. It supports first and second order gradients,
+    but not third order gradients.
+
+    Args:
+        x:      Input activation tensor. Can have any shape, but if `b` is defined, the
+                dimension corresponding to `axis`, as well as the rank, must be known.
+        b:      Bias vector, or `None` to disable. Must be a 1D tensor of the same type
+                as `x`. The shape must be known, and it must match the dimension of `x`
+                corresponding to `axis`.
+        axis:   The dimension in `x` corresponding to the elements of `b`.
+                The value of `axis` is ignored if `b` is not specified.
+        act:    Name of the activation function to evaluate, or `"linear"` to disable.
+                Can be e.g. `"relu"`, `"lrelu"`, `"tanh"`, `"sigmoid"`, `"swish"`, etc.
+                See `activation_funcs` for a full list. `None` is not allowed.
+        alpha:  Shape parameter for the activation function, or `None` to use the default.
+        gain:   Scaling factor for the output tensor, or `None` to use default.
+                See `activation_funcs` for the default scaling of each activation function.
+                If unsure, consider specifying `1.0`.
+        clamp:  Clamp the output values to `[-clamp, +clamp]`, or `None` to disable
+                the clamping (default).
+        impl:   Name of the implementation to use. Can be `"ref"` or `"cuda"` (default).
+
+    Returns:
+        Tensor of the same shape and datatype as `x`.
+    """
+
+    impl_dict = {
+        'ref':  _fused_bias_act_ref,
+        'cuda': _fused_bias_act_cuda,
+    }
+    return impl_dict[impl](x=x, b=b, axis=axis, act=act, alpha=alpha, gain=gain, clamp=clamp)
+
+#----------------------------------------------------------------------------
+
+def _fused_bias_act_ref(x, b, axis, act, alpha, gain, clamp):
+    """Slow reference implementation of `fused_bias_act()` using standard TensorFlow ops."""
+
+    # Validate arguments.
+    x = tf.convert_to_tensor(x)
+    b = tf.convert_to_tensor(b) if b is not None else tf.constant([], dtype=x.dtype)
+    act_spec = activation_funcs[act]
+    assert b.shape.rank == 1 and (b.shape[0] == 0 or b.shape[0] == x.shape[axis])
+    assert b.shape[0] == 0 or 0 <= axis < x.shape.rank
+    if alpha is None:
+        alpha = act_spec.def_alpha
+    if gain is None:
+        gain = act_spec.def_gain
+
+    # Add bias.
+    if b.shape[0] != 0:
+        x += tf.reshape(b, [-1 if i == axis else 1 for i in range(x.shape.rank)])
+
+    # Evaluate activation function.
+    x = act_spec.func(x, alpha=alpha)
+
+    # Scale by gain.
+    if gain != 1:
+        x *= gain
+
+    # Clamp.
+    if clamp is not None:
+        clamp = np.asarray(clamp, dtype=x.dtype.name)
+        assert clamp.shape == () and clamp >= 0
+        x = tf.clip_by_value(x, -clamp, clamp)
+    return x
+
+#----------------------------------------------------------------------------
+
+def _fused_bias_act_cuda(x, b, axis, act, alpha, gain, clamp):
+    """Fast CUDA implementation of `fused_bias_act()` using custom ops."""
+
+    # Validate arguments.
+    x = tf.convert_to_tensor(x)
+    empty_tensor = tf.constant([], dtype=x.dtype)
+    b = tf.convert_to_tensor(b) if b is not None else empty_tensor
+    act_spec = activation_funcs[act]
+    assert b.shape.rank == 1 and (b.shape[0] == 0 or b.shape[0] == x.shape[axis])
+    assert b.shape[0] == 0 or 0 <= axis < x.shape.rank
+    if alpha is None:
+        alpha = act_spec.def_alpha
+    if gain is None:
+        gain = act_spec.def_gain
+
+    # Special cases.
+    if act == 'linear' and b is None and gain == 1.0:
+        return x
+    if act_spec.cuda_idx is None:
+        return _fused_bias_act_ref(x=x, b=b, axis=axis, act=act, alpha=alpha, gain=gain, clamp=clamp)
+
+    # CUDA op.
+    cuda_op = _get_plugin().fused_bias_act
+    cuda_kwargs = dict(axis=int(axis), act=int(act_spec.cuda_idx), gain=float(gain))
+    if alpha is not None:
+        cuda_kwargs['alpha'] = float(alpha)
+    if clamp is not None:
+        clamp = np.asarray(clamp, dtype=x.dtype.name)
+        assert clamp.shape == () and clamp >= 0
+        cuda_kwargs['clamp'] = float(clamp.astype(np.float32))
+    def ref(tensor, name):
+        return tensor if act_spec.ref == name else empty_tensor
+
+    # Forward pass: y = func(x, b).
+    def func_y(x, b):
+        y = cuda_op(x=x, b=b, xref=empty_tensor, yref=empty_tensor, grad=0, **cuda_kwargs)
+        y.set_shape(x.shape)
+        return y
+
+    # Backward pass: dx, db = grad(dy, x, y)
+    def grad_dx(dy, x, y):
+        dx = cuda_op(x=dy, b=empty_tensor, xref=ref(x,'x'), yref=ref(y,'y'), grad=1, **cuda_kwargs)
+        dx.set_shape(x.shape)
+        return dx
+    def grad_db(dx):
+        if b.shape[0] == 0:
+            return empty_tensor
+        db = dx
+        if axis < x.shape.rank - 1:
+            db = tf.reduce_sum(db, list(range(axis + 1, x.shape.rank)))
+        if axis > 0:
+            db = tf.reduce_sum(db, list(range(axis)))
+        db.set_shape(b.shape)
+        return db
+
+    # Second order gradients: d_dy, d_x = grad2(d_dx, d_db, x, y)
+    def grad2_d_dy(d_dx, d_db, x, y):
+        d_dy = cuda_op(x=d_dx, b=d_db, xref=ref(x,'x'), yref=ref(y,'y'), grad=1, **cuda_kwargs)
+        d_dy.set_shape(x.shape)
+        return d_dy
+    def grad2_d_x(d_dx, d_db, x, y):
+        d_x = cuda_op(x=d_dx, b=d_db, xref=ref(x,'x'), yref=ref(y,'y'), grad=2, **cuda_kwargs)
+        d_x.set_shape(x.shape)
+        return d_x
+
+    # Fast version for piecewise-linear activation funcs.
+    @tf.custom_gradient
+    def func_zero_2nd_grad(x, b):
+        y = func_y(x, b)
+        @tf.custom_gradient
+        def grad(dy):
+            dx = grad_dx(dy, x, y)
+            db = grad_db(dx)
+            def grad2(d_dx, d_db):
+                d_dy = grad2_d_dy(d_dx, d_db, x, y)
+                return d_dy
+            return (dx, db), grad2
+        return y, grad
+
+    # Slow version for general activation funcs.
+    @tf.custom_gradient
+    def func_nonzero_2nd_grad(x, b):
+        y = func_y(x, b)
+        def grad_wrap(dy):
+            @tf.custom_gradient
+            def grad_impl(dy, x):
+                dx = grad_dx(dy, x, y)
+                db = grad_db(dx)
+                def grad2(d_dx, d_db):
+                    d_dy = grad2_d_dy(d_dx, d_db, x, y)
+                    d_x = grad2_d_x(d_dx, d_db, x, y)
+                    return d_dy, d_x
+                return (dx, db), grad2
+            return grad_impl(dy, x)
+        return y, grad_wrap
+
+    # Which version to use?
+    if act_spec.zero_2nd_grad:
+        return func_zero_2nd_grad(x, b)
+    return func_nonzero_2nd_grad(x, b)
+
+#----------------------------------------------------------------------------

dnnlib/tflib/ops/upfirdn_2d.cu

@@ -0,0 +1,359 @@
+// Copyright (c) 2020, NVIDIA CORPORATION.  All rights reserved.
+//
+// NVIDIA CORPORATION and its licensors retain all intellectual property
+// and proprietary rights in and to this software, related documentation
+// and any modifications thereto.  Any use, reproduction, disclosure or
+// distribution of this software and related documentation without an express
+// license agreement from NVIDIA CORPORATION is strictly prohibited.
+
+#define EIGEN_USE_GPU
+#define __CUDA_INCLUDE_COMPILER_INTERNAL_HEADERS__
+#include "tensorflow/core/framework/op.h"
+#include "tensorflow/core/framework/op_kernel.h"
+#include "tensorflow/core/framework/shape_inference.h"
+#include <stdio.h>
+
+using namespace tensorflow;
+using namespace tensorflow::shape_inference;
+
+//------------------------------------------------------------------------
+// Helpers.
+
+#define OP_CHECK_CUDA_ERROR(CTX, CUDA_CALL) do { cudaError_t err = CUDA_CALL; OP_REQUIRES(CTX, err == cudaSuccess, errors::Internal(cudaGetErrorName(err))); } while (false)
+
+static __host__ __device__ __forceinline__ int floorDiv(int a, int b)
+{
+    int t = 1 - a / b;
+    return (a + t * b) / b - t;
+}
+
+//------------------------------------------------------------------------
+// CUDA kernel params.
+
+template <class T>
+struct UpFirDn2DKernelParams
+{
+    const T*    x;          // [majorDim, inH, inW, minorDim]
+    const T*    k;          // [kernelH, kernelW]
+    T*          y;          // [majorDim, outH, outW, minorDim]
+
+    int         upx;
+    int         upy;
+    int         downx;
+    int         downy;
+    int         padx0;
+    int         padx1;
+    int         pady0;
+    int         pady1;
+
+    int         majorDim;
+    int         inH;
+    int         inW;
+    int         minorDim;
+    int         kernelH;
+    int         kernelW;
+    int         outH;
+    int         outW;
+    int         loopMajor;
+    int         loopX;
+};
+
+//------------------------------------------------------------------------
+// General CUDA implementation for large filter kernels.
+
+template <class T>
+static __global__ void UpFirDn2DKernel_large(const UpFirDn2DKernelParams<T> p)
+{
+    // Calculate thread index.
+    int minorIdx = blockIdx.x * blockDim.x + threadIdx.x;
+    int outY = minorIdx / p.minorDim;
+    minorIdx -= outY * p.minorDim;
+    int outXBase = blockIdx.y * p.loopX * blockDim.y + threadIdx.y;
+    int majorIdxBase = blockIdx.z * p.loopMajor;
+    if (outXBase >= p.outW || outY >= p.outH || majorIdxBase >= p.majorDim)
+        return;
+
+    // Setup Y receptive field.
+    int midY = outY * p.downy + p.upy - 1 - p.pady0;
+    int inY = min(max(floorDiv(midY, p.upy), 0), p.inH);
+    int h = min(max(floorDiv(midY + p.kernelH, p.upy), 0), p.inH) - inY;
+    int kernelY = midY + p.kernelH - (inY + 1) * p.upy;
+
+    // Loop over majorDim and outX.
+    for (int loopMajor = 0, majorIdx = majorIdxBase; loopMajor < p.loopMajor && majorIdx < p.majorDim; loopMajor++, majorIdx++)
+    for (int loopX = 0, outX = outXBase; loopX < p.loopX && outX < p.outW; loopX++, outX += blockDim.y)
+    {
+        // Setup X receptive field.
+        int midX = outX * p.downx + p.upx - 1 - p.padx0;
+        int inX = min(max(floorDiv(midX, p.upx), 0), p.inW);
+        int w = min(max(floorDiv(midX + p.kernelW, p.upx), 0), p.inW) - inX;
+        int kernelX = midX + p.kernelW - (inX + 1) * p.upx;
+
+        // Initialize pointers.
+        const T* xp = &p.x[((majorIdx * p.inH + inY) * p.inW + inX) * p.minorDim + minorIdx];
+        const T* kp = &p.k[kernelY * p.kernelW + kernelX];
+        int xpx = p.minorDim;
+        int kpx = -p.upx;
+        int xpy = p.inW * p.minorDim;
+        int kpy = -p.upy * p.kernelW;
+
+        // Inner loop.
+        float v = 0.0f;
+        for (int y = 0; y < h; y++)
+        {
+            for (int x = 0; x < w; x++)
+            {
+                v += (float)(*xp) * (float)(*kp);
+                xp += xpx;
+                kp += kpx;
+            }
+            xp += xpy - w * xpx;
+            kp += kpy - w * kpx;
+        }
+
+        // Store result.
+        p.y[((majorIdx * p.outH + outY) * p.outW + outX) * p.minorDim + minorIdx] = (T)v;
+    }
+}
+
+//------------------------------------------------------------------------
+// Specialized CUDA implementation for small filter kernels.
+
+template <class T, int upx, int upy, int downx, int downy, int kernelW, int kernelH, int tileOutW, int tileOutH>
+static __global__ void UpFirDn2DKernel_small(const UpFirDn2DKernelParams<T> p)
+{
+    //assert(kernelW % upx == 0);
+    //assert(kernelH % upy == 0);
+    const int tileInW = ((tileOutW - 1) * downx + kernelW - 1) / upx + 1;
+    const int tileInH = ((tileOutH - 1) * downy + kernelH - 1) / upy + 1;
+    __shared__ volatile float sk[kernelH][kernelW];
+    __shared__ volatile float sx[tileInH][tileInW];
+
+    // Calculate tile index.
+    int minorIdx = blockIdx.x;
+    int tileOutY = minorIdx / p.minorDim;
+    minorIdx -= tileOutY * p.minorDim;
+    tileOutY *= tileOutH;
+    int tileOutXBase = blockIdx.y * p.loopX * tileOutW;
+    int majorIdxBase = blockIdx.z * p.loopMajor;
+    if (tileOutXBase >= p.outW | tileOutY >= p.outH | majorIdxBase >= p.majorDim)
+        return;
+
+    // Load filter kernel (flipped).
+    for (int tapIdx = threadIdx.x; tapIdx < kernelH * kernelW; tapIdx += blockDim.x)
+    {
+        int ky = tapIdx / kernelW;
+        int kx = tapIdx - ky * kernelW;
+        float v = 0.0f;
+        if (kx < p.kernelW & ky < p.kernelH)
+            v = (float)p.k[(p.kernelH - 1 - ky) * p.kernelW + (p.kernelW - 1 - kx)];
+        sk[ky][kx] = v;
+    }
+
+    // Loop over majorDim and outX.
+    for (int loopMajor = 0, majorIdx = majorIdxBase; loopMajor < p.loopMajor & majorIdx < p.majorDim; loopMajor++, majorIdx++)
+    for (int loopX = 0, tileOutX = tileOutXBase; loopX < p.loopX & tileOutX < p.outW; loopX++, tileOutX += tileOutW)
+    {
+        // Load input pixels.
+        int tileMidX = tileOutX * downx + upx - 1 - p.padx0;
+        int tileMidY = tileOutY * downy + upy - 1 - p.pady0;
+        int tileInX = floorDiv(tileMidX, upx);
+        int tileInY = floorDiv(tileMidY, upy);
+        __syncthreads();
+        for (int inIdx = threadIdx.x; inIdx < tileInH * tileInW; inIdx += blockDim.x)
+        {
+            int relInY = inIdx / tileInW;
+            int relInX = inIdx - relInY * tileInW;
+            int inX = relInX + tileInX;
+            int inY = relInY + tileInY;
+            float v = 0.0f;
+            if (inX >= 0 & inY >= 0 & inX < p.inW & inY < p.inH)
+                v = (float)p.x[((majorIdx * p.inH + inY) * p.inW + inX) * p.minorDim + minorIdx];
+            sx[relInY][relInX] = v;
+        }
+
+        // Loop over output pixels.
+        __syncthreads();
+        for (int outIdx = threadIdx.x; outIdx < tileOutH * tileOutW; outIdx += blockDim.x)
+        {
+            int relOutY = outIdx / tileOutW;
+            int relOutX = outIdx - relOutY * tileOutW;
+            int outX = relOutX + tileOutX;
+            int outY = relOutY + tileOutY;
+
+            // Setup receptive field.
+            int midX = tileMidX + relOutX * downx;
+            int midY = tileMidY + relOutY * downy;
+            int inX = floorDiv(midX, upx);
+            int inY = floorDiv(midY, upy);
+            int relInX = inX - tileInX;
+            int relInY = inY - tileInY;
+            int kernelX = (inX + 1) * upx - midX - 1; // flipped
+            int kernelY = (inY + 1) * upy - midY - 1; // flipped
+
+            // Inner loop.
+            float v = 0.0f;
+            #pragma unroll
+            for (int y = 0; y < kernelH / upy; y++)
+                #pragma unroll
+                for (int x = 0; x < kernelW / upx; x++)
+                    v += sx[relInY + y][relInX + x] * sk[kernelY + y * upy][kernelX + x * upx];
+
+            // Store result.
+            if (outX < p.outW & outY < p.outH)
+                p.y[((majorIdx * p.outH + outY) * p.outW + outX) * p.minorDim + minorIdx] = (T)v;
+        }
+    }
+}
+
+//------------------------------------------------------------------------
+// TensorFlow op.
+
+template <class T>
+struct UpFirDn2DOp : public OpKernel
+{
+    UpFirDn2DKernelParams<T> m_attribs;
+
+    UpFirDn2DOp(OpKernelConstruction* ctx) : OpKernel(ctx)
+    {
+        memset(&m_attribs, 0, sizeof(m_attribs));
+        OP_REQUIRES_OK(ctx, ctx->GetAttr("upx", &m_attribs.upx));
+        OP_REQUIRES_OK(ctx, ctx->GetAttr("upy", &m_attribs.upy));
+        OP_REQUIRES_OK(ctx, ctx->GetAttr("downx", &m_attribs.downx));
+        OP_REQUIRES_OK(ctx, ctx->GetAttr("downy", &m_attribs.downy));
+        OP_REQUIRES_OK(ctx, ctx->GetAttr("padx0", &m_attribs.padx0));
+        OP_REQUIRES_OK(ctx, ctx->GetAttr("padx1", &m_attribs.padx1));
+        OP_REQUIRES_OK(ctx, ctx->GetAttr("pady0", &m_attribs.pady0));
+        OP_REQUIRES_OK(ctx, ctx->GetAttr("pady1", &m_attribs.pady1));
+        OP_REQUIRES(ctx, m_attribs.upx >= 1 && m_attribs.upy >= 1, errors::InvalidArgument("upx and upy must be at least 1x1"));
+        OP_REQUIRES(ctx, m_attribs.downx >= 1 && m_attribs.downy >= 1, errors::InvalidArgument("downx and downy must be at least 1x1"));
+    }
+
+    void Compute(OpKernelContext* ctx)
+    {
+        UpFirDn2DKernelParams<T> p = m_attribs;
+        cudaStream_t stream = ctx->eigen_device<Eigen::GpuDevice>().stream();
+
+        const Tensor& x = ctx->input(0); // [majorDim, inH, inW, minorDim]
+        const Tensor& k = ctx->input(1); // [kernelH, kernelW]
+        p.x = x.flat<T>().data();
+        p.k = k.flat<T>().data();
+        OP_REQUIRES(ctx, x.dims() == 4, errors::InvalidArgument("input must have rank 4"));
+        OP_REQUIRES(ctx, k.dims() == 2, errors::InvalidArgument("kernel must have rank 2"));
+        OP_REQUIRES(ctx, x.NumElements() <= kint32max, errors::InvalidArgument("input too large"));
+        OP_REQUIRES(ctx, k.NumElements() <= kint32max, errors::InvalidArgument("kernel too large"));
+
+        p.majorDim  = (int)x.dim_size(0);
+        p.inH       = (int)x.dim_size(1);
+        p.inW       = (int)x.dim_size(2);
+        p.minorDim  = (int)x.dim_size(3);
+        p.kernelH   = (int)k.dim_size(0);
+        p.kernelW   = (int)k.dim_size(1);
+        OP_REQUIRES(ctx, p.kernelW >= 1 && p.kernelH >= 1, errors::InvalidArgument("kernel must be at least 1x1"));
+
+        p.outW = (p.inW * p.upx + p.padx0 + p.padx1 - p.kernelW + p.downx) / p.downx;
+        p.outH = (p.inH * p.upy + p.pady0 + p.pady1 - p.kernelH + p.downy) / p.downy;
+        OP_REQUIRES(ctx, p.outW >= 1 && p.outH >= 1, errors::InvalidArgument("output must be at least 1x1"));
+
+        Tensor* y = NULL; // [majorDim, outH, outW, minorDim]
+        TensorShape ys;
+        ys.AddDim(p.majorDim);
+        ys.AddDim(p.outH);
+        ys.AddDim(p.outW);
+        ys.AddDim(p.minorDim);
+        OP_REQUIRES_OK(ctx, ctx->allocate_output(0, ys, &y));
+        p.y = y->flat<T>().data();
+        OP_REQUIRES(ctx, y->NumElements() <= kint32max, errors::InvalidArgument("output too large"));
+
+        // Choose CUDA kernel to use.
+        void* cudaKernel = (void*)UpFirDn2DKernel_large<T>;
+        int tileOutW = -1;
+        int tileOutH = -1;
+
+        if (p.upx == 1 && p.upy == 1 && p.downx == 1 && p.downy == 1 && p.kernelW <= 7  && p.kernelH <= 7 ) { cudaKernel = (void*)UpFirDn2DKernel_small<T, 1,1, 1,1, 7,7,  64,16>; tileOutW = 64;  tileOutH = 16; }
+        if (p.upx == 1 && p.upy == 1 && p.downx == 1 && p.downy == 1 && p.kernelW <= 6  && p.kernelH <= 6 ) { cudaKernel = (void*)UpFirDn2DKernel_small<T, 1,1, 1,1, 6,6,  64,16>; tileOutW = 64;  tileOutH = 16; }
+        if (p.upx == 1 && p.upy == 1 && p.downx == 1 && p.downy == 1 && p.kernelW <= 5  && p.kernelH <= 5 ) { cudaKernel = (void*)UpFirDn2DKernel_small<T, 1,1, 1,1, 5,5,  64,16>; tileOutW = 64;  tileOutH = 16; }
+        if (p.upx == 1 && p.upy == 1 && p.downx == 1 && p.downy == 1 && p.kernelW <= 4  && p.kernelH <= 4 ) { cudaKernel = (void*)UpFirDn2DKernel_small<T, 1,1, 1,1, 4,4,  64,16>; tileOutW = 64;  tileOutH = 16; }
+        if (p.upx == 1 && p.upy == 1 && p.downx == 1 && p.downy == 1 && p.kernelW <= 3  && p.kernelH <= 3 ) { cudaKernel = (void*)UpFirDn2DKernel_small<T, 1,1, 1,1, 3,3,  64,16>; tileOutW = 64;  tileOutH = 16; }
+        if (p.upx == 1 && p.upy == 1 && p.downx == 1 && p.downy == 1 && p.kernelW <= 24 && p.kernelH <= 1 ) { cudaKernel = (void*)UpFirDn2DKernel_small<T, 1,1, 1,1, 24,1, 128,8>; tileOutW = 128; tileOutH = 8;  }
+        if (p.upx == 1 && p.upy == 1 && p.downx == 1 && p.downy == 1 && p.kernelW <= 20 && p.kernelH <= 1 ) { cudaKernel = (void*)UpFirDn2DKernel_small<T, 1,1, 1,1, 20,1, 128,8>; tileOutW = 128; tileOutH = 8;  }
+        if (p.upx == 1 && p.upy == 1 && p.downx == 1 && p.downy == 1 && p.kernelW <= 16 && p.kernelH <= 1 ) { cudaKernel = (void*)UpFirDn2DKernel_small<T, 1,1, 1,1, 16,1, 128,8>; tileOutW = 128; tileOutH = 8;  }
+        if (p.upx == 1 && p.upy == 1 && p.downx == 1 && p.downy == 1 && p.kernelW <= 12 && p.kernelH <= 1 ) { cudaKernel = (void*)UpFirDn2DKernel_small<T, 1,1, 1,1, 12,1, 128,8>; tileOutW = 128; tileOutH = 8;  }
+        if (p.upx == 1 && p.upy == 1 && p.downx == 1 && p.downy == 1 && p.kernelW <= 8  && p.kernelH <= 1 ) { cudaKernel = (void*)UpFirDn2DKernel_small<T, 1,1, 1,1, 8,1,  128,8>; tileOutW = 128; tileOutH = 8;  }
+        if (p.upx == 1 && p.upy == 1 && p.downx == 1 && p.downy == 1 && p.kernelW <= 1  && p.kernelH <= 24) { cudaKernel = (void*)UpFirDn2DKernel_small<T, 1,1, 1,1, 1,24, 32,32>; tileOutW = 32;  tileOutH = 32; }
+        if (p.upx == 1 && p.upy == 1 && p.downx == 1 && p.downy == 1 && p.kernelW <= 1  && p.kernelH <= 20) { cudaKernel = (void*)UpFirDn2DKernel_small<T, 1,1, 1,1, 1,20, 32,32>; tileOutW = 32;  tileOutH = 32; }
+        if (p.upx == 1 && p.upy == 1 && p.downx == 1 && p.downy == 1 && p.kernelW <= 1  && p.kernelH <= 16) { cudaKernel = (void*)UpFirDn2DKernel_small<T, 1,1, 1,1, 1,16, 32,32>; tileOutW = 32;  tileOutH = 32; }
+        if (p.upx == 1 && p.upy == 1 && p.downx == 1 && p.downy == 1 && p.kernelW <= 1  && p.kernelH <= 12) { cudaKernel = (void*)UpFirDn2DKernel_small<T, 1,1, 1,1, 1,12, 32,32>; tileOutW = 32;  tileOutH = 32; }
+        if (p.upx == 1 && p.upy == 1 && p.downx == 1 && p.downy == 1 && p.kernelW <= 1  && p.kernelH <= 8 ) { cudaKernel = (void*)UpFirDn2DKernel_small<T, 1,1, 1,1, 1,8,  32,32>; tileOutW = 32;  tileOutH = 32; }
+
+        if (p.upx == 2 && p.upy == 2 && p.downx == 1 && p.downy == 1 && p.kernelW <= 8  && p.kernelH <= 8 ) { cudaKernel = (void*)UpFirDn2DKernel_small<T, 2,2, 1,1, 8,8,  64,16>; tileOutW = 64;  tileOutH = 16; }
+        if (p.upx == 2 && p.upy == 2 && p.downx == 1 && p.downy == 1 && p.kernelW <= 6  && p.kernelH <= 6 ) { cudaKernel = (void*)UpFirDn2DKernel_small<T, 2,2, 1,1, 6,6,  64,16>; tileOutW = 64;  tileOutH = 16; }
+        if (p.upx == 2 && p.upy == 2 && p.downx == 1 && p.downy == 1 && p.kernelW <= 4  && p.kernelH <= 4 ) { cudaKernel = (void*)UpFirDn2DKernel_small<T, 2,2, 1,1, 4,4,  64,16>; tileOutW = 64;  tileOutH = 16; }
+        if (p.upx == 2 && p.upy == 2 && p.downx == 1 && p.downy == 1 && p.kernelW <= 2  && p.kernelH <= 2 ) { cudaKernel = (void*)UpFirDn2DKernel_small<T, 2,2, 1,1, 2,2,  64,16>; tileOutW = 64;  tileOutH = 16; }
+        if (p.upx == 2 && p.upy == 1 && p.downx == 1 && p.downy == 1 && p.kernelW <= 24 && p.kernelH <= 1 ) { cudaKernel = (void*)UpFirDn2DKernel_small<T, 2,1, 1,1, 24,1, 128,8>; tileOutW = 128; tileOutH = 8;  }
+        if (p.upx == 2 && p.upy == 1 && p.downx == 1 && p.downy == 1 && p.kernelW <= 20 && p.kernelH <= 1 ) { cudaKernel = (void*)UpFirDn2DKernel_small<T, 2,1, 1,1, 20,1, 128,8>; tileOutW = 128; tileOutH = 8;  }
+        if (p.upx == 2 && p.upy == 1 && p.downx == 1 && p.downy == 1 && p.kernelW <= 16 && p.kernelH <= 1 ) { cudaKernel = (void*)UpFirDn2DKernel_small<T, 2,1, 1,1, 16,1, 128,8>; tileOutW = 128; tileOutH = 8;  }
+        if (p.upx == 2 && p.upy == 1 && p.downx == 1 && p.downy == 1 && p.kernelW <= 12 && p.kernelH <= 1 ) { cudaKernel = (void*)UpFirDn2DKernel_small<T, 2,1, 1,1, 12,1, 128,8>; tileOutW = 128; tileOutH = 8;  }
+        if (p.upx == 2 && p.upy == 1 && p.downx == 1 && p.downy == 1 && p.kernelW <= 8  && p.kernelH <= 1 ) { cudaKernel = (void*)UpFirDn2DKernel_small<T, 2,1, 1,1, 8,1,  128,8>; tileOutW = 128; tileOutH = 8;  }
+        if (p.upx == 1 && p.upy == 2 && p.downx == 1 && p.downy == 1 && p.kernelW <= 1  && p.kernelH <= 24) { cudaKernel = (void*)UpFirDn2DKernel_small<T, 1,2, 1,1, 1,24, 32,32>; tileOutW = 32;  tileOutH = 32; }
+        if (p.upx == 1 && p.upy == 2 && p.downx == 1 && p.downy == 1 && p.kernelW <= 1  && p.kernelH <= 20) { cudaKernel = (void*)UpFirDn2DKernel_small<T, 1,2, 1,1, 1,20, 32,32>; tileOutW = 32;  tileOutH = 32; }
+        if (p.upx == 1 && p.upy == 2 && p.downx == 1 && p.downy == 1 && p.kernelW <= 1  && p.kernelH <= 16) { cudaKernel = (void*)UpFirDn2DKernel_small<T, 1,2, 1,1, 1,16, 32,32>; tileOutW = 32;  tileOutH = 32; }
+        if (p.upx == 1 && p.upy == 2 && p.downx == 1 && p.downy == 1 && p.kernelW <= 1  && p.kernelH <= 12) { cudaKernel = (void*)UpFirDn2DKernel_small<T, 1,2, 1,1, 1,12, 32,32>; tileOutW = 32;  tileOutH = 32; }
+        if (p.upx == 1 && p.upy == 2 && p.downx == 1 && p.downy == 1 && p.kernelW <= 1  && p.kernelH <= 8 ) { cudaKernel = (void*)UpFirDn2DKernel_small<T, 1,2, 1,1, 1,8,  32,32>; tileOutW = 32;  tileOutH = 32; }
+
+        if (p.upx == 1 && p.upy == 1 && p.downx == 2 && p.downy == 2 && p.kernelW <= 8  && p.kernelH <= 8 ) { cudaKernel = (void*)UpFirDn2DKernel_small<T, 1,1, 2,2, 8,8,  32,8 >; tileOutW = 32;  tileOutH = 8;  }
+        if (p.upx == 1 && p.upy == 1 && p.downx == 2 && p.downy == 2 && p.kernelW <= 6  && p.kernelH <= 6 ) { cudaKernel = (void*)UpFirDn2DKernel_small<T, 1,1, 2,2, 6,6,  32,8 >; tileOutW = 32;  tileOutH = 8;  }
+        if (p.upx == 1 && p.upy == 1 && p.downx == 2 && p.downy == 2 && p.kernelW <= 4  && p.kernelH <= 4 ) { cudaKernel = (void*)UpFirDn2DKernel_small<T, 1,1, 2,2, 4,4,  32,8 >; tileOutW = 32;  tileOutH = 8;  }
+        if (p.upx == 1 && p.upy == 1 && p.downx == 2 && p.downy == 2 && p.kernelW <= 2  && p.kernelH <= 2 ) { cudaKernel = (void*)UpFirDn2DKernel_small<T, 1,1, 2,2, 2,2,  32,8 >; tileOutW = 32;  tileOutH = 8;  }
+        if (p.upx == 1 && p.upy == 1 && p.downx == 2 && p.downy == 1 && p.kernelW <= 24 && p.kernelH <= 1 ) { cudaKernel = (void*)UpFirDn2DKernel_small<T, 1,1, 2,1, 24,1, 64,8 >; tileOutW = 64;  tileOutH = 8;  }
+        if (p.upx == 1 && p.upy == 1 && p.downx == 2 && p.downy == 1 && p.kernelW <= 20 && p.kernelH <= 1 ) { cudaKernel = (void*)UpFirDn2DKernel_small<T, 1,1, 2,1, 20,1, 64,8 >; tileOutW = 64;  tileOutH = 8;  }
+        if (p.upx == 1 && p.upy == 1 && p.downx == 2 && p.downy == 1 && p.kernelW <= 16 && p.kernelH <= 1 ) { cudaKernel = (void*)UpFirDn2DKernel_small<T, 1,1, 2,1, 16,1, 64,8 >; tileOutW = 64;  tileOutH = 8;  }
+        if (p.upx == 1 && p.upy == 1 && p.downx == 2 && p.downy == 1 && p.kernelW <= 12 && p.kernelH <= 1 ) { cudaKernel = (void*)UpFirDn2DKernel_small<T, 1,1, 2,1, 12,1, 64,8 >; tileOutW = 64;  tileOutH = 8;  }
+        if (p.upx == 1 && p.upy == 1 && p.downx == 2 && p.downy == 1 && p.kernelW <= 8  && p.kernelH <= 1 ) { cudaKernel = (void*)UpFirDn2DKernel_small<T, 1,1, 2,1, 8,1,  64,8 >; tileOutW = 64;  tileOutH = 8;  }
+        if (p.upx == 1 && p.upy == 1 && p.downx == 1 && p.downy == 2 && p.kernelW <= 1  && p.kernelH <= 24) { cudaKernel = (void*)UpFirDn2DKernel_small<T, 1,1, 1,2, 1,24, 32,16>; tileOutW = 32;  tileOutH = 16; }
+        if (p.upx == 1 && p.upy == 1 && p.downx == 1 && p.downy == 2 && p.kernelW <= 1  && p.kernelH <= 20) { cudaKernel = (void*)UpFirDn2DKernel_small<T, 1,1, 1,2, 1,20, 32,16>; tileOutW = 32;  tileOutH = 16; }
+        if (p.upx == 1 && p.upy == 1 && p.downx == 1 && p.downy == 2 && p.kernelW <= 1  && p.kernelH <= 16) { cudaKernel = (void*)UpFirDn2DKernel_small<T, 1,1, 1,2, 1,16, 32,16>; tileOutW = 32;  tileOutH = 16; }
+        if (p.upx == 1 && p.upy == 1 && p.downx == 1 && p.downy == 2 && p.kernelW <= 1  && p.kernelH <= 12) { cudaKernel = (void*)UpFirDn2DKernel_small<T, 1,1, 1,2, 1,12, 32,16>; tileOutW = 32;  tileOutH = 16; }
+        if (p.upx == 1 && p.upy == 1 && p.downx == 1 && p.downy == 2 && p.kernelW <= 1  && p.kernelH <= 8 ) { cudaKernel = (void*)UpFirDn2DKernel_small<T, 1,1, 1,2, 1,8,  32,16>; tileOutW = 32;  tileOutH = 16; }
+
+        // Choose launch params.
+        dim3 blockSize;
+        dim3 gridSize;
+        if (tileOutW > 0 && tileOutH > 0) // small
+        {
+            p.loopMajor = (p.majorDim - 1) / 16384 + 1;
+            p.loopX = 1;
+            blockSize = dim3(32 * 8, 1, 1);
+            gridSize = dim3(((p.outH - 1) / tileOutH + 1) * p.minorDim, (p.outW - 1) / (p.loopX * tileOutW) + 1, (p.majorDim - 1) / p.loopMajor + 1);
+        }
+        else // large
+        {
+            p.loopMajor = (p.majorDim - 1) / 16384 + 1;
+            p.loopX = 4;
+            blockSize = dim3(4, 32, 1);
+            gridSize = dim3((p.outH * p.minorDim - 1) / blockSize.x + 1, (p.outW - 1) / (p.loopX * blockSize.y) + 1, (p.majorDim - 1) / p.loopMajor + 1);
+        }
+
+        // Launch CUDA kernel.
+        void* args[] = {&p};
+        OP_CHECK_CUDA_ERROR(ctx, cudaLaunchKernel(cudaKernel, gridSize, blockSize, args, 0, stream));
+    }
+};
+
+REGISTER_OP("UpFirDn2D")
+    .Input      ("x: T")
+    .Input      ("k: T")
+    .Output     ("y: T")
+    .Attr       ("T: {float, half}")
+    .Attr       ("upx: int = 1")
+    .Attr       ("upy: int = 1")
+    .Attr       ("downx: int = 1")
+    .Attr       ("downy: int = 1")
+    .Attr       ("padx0: int = 0")
+    .Attr       ("padx1: int = 0")
+    .Attr       ("pady0: int = 0")
+    .Attr       ("pady1: int = 0");
+REGISTER_KERNEL_BUILDER(Name("UpFirDn2D").Device(DEVICE_GPU).TypeConstraint<float>("T"), UpFirDn2DOp<float>);
+REGISTER_KERNEL_BUILDER(Name("UpFirDn2D").Device(DEVICE_GPU).TypeConstraint<Eigen::half>("T"), UpFirDn2DOp<Eigen::half>);
+
+//------------------------------------------------------------------------

dnnlib/tflib/ops/upfirdn_2d.py

@@ -0,0 +1,418 @@
+# Copyright (c) 2020, NVIDIA CORPORATION.  All rights reserved.
+#
+# NVIDIA CORPORATION and its licensors retain all intellectual property
+# and proprietary rights in and to this software, related documentation
+# and any modifications thereto.  Any use, reproduction, disclosure or
+# distribution of this software and related documentation without an express
+# license agreement from NVIDIA CORPORATION is strictly prohibited.
+
+"""Custom TensorFlow ops for efficient resampling of 2D images."""
+
+import os
+import numpy as np
+import tensorflow as tf
+from .. import custom_ops
+
+def _get_plugin():
+    return custom_ops.get_plugin(os.path.splitext(__file__)[0] + '.cu')
+
+#----------------------------------------------------------------------------
+
+def upfirdn_2d(x, k, upx=1, upy=1, downx=1, downy=1, padx0=0, padx1=0, pady0=0, pady1=0, impl='cuda'):
+    r"""Pad, upsample, FIR filter, and downsample a batch of 2D images.
+
+    Accepts a batch of 2D images of the shape `[majorDim, inH, inW, minorDim]`
+    and performs the following operations for each image, batched across
+    `majorDim` and `minorDim`:
+
+    1. Upsample the image by inserting the zeros after each pixel (`upx`, `upy`).
+
+    2. Pad the image with zeros by the specified number of pixels on each side
+       (`padx0`, `padx1`, `pady0`, `pady1`). Specifying a negative value
+       corresponds to cropping the image.
+
+    3. Convolve the image with the specified 2D FIR filter (`k`), shrinking the
+       image so that the footprint of all output pixels lies within the input image.
+
+    4. Downsample the image by throwing away pixels (`downx`, `downy`).
+
+    This sequence of operations bears close resemblance to scipy.signal.upfirdn().
+    The fused op is considerably more efficient than performing the same calculation
+    using standard TensorFlow ops. It supports gradients of arbitrary order.
+
+    Args:
+        x:      Input tensor of the shape `[majorDim, inH, inW, minorDim]`.
+        k:      2D FIR filter of the shape `[firH, firW]`.
+        upx:    Integer upsampling factor along the X-axis (default: 1).
+        upy:    Integer upsampling factor along the Y-axis (default: 1).
+        downx:  Integer downsampling factor along the X-axis (default: 1).
+        downy:  Integer downsampling factor along the Y-axis (default: 1).
+        padx0:  Number of pixels to pad on the left side (default: 0).
+        padx1:  Number of pixels to pad on the right side (default: 0).
+        pady0:  Number of pixels to pad on the top side (default: 0).
+        pady1:  Number of pixels to pad on the bottom side (default: 0).
+        impl:   Name of the implementation to use. Can be `"ref"` or `"cuda"` (default).
+
+    Returns:
+        Tensor of the shape `[majorDim, outH, outW, minorDim]`, and same datatype as `x`.
+    """
+
+    impl_dict = {
+        'ref':  _upfirdn_2d_ref,
+        'cuda': _upfirdn_2d_cuda,
+    }
+    return impl_dict[impl](x=x, k=k, upx=upx, upy=upy, downx=downx, downy=downy, padx0=padx0, padx1=padx1, pady0=pady0, pady1=pady1)
+
+#----------------------------------------------------------------------------
+
+def _upfirdn_2d_ref(x, k, upx, upy, downx, downy, padx0, padx1, pady0, pady1):
+    """Slow reference implementation of `upfirdn_2d()` using standard TensorFlow ops."""
+
+    x = tf.convert_to_tensor(x)
+    k = np.asarray(k, dtype=np.float32)
+    assert x.shape.rank == 4
+    inH = x.shape[1].value
+    inW = x.shape[2].value
+    minorDim = _shape(x, 3)
+    kernelH, kernelW = k.shape
+    assert inW >= 1 and inH >= 1
+    assert kernelW >= 1 and kernelH >= 1
+    assert isinstance(upx, int) and isinstance(upy, int)
+    assert isinstance(downx, int) and isinstance(downy, int)
+    assert isinstance(padx0, int) and isinstance(padx1, int)
+    assert isinstance(pady0, int) and isinstance(pady1, int)
+
+    # Upsample (insert zeros).
+    x = tf.reshape(x, [-1, inH, 1, inW, 1, minorDim])
+    x = tf.pad(x, [[0, 0], [0, 0], [0, upy - 1], [0, 0], [0, upx - 1], [0, 0]])
+    x = tf.reshape(x, [-1, inH * upy, inW * upx, minorDim])
+
+    # Pad (crop if negative).
+    x = tf.pad(x, [[0, 0], [max(pady0, 0), max(pady1, 0)], [max(padx0, 0), max(padx1, 0)], [0, 0]])
+    x = x[:, max(-pady0, 0) : x.shape[1].value - max(-pady1, 0), max(-padx0, 0) : x.shape[2].value - max(-padx1, 0), :]
+
+    # Convolve with filter.
+    x = tf.transpose(x, [0, 3, 1, 2])
+    x = tf.reshape(x, [-1, 1, inH * upy + pady0 + pady1, inW * upx + padx0 + padx1])
+    w = tf.constant(k[::-1, ::-1, np.newaxis, np.newaxis], dtype=x.dtype)
+    x = tf.nn.conv2d(x, w, strides=[1,1,1,1], padding='VALID', data_format='NCHW')
+    x = tf.reshape(x, [-1, minorDim, inH * upy + pady0 + pady1 - kernelH + 1, inW * upx + padx0 + padx1 - kernelW + 1])
+    x = tf.transpose(x, [0, 2, 3, 1])
+
+    # Downsample (throw away pixels).
+    return x[:, ::downy, ::downx, :]
+
+#----------------------------------------------------------------------------
+
+def _upfirdn_2d_cuda(x, k, upx, upy, downx, downy, padx0, padx1, pady0, pady1):
+    """Fast CUDA implementation of `upfirdn_2d()` using custom ops."""
+
+    x = tf.convert_to_tensor(x)
+    k = np.asarray(k, dtype=np.float32)
+    majorDim, inH, inW, minorDim = x.shape.as_list()
+    kernelH, kernelW = k.shape
+    assert inW >= 1 and inH >= 1
+    assert kernelW >= 1 and kernelH >= 1
+    assert isinstance(upx, int) and isinstance(upy, int)
+    assert isinstance(downx, int) and isinstance(downy, int)
+    assert isinstance(padx0, int) and isinstance(padx1, int)
+    assert isinstance(pady0, int) and isinstance(pady1, int)
+
+    outW = (inW * upx + padx0 + padx1 - kernelW) // downx + 1
+    outH = (inH * upy + pady0 + pady1 - kernelH) // downy + 1
+    assert outW >= 1 and outH >= 1
+
+    cuda_op = _get_plugin().up_fir_dn2d
+    kc = tf.constant(k, dtype=x.dtype)
+    gkc = tf.constant(k[::-1, ::-1], dtype=x.dtype)
+    gpadx0 = kernelW - padx0 - 1
+    gpady0 = kernelH - pady0 - 1
+    gpadx1 = inW * upx - outW * downx + padx0 - upx + 1
+    gpady1 = inH * upy - outH * downy + pady0 - upy + 1
+
+    @tf.custom_gradient
+    def func(x):
+        y = cuda_op(x=x, k=kc, upx=int(upx), upy=int(upy), downx=int(downx), downy=int(downy), padx0=int(padx0), padx1=int(padx1), pady0=int(pady0), pady1=int(pady1))
+        y.set_shape([majorDim, outH, outW, minorDim])
+        @tf.custom_gradient
+        def grad(dy):
+            dx = cuda_op(x=dy, k=gkc, upx=int(downx), upy=int(downy), downx=int(upx), downy=int(upy), padx0=int(gpadx0), padx1=int(gpadx1), pady0=int(gpady0), pady1=int(gpady1))
+            dx.set_shape([majorDim, inH, inW, minorDim])
+            return dx, func
+        return y, grad
+    return func(x)
+
+#----------------------------------------------------------------------------
+
+def filter_2d(x, k, gain=1, padding=0, data_format='NCHW', impl='cuda'):
+    r"""Filter a batch of 2D images with the given FIR filter.
+
+    Accepts a batch of 2D images of the shape `[N, C, H, W]` or `[N, H, W, C]`
+    and filters each image with the given filter. The filter is normalized so that
+    if the input pixels are constant, they will be scaled by the specified `gain`.
+    Pixels outside the image are assumed to be zero.
+
+    Args:
+        x:            Input tensor of the shape `[N, C, H, W]` or `[N, H, W, C]`.
+        k:            FIR filter of the shape `[firH, firW]` or `[firN]` (separable).
+        gain:         Scaling factor for signal magnitude (default: 1.0).
+        padding:      Number of pixels to pad or crop the output on each side (default: 0).
+        data_format:  `'NCHW'` or `'NHWC'` (default: `'NCHW'`).
+        impl:         Name of the implementation to use. Can be `"ref"` or `"cuda"` (default).
+
+    Returns:
+        Tensor of the same shape and datatype as `x`.
+    """
+
+    assert isinstance(padding, int)
+    k = _FilterKernel(k=k, gain=gain)
+    assert k.w == k.h
+    pad0 = k.w // 2 + padding
+    pad1 = (k.w - 1) // 2 + padding
+    return _simple_upfirdn_2d(x, k, pad0=pad0, pad1=pad1, data_format=data_format, impl=impl)
+
+#----------------------------------------------------------------------------
+
+def upsample_2d(x, k=None, factor=2, gain=1, padding=0, data_format='NCHW', impl='cuda'):
+    r"""Upsample a batch of 2D images with the given filter.
+
+    Accepts a batch of 2D images of the shape `[N, C, H, W]` or `[N, H, W, C]`
+    and upsamples each image with the given filter. The filter is normalized so that
+    if the input pixels are constant, they will be scaled by the specified `gain`.
+    Pixels outside the image are assumed to be zero, and the filter is padded with
+    zeros so that its shape is a multiple of the upsampling factor.
+
+    Args:
+        x:            Input tensor of the shape `[N, C, H, W]` or `[N, H, W, C]`.
+        k:            FIR filter of the shape `[firH, firW]` or `[firN]` (separable).
+                      The default is `[1] * factor`, which corresponds to nearest-neighbor
+                      upsampling.
+        factor:       Integer upsampling factor (default: 2).
+        gain:         Scaling factor for signal magnitude (default: 1.0).
+        padding:      Number of pixels to pad or crop the output on each side (default: 0).
+        data_format:  `'NCHW'` or `'NHWC'` (default: `'NCHW'`).
+        impl:         Name of the implementation to use. Can be `"ref"` or `"cuda"` (default).
+
+    Returns:
+        Tensor of the shape `[N, C, H * factor, W * factor]` or
+        `[N, H * factor, W * factor, C]`, and same datatype as `x`.
+    """
+
+    assert isinstance(factor, int) and factor >= 1
+    assert isinstance(padding, int)
+    k = _FilterKernel(k if k is not None else [1] * factor, gain * (factor ** 2))
+    assert k.w == k.h
+    pad0 = (k.w + factor - 1) // 2 + padding
+    pad1 = (k.w - factor) // 2 + padding
+    return _simple_upfirdn_2d(x, k, up=factor, pad0=pad0, pad1=pad1, data_format=data_format, impl=impl)
+
+#----------------------------------------------------------------------------
+
+def downsample_2d(x, k=None, factor=2, gain=1, padding=0, data_format='NCHW', impl='cuda'):
+    r"""Downsample a batch of 2D images with the given filter.
+
+    Accepts a batch of 2D images of the shape `[N, C, H, W]` or `[N, H, W, C]`
+    and downsamples each image with the given filter. The filter is normalized so that
+    if the input pixels are constant, they will be scaled by the specified `gain`.
+    Pixels outside the image are assumed to be zero, and the filter is padded with
+    zeros so that its shape is a multiple of the downsampling factor.
+
+    Args:
+        x:            Input tensor of the shape `[N, C, H, W]` or `[N, H, W, C]`.
+        k:            FIR filter of the shape `[firH, firW]` or `[firN]` (separable).
+                      The default is `[1] * factor`, which corresponds to average pooling.
+        factor:       Integer downsampling factor (default: 2).
+        gain:         Scaling factor for signal magnitude (default: 1.0).
+        padding:      Number of pixels to pad or crop the output on each side (default: 0).
+        data_format:  `'NCHW'` or `'NHWC'` (default: `'NCHW'`).
+        impl:         Name of the implementation to use. Can be `"ref"` or `"cuda"` (default).
+
+    Returns:
+        Tensor of the shape `[N, C, H // factor, W // factor]` or
+        `[N, H // factor, W // factor, C]`, and same datatype as `x`.
+    """
+
+    assert isinstance(factor, int) and factor >= 1
+    assert isinstance(padding, int)
+    k = _FilterKernel(k if k is not None else [1] * factor, gain)
+    assert k.w == k.h
+    pad0 = (k.w - factor + 1) // 2 + padding * factor
+    pad1 = (k.w - factor) // 2 + padding * factor
+    return _simple_upfirdn_2d(x, k, down=factor, pad0=pad0, pad1=pad1, data_format=data_format, impl=impl)
+
+#----------------------------------------------------------------------------
+
+def upsample_conv_2d(x, w, k=None, factor=2, gain=1, padding=0, data_format='NCHW', impl='cuda'):
+    r"""Fused `upsample_2d()` followed by `tf.nn.conv2d()`.
+
+    Padding is performed only once at the beginning, not between the operations.
+    The fused op is considerably more efficient than performing the same calculation
+    using standard TensorFlow ops. It supports gradients of arbitrary order.
+
+    Args:
+        x:            Input tensor of the shape `[N, C, H, W]` or `[N, H, W, C]`.
+        w:            Weight tensor of the shape `[filterH, filterW, inChannels, outChannels]`.
+                      Grouped convolution can be performed by `inChannels = x.shape[0] // numGroups`.
+        k:            FIR filter of the shape `[firH, firW]` or `[firN]` (separable).
+                      The default is `[1] * factor`, which corresponds to nearest-neighbor
+                      upsampling.
+        factor:       Integer upsampling factor (default: 2).
+        gain:         Scaling factor for signal magnitude (default: 1.0).
+        padding:      Number of pixels to pad or crop the output on each side (default: 0).
+        data_format:  `'NCHW'` or `'NHWC'` (default: `'NCHW'`).
+        impl:         Name of the implementation to use. Can be `"ref"` or `"cuda"` (default).
+
+    Returns:
+        Tensor of the shape `[N, C, H * factor, W * factor]` or
+        `[N, H * factor, W * factor, C]`, and same datatype as `x`.
+    """
+
+    assert isinstance(factor, int) and factor >= 1
+    assert isinstance(padding, int)
+
+    # Check weight shape.
+    w = tf.convert_to_tensor(w)
+    ch, cw, _inC, _outC = w.shape.as_list()
+    inC = _shape(w, 2)
+    outC = _shape(w, 3)
+    assert cw == ch
+
+    # Fast path for 1x1 convolution.
+    if cw == 1 and ch == 1:
+        x = tf.nn.conv2d(x, w, data_format=data_format, strides=[1,1,1,1], padding='VALID')
+        x = upsample_2d(x, k, factor=factor, gain=gain, padding=padding, data_format=data_format, impl=impl)
+        return x
+
+    # Setup filter kernel.
+    k = _FilterKernel(k if k is not None else [1] * factor, gain * (factor ** 2))
+    assert k.w == k.h
+
+    # Determine data dimensions.
+    if data_format == 'NCHW':
+        stride = [1, 1, factor, factor]
+        output_shape = [_shape(x, 0), outC, (_shape(x, 2) - 1) * factor + ch, (_shape(x, 3) - 1) * factor + cw]
+        num_groups = _shape(x, 1) // inC
+    else:
+        stride = [1, factor, factor, 1]
+        output_shape = [_shape(x, 0), (_shape(x, 1) - 1) * factor + ch, (_shape(x, 2) - 1) * factor + cw, outC]
+        num_groups = _shape(x, 3) // inC
+
+    # Transpose weights.
+    w = tf.reshape(w, [ch, cw, inC, num_groups, -1])
+    w = tf.transpose(w[::-1, ::-1], [0, 1, 4, 3, 2])
+    w = tf.reshape(w, [ch, cw, -1, num_groups * inC])
+
+    # Execute.
+    x = tf.nn.conv2d_transpose(x, w, output_shape=output_shape, strides=stride, padding='VALID', data_format=data_format)
+    pad0 = (k.w + factor - cw) // 2 + padding
+    pad1 = (k.w - factor - cw + 3) // 2 + padding
+    return _simple_upfirdn_2d(x, k, pad0=pad0, pad1=pad1, data_format=data_format, impl=impl)
+
+#----------------------------------------------------------------------------
+
+def conv_downsample_2d(x, w, k=None, factor=2, gain=1, padding=0, data_format='NCHW', impl='cuda'):
+    r"""Fused `tf.nn.conv2d()` followed by `downsample_2d()`.
+
+    Padding is performed only once at the beginning, not between the operations.
+    The fused op is considerably more efficient than performing the same calculation
+    using standard TensorFlow ops. It supports gradients of arbitrary order.
+
+    Args:
+        x:            Input tensor of the shape `[N, C, H, W]` or `[N, H, W, C]`.
+        w:            Weight tensor of the shape `[filterH, filterW, inChannels, outChannels]`.
+                      Grouped convolution can be performed by `inChannels = x.shape[0] // numGroups`.
+        k:            FIR filter of the shape `[firH, firW]` or `[firN]` (separable).
+                      The default is `[1] * factor`, which corresponds to average pooling.
+        factor:       Integer downsampling factor (default: 2).
+        gain:         Scaling factor for signal magnitude (default: 1.0).
+        padding:      Number of pixels to pad or crop the output on each side (default: 0).
+        data_format:  `'NCHW'` or `'NHWC'` (default: `'NCHW'`).
+        impl:         Name of the implementation to use. Can be `"ref"` or `"cuda"` (default).
+
+    Returns:
+        Tensor of the shape `[N, C, H // factor, W // factor]` or
+        `[N, H // factor, W // factor, C]`, and same datatype as `x`.
+    """
+
+    assert isinstance(factor, int) and factor >= 1
+    assert isinstance(padding, int)
+
+    # Check weight shape.
+    w = tf.convert_to_tensor(w)
+    ch, cw, _inC, _outC = w.shape.as_list()
+    assert cw == ch
+
+    # Fast path for 1x1 convolution.
+    if cw == 1 and ch == 1:
+        x = downsample_2d(x, k, factor=factor, gain=gain, padding=padding, data_format=data_format, impl=impl)
+        x = tf.nn.conv2d(x, w, data_format=data_format, strides=[1,1,1,1], padding='VALID')
+        return x
+
+    # Setup filter kernel.
+    k = _FilterKernel(k if k is not None else [1] * factor, gain)
+    assert k.w == k.h
+
+    # Determine stride.
+    if data_format == 'NCHW':
+        s = [1, 1, factor, factor]
+    else:
+        s = [1, factor, factor, 1]
+
+    # Execute.
+    pad0 = (k.w - factor + cw) // 2 + padding * factor
+    pad1 = (k.w - factor + cw - 1) // 2 + padding * factor
+    x = _simple_upfirdn_2d(x, k, pad0=pad0, pad1=pad1, data_format=data_format, impl=impl)
+    return tf.nn.conv2d(x, w, strides=s, padding='VALID', data_format=data_format)
+
+#----------------------------------------------------------------------------
+# Internal helpers.
+
+class _FilterKernel:
+    def __init__(self, k, gain=1):
+        k = np.asarray(k, dtype=np.float32)
+        k /= np.sum(k)
+
+        # Separable.
+        if k.ndim == 1 and k.size >= 8:
+            self.w = k.size
+            self.h = k.size
+            self.kx = k[np.newaxis, :]
+            self.ky = k[:, np.newaxis] * gain
+            self.kxy = None
+
+        # Non-separable.
+        else:
+            if k.ndim == 1:
+                k = np.outer(k, k)
+            assert k.ndim == 2
+            self.w = k.shape[1]
+            self.h = k.shape[0]
+            self.kx = None
+            self.ky = None
+            self.kxy = k * gain
+
+def _simple_upfirdn_2d(x, k, up=1, down=1, pad0=0, pad1=0, data_format='NCHW', impl='cuda'):
+    assert isinstance(k, _FilterKernel)
+    assert data_format in ['NCHW', 'NHWC']
+    assert x.shape.rank == 4
+    y = x
+    if data_format == 'NCHW':
+        y = tf.reshape(y, [-1, _shape(y, 2), _shape(y, 3), 1])
+    if k.kx is not None:
+        y = upfirdn_2d(y, k.kx, upx=up, downx=down, padx0=pad0, padx1=pad1, impl=impl)
+    if k.ky is not None:
+        y = upfirdn_2d(y, k.ky, upy=up, downy=down, pady0=pad0, pady1=pad1, impl=impl)
+    if k.kxy is not None:
+        y = upfirdn_2d(y, k.kxy, upx=up, upy=up, downx=down, downy=down, padx0=pad0, padx1=pad1, pady0=pad0, pady1=pad1, impl=impl)
+    if data_format == 'NCHW':
+        y = tf.reshape(y, [-1, _shape(x, 1), _shape(y, 1), _shape(y, 2)])
+    return y
+
+def _shape(tf_expr, dim_idx):
+    if tf_expr.shape.rank is not None:
+        dim = tf_expr.shape[dim_idx].value
+        if dim is not None:
+            return dim
+    return tf.shape(tf_expr)[dim_idx]
+
+#----------------------------------------------------------------------------

dnnlib/tflib/optimizer.py

@@ -0,0 +1,372 @@
+# Copyright (c) 2020, NVIDIA CORPORATION.  All rights reserved.
+#
+# NVIDIA CORPORATION and its licensors retain all intellectual property
+# and proprietary rights in and to this software, related documentation
+# and any modifications thereto.  Any use, reproduction, disclosure or
+# distribution of this software and related documentation without an express
+# license agreement from NVIDIA CORPORATION is strictly prohibited.
+
+"""Helper wrapper for a Tensorflow optimizer."""
+
+import platform
+import numpy as np
+import tensorflow as tf
+
+from collections import OrderedDict
+from typing import List, Union
+
+from . import autosummary
+from . import tfutil
+from .. import util
+
+from .tfutil import TfExpression, TfExpressionEx
+
+_collective_ops_warning_printed = False
+_collective_ops_group_key       = 831766147
+_collective_ops_instance_key    = 436340067
+
+class Optimizer:
+    """A Wrapper for tf.train.Optimizer.
+
+    Automatically takes care of:
+    - Gradient averaging for multi-GPU training.
+    - Gradient accumulation for arbitrarily large minibatches.
+    - Dynamic loss scaling and typecasts for FP16 training.
+    - Ignoring corrupted gradients that contain NaNs/Infs.
+    - Reporting statistics.
+    - Well-chosen default settings.
+    """
+
+    def __init__(self,
+        name:                   str             = "Train",                  # Name string that will appear in TensorFlow graph.
+        tf_optimizer:           str             = "tf.train.AdamOptimizer", # Underlying optimizer class.
+        learning_rate:          TfExpressionEx  = 0.001,                    # Learning rate. Can vary over time.
+        minibatch_multiplier:   TfExpressionEx  = None,                     # Treat N consecutive minibatches as one by accumulating gradients.
+        share:                  "Optimizer"     = None,                     # Share internal state with a previously created optimizer?
+        use_loss_scaling:       bool            = False,                    # Enable dynamic loss scaling for robust mixed-precision training?
+        loss_scaling_init:      float           = 64.0,                     # Log2 of initial loss scaling factor.
+        loss_scaling_inc:       float           = 0.0005,                   # Log2 of per-minibatch loss scaling increment when there is no overflow.
+        loss_scaling_dec:       float           = 1.0,                      # Log2 of per-minibatch loss scaling decrement when there is an overflow.
+        report_mem_usage:       bool            = False,                    # Report fine-grained memory usage statistics in TensorBoard?
+        **kwargs):
+
+        # Public fields.
+        self.name                   = name
+        self.learning_rate          = learning_rate
+        self.minibatch_multiplier   = minibatch_multiplier
+        self.id                     = self.name.replace("/", ".")
+        self.scope                  = tf.get_default_graph().unique_name(self.id)
+        self.optimizer_class        = util.get_obj_by_name(tf_optimizer)
+        self.optimizer_kwargs       = dict(kwargs)
+        self.use_loss_scaling       = use_loss_scaling
+        self.loss_scaling_init      = loss_scaling_init
+        self.loss_scaling_inc       = loss_scaling_inc
+        self.loss_scaling_dec       = loss_scaling_dec
+
+        # Private fields.
+        self._updates_applied       = False
+        self._devices               = OrderedDict() # device_name => EasyDict()
+        self._shared_optimizers     = OrderedDict() # device_name => optimizer_class
+        self._gradient_shapes       = None          # [shape, ...]
+        self._report_mem_usage      = report_mem_usage
+
+        # Validate arguments.
+        assert callable(self.optimizer_class)
+
+        # Share internal state if requested.
+        if share is not None:
+            assert isinstance(share, Optimizer)
+            assert self.optimizer_class is share.optimizer_class
+            assert self.learning_rate is share.learning_rate
+            assert self.optimizer_kwargs == share.optimizer_kwargs
+            self._shared_optimizers = share._shared_optimizers # pylint: disable=protected-access
+
+    def _get_device(self, device_name: str):
+        """Get internal state for the given TensorFlow device."""
+        tfutil.assert_tf_initialized()
+        if device_name in self._devices:
+            return self._devices[device_name]
+
+        # Initialize fields.
+        device = util.EasyDict()
+        device.name             = device_name
+        device.optimizer        = None          # Underlying optimizer:     optimizer_class
+        device.loss_scaling_var = None          # Log2 of loss scaling:     tf.Variable
+        device.grad_raw         = OrderedDict() # Raw gradients:            var => [grad, ...]
+        device.grad_clean       = OrderedDict() # Clean gradients:          var => grad
+        device.grad_acc_vars    = OrderedDict() # Accumulation sums:        var => tf.Variable
+        device.grad_acc_count   = None          # Accumulation counter:     tf.Variable
+        device.grad_acc         = OrderedDict() # Accumulated gradients:    var => grad
+
+        # Setup TensorFlow objects.
+        with tfutil.absolute_name_scope(self.scope + "/Devices"), tf.device(device_name), tf.control_dependencies(None):
+            if device_name not in self._shared_optimizers:
+                optimizer_name = self.scope.replace("/", "_") + "_opt%d" % len(self._shared_optimizers)
+                self._shared_optimizers[device_name] = self.optimizer_class(name=optimizer_name, learning_rate=self.learning_rate, **self.optimizer_kwargs)
+            device.optimizer = self._shared_optimizers[device_name]
+            if self.use_loss_scaling:
+                device.loss_scaling_var = tf.Variable(np.float32(self.loss_scaling_init), trainable=False, name="loss_scaling_var")
+
+        # Register device.
+        self._devices[device_name] = device
+        return device
+
+    def register_gradients(self, loss: TfExpression, trainable_vars: Union[List, dict]) -> None:
+        """Register the gradients of the given loss function with respect to the given variables.
+        Intended to be called once per GPU."""
+        tfutil.assert_tf_initialized()
+        assert not self._updates_applied
+        device = self._get_device(loss.device)
+
+        # Validate trainables.
+        if isinstance(trainable_vars, dict):
+            trainable_vars = list(trainable_vars.values())  # allow passing in Network.trainables as vars
+        assert isinstance(trainable_vars, list) and len(trainable_vars) >= 1
+        assert all(tfutil.is_tf_expression(expr) for expr in trainable_vars + [loss])
+        assert all(var.device == device.name for var in trainable_vars)
+
+        # Validate shapes.
+        if self._gradient_shapes is None:
+            self._gradient_shapes = [var.shape.as_list() for var in trainable_vars]
+        assert len(trainable_vars) == len(self._gradient_shapes)
+        assert all(var.shape.as_list() == var_shape for var, var_shape in zip(trainable_vars, self._gradient_shapes))
+
+        # Report memory usage if requested.
+        deps = [loss]
+        if self._report_mem_usage:
+            self._report_mem_usage = False
+            try:
+                with tf.name_scope(self.id + '_mem'), tf.device(device.name), tf.control_dependencies([loss]):
+                    deps.append(autosummary.autosummary(self.id + "/mem_usage_gb", tf.contrib.memory_stats.BytesInUse() / 2**30))
+            except tf.errors.NotFoundError:
+                pass
+
+        # Compute gradients.
+        with tf.name_scope(self.id + "_grad"), tf.device(device.name), tf.control_dependencies(deps):
+            loss = self.apply_loss_scaling(tf.cast(loss, tf.float32))
+            gate = tf.train.Optimizer.GATE_NONE  # disable gating to reduce memory usage
+            grad_list = device.optimizer.compute_gradients(loss=loss, var_list=trainable_vars, gate_gradients=gate)
+
+        # Register gradients.
+        for grad, var in grad_list:
+            if var not in device.grad_raw:
+                device.grad_raw[var] = []
+            device.grad_raw[var].append(grad)
+
+    def apply_updates(self, allow_no_op: bool = False) -> tf.Operation:
+        """Construct training op to update the registered variables based on their gradients."""
+        tfutil.assert_tf_initialized()
+        assert not self._updates_applied
+        self._updates_applied = True
+        all_ops = []
+
+        # Check for no-op.
+        if allow_no_op and len(self._devices) == 0:
+            with tfutil.absolute_name_scope(self.scope):
+                return tf.no_op(name='TrainingOp')
+
+        # Clean up gradients.
+        for device_idx, device in enumerate(self._devices.values()):
+            with tfutil.absolute_name_scope(self.scope + "/Clean%d" % device_idx), tf.device(device.name):
+                for var, grad in device.grad_raw.items():
+
+                    # Filter out disconnected gradients and convert to float32.
+                    grad = [g for g in grad if g is not None]
+                    grad = [tf.cast(g, tf.float32) for g in grad]
+
+                    # Sum within the device.
+                    if len(grad) == 0:
+                        grad = tf.zeros(var.shape)  # No gradients => zero.
+                    elif len(grad) == 1:
+                        grad = grad[0]              # Single gradient => use as is.
+                    else:
+                        grad = tf.add_n(grad)       # Multiple gradients => sum.
+
+                    # Scale as needed.
+                    scale = 1.0 / len(device.grad_raw[var]) / len(self._devices)
+                    scale = tf.constant(scale, dtype=tf.float32, name="scale")
+                    if self.minibatch_multiplier is not None:
+                        scale /= tf.cast(self.minibatch_multiplier, tf.float32)
+                    scale = self.undo_loss_scaling(scale)
+                    device.grad_clean[var] = grad * scale
+
+        # Sum gradients across devices.
+        if len(self._devices) > 1:
+            with tfutil.absolute_name_scope(self.scope + "/Broadcast"), tf.device(None):
+                if platform.system() == "Windows":    # Windows => NCCL ops are not available.
+                    self._broadcast_fallback()
+                elif tf.VERSION.startswith("1.15."):  # TF 1.15 => NCCL ops are broken: https://github.com/tensorflow/tensorflow/issues/41539
+                    self._broadcast_fallback()
+                else:                                 # Otherwise => NCCL ops are safe to use.
+                    self._broadcast_nccl()
+
+        # Apply updates separately on each device.
+        for device_idx, device in enumerate(self._devices.values()):
+            with tfutil.absolute_name_scope(self.scope + "/Apply%d" % device_idx), tf.device(device.name):
+                # pylint: disable=cell-var-from-loop
+
+                # Accumulate gradients over time.
+                if self.minibatch_multiplier is None:
+                    acc_ok = tf.constant(True, name='acc_ok')
+                    device.grad_acc = OrderedDict(device.grad_clean)
+                else:
+                    # Create variables.
+                    with tf.control_dependencies(None):
+                        for var in device.grad_clean.keys():
+                            device.grad_acc_vars[var] = tf.Variable(tf.zeros(var.shape), trainable=False, name="grad_acc_var")
+                        device.grad_acc_count = tf.Variable(tf.zeros([]), trainable=False, name="grad_acc_count")
+
+                    # Track counter.
+                    count_cur = device.grad_acc_count + 1.0
+                    count_inc_op = lambda: tf.assign(device.grad_acc_count, count_cur)
+                    count_reset_op = lambda: tf.assign(device.grad_acc_count, tf.zeros([]))
+                    acc_ok = (count_cur >= tf.cast(self.minibatch_multiplier, tf.float32))
+                    all_ops.append(tf.cond(acc_ok, count_reset_op, count_inc_op))
+
+                    # Track gradients.
+                    for var, grad in device.grad_clean.items():
+                        acc_var = device.grad_acc_vars[var]
+                        acc_cur = acc_var + grad
+                        device.grad_acc[var] = acc_cur
+                        with tf.control_dependencies([acc_cur]):
+                            acc_inc_op = lambda: tf.assign(acc_var, acc_cur)
+                            acc_reset_op = lambda: tf.assign(acc_var, tf.zeros(var.shape))
+                            all_ops.append(tf.cond(acc_ok, acc_reset_op, acc_inc_op))
+
+                # No overflow => apply gradients.
+                all_ok = tf.reduce_all(tf.stack([acc_ok] + [tf.reduce_all(tf.is_finite(g)) for g in device.grad_acc.values()]))
+                apply_op = lambda: device.optimizer.apply_gradients([(tf.cast(grad, var.dtype), var) for var, grad in device.grad_acc.items()])
+                all_ops.append(tf.cond(all_ok, apply_op, tf.no_op))
+
+                # Adjust loss scaling.
+                if self.use_loss_scaling:
+                    ls_inc_op = lambda: tf.assign_add(device.loss_scaling_var, self.loss_scaling_inc)
+                    ls_dec_op = lambda: tf.assign_sub(device.loss_scaling_var, self.loss_scaling_dec)
+                    ls_update_op = lambda: tf.group(tf.cond(all_ok, ls_inc_op, ls_dec_op))
+                    all_ops.append(tf.cond(acc_ok, ls_update_op, tf.no_op))
+
+                # Last device => report statistics.
+                if device_idx == len(self._devices) - 1:
+                    all_ops.append(autosummary.autosummary(self.id + "/learning_rate", tf.convert_to_tensor(self.learning_rate)))
+                    all_ops.append(autosummary.autosummary(self.id + "/overflow_frequency", tf.where(all_ok, 0, 1), condition=acc_ok))
+                    if self.use_loss_scaling:
+                        all_ops.append(autosummary.autosummary(self.id + "/loss_scaling_log2", device.loss_scaling_var))
+
+        # Initialize variables.
+        self.reset_optimizer_state()
+        if self.use_loss_scaling:
+            tfutil.init_uninitialized_vars([device.loss_scaling_var for device in self._devices.values()])
+        if self.minibatch_multiplier is not None:
+            tfutil.run([var.initializer for device in self._devices.values() for var in list(device.grad_acc_vars.values()) + [device.grad_acc_count]])
+
+        # Group everything into a single op.
+        with tfutil.absolute_name_scope(self.scope):
+            return tf.group(*all_ops, name="TrainingOp")
+
+    def reset_optimizer_state(self) -> None:
+        """Reset internal state of the underlying optimizer."""
+        tfutil.assert_tf_initialized()
+        tfutil.run([var.initializer for device in self._devices.values() for var in device.optimizer.variables()])
+
+    def get_loss_scaling_var(self, device: str) -> Union[tf.Variable, None]:
+        """Get or create variable representing log2 of the current dynamic loss scaling factor."""
+        return self._get_device(device).loss_scaling_var
+
+    def apply_loss_scaling(self, value: TfExpression) -> TfExpression:
+        """Apply dynamic loss scaling for the given expression."""
+        assert tfutil.is_tf_expression(value)
+        if not self.use_loss_scaling:
+            return value
+        return value * tfutil.exp2(self.get_loss_scaling_var(value.device))
+
+    def undo_loss_scaling(self, value: TfExpression) -> TfExpression:
+        """Undo the effect of dynamic loss scaling for the given expression."""
+        assert tfutil.is_tf_expression(value)
+        if not self.use_loss_scaling:
+            return value
+        return value * tfutil.exp2(-self.get_loss_scaling_var(value.device)) # pylint: disable=invalid-unary-operand-type
+
+    def _broadcast_nccl(self):
+        """Sum gradients across devices using NCCL ops (fast path)."""
+        from tensorflow.python.ops import nccl_ops # pylint: disable=no-name-in-module
+        for all_vars in zip(*[device.grad_clean.keys() for device in self._devices.values()]):
+            if any(x.shape.num_elements() > 0 for x in all_vars):
+                all_grads = [device.grad_clean[var] for device, var in zip(self._devices.values(), all_vars)]
+                all_grads = nccl_ops.all_sum(all_grads)
+                for device, var, grad in zip(self._devices.values(), all_vars, all_grads):
+                    device.grad_clean[var] = grad
+
+    def _broadcast_fallback(self):
+        """Sum gradients across devices using TensorFlow collective ops (slow fallback path)."""
+        from tensorflow.python.ops import collective_ops # pylint: disable=no-name-in-module
+        global _collective_ops_warning_printed, _collective_ops_group_key, _collective_ops_instance_key
+        if all(x.shape.num_elements() == 0 for device in self._devices.values() for x in device.grad_clean.values()):
+            return
+        if not _collective_ops_warning_printed:
+            print("------------------------------------------------------------------------")
+            print("WARNING: Using slow fallback implementation for inter-GPU communication.")
+            print("Please use TensorFlow 1.14 on Linux for optimal training performance.")
+            print("------------------------------------------------------------------------")
+            _collective_ops_warning_printed = True
+        for device in self._devices.values():
+            with tf.device(device.name):
+                combo = [tf.reshape(x, [x.shape.num_elements()]) for x in device.grad_clean.values()]
+                combo = tf.concat(combo, axis=0)
+                combo = collective_ops.all_reduce(combo, merge_op='Add', final_op='Id',
+                    group_size=len(self._devices), group_key=_collective_ops_group_key,
+                    instance_key=_collective_ops_instance_key)
+                cur_ofs = 0
+                for var, grad_old in device.grad_clean.items():
+                    grad_new = tf.reshape(combo[cur_ofs : cur_ofs + grad_old.shape.num_elements()], grad_old.shape)
+                    cur_ofs += grad_old.shape.num_elements()
+                    device.grad_clean[var] = grad_new
+        _collective_ops_instance_key += 1
+
+
+class SimpleAdam:
+    """Simplified version of tf.train.AdamOptimizer that behaves identically when used with dnnlib.tflib.Optimizer."""
+
+    def __init__(self, name="Adam", learning_rate=0.001, beta1=0.9, beta2=0.999, epsilon=1e-8):
+        self.name = name
+        self.learning_rate = learning_rate
+        self.beta1 = beta1
+        self.beta2 = beta2
+        self.epsilon = epsilon
+        self.all_state_vars = []
+
+    def variables(self):
+        return self.all_state_vars
+
+    def compute_gradients(self, loss, var_list, gate_gradients=tf.train.Optimizer.GATE_NONE):
+        assert gate_gradients == tf.train.Optimizer.GATE_NONE
+        return list(zip(tf.gradients(loss, var_list), var_list))
+
+    def apply_gradients(self, grads_and_vars):
+        with tf.name_scope(self.name):
+            state_vars = []
+            update_ops = []
+
+            # Adjust learning rate to deal with startup bias.
+            with tf.control_dependencies(None):
+                b1pow_var = tf.Variable(dtype=tf.float32, initial_value=1, trainable=False)
+                b2pow_var = tf.Variable(dtype=tf.float32, initial_value=1, trainable=False)
+                state_vars += [b1pow_var, b2pow_var]
+            b1pow_new = b1pow_var * self.beta1
+            b2pow_new = b2pow_var * self.beta2
+            update_ops += [tf.assign(b1pow_var, b1pow_new), tf.assign(b2pow_var, b2pow_new)]
+            lr_new = self.learning_rate * tf.sqrt(1 - b2pow_new) / (1 - b1pow_new)
+
+            # Construct ops to update each variable.
+            for grad, var in grads_and_vars:
+                with tf.control_dependencies(None):
+                    m_var = tf.Variable(dtype=tf.float32, initial_value=tf.zeros_like(var), trainable=False)
+                    v_var = tf.Variable(dtype=tf.float32, initial_value=tf.zeros_like(var), trainable=False)
+                    state_vars += [m_var, v_var]
+                m_new = self.beta1 * m_var + (1 - self.beta1) * grad
+                v_new = self.beta2 * v_var + (1 - self.beta2) * tf.square(grad)
+                var_delta = lr_new * m_new / (tf.sqrt(v_new) + self.epsilon)
+                update_ops += [tf.assign(m_var, m_new), tf.assign(v_var, v_new), tf.assign_sub(var, var_delta)]
+
+            # Group everything together.
+            self.all_state_vars += state_vars
+            return tf.group(*update_ops)

dnnlib/tflib/tfutil.py

@@ -0,0 +1,264 @@
+# Copyright (c) 2020, NVIDIA CORPORATION.  All rights reserved.
+#
+# NVIDIA CORPORATION and its licensors retain all intellectual property
+# and proprietary rights in and to this software, related documentation
+# and any modifications thereto.  Any use, reproduction, disclosure or
+# distribution of this software and related documentation without an express
+# license agreement from NVIDIA CORPORATION is strictly prohibited.
+
+"""Miscellaneous helper utils for Tensorflow."""
+
+import os
+import numpy as np
+import tensorflow as tf
+
+# Silence deprecation warnings from TensorFlow 1.13 onwards
+import logging
+logging.getLogger('tensorflow').setLevel(logging.ERROR)
+import tensorflow.contrib   # requires TensorFlow 1.x!
+tf.contrib = tensorflow.contrib
+
+from typing import Any, Iterable, List, Union
+
+TfExpression = Union[tf.Tensor, tf.Variable, tf.Operation]
+"""A type that represents a valid Tensorflow expression."""
+
+TfExpressionEx = Union[TfExpression, int, float, np.ndarray]
+"""A type that can be converted to a valid Tensorflow expression."""
+
+
+def run(*args, **kwargs) -> Any:
+    """Run the specified ops in the default session."""
+    assert_tf_initialized()
+    return tf.get_default_session().run(*args, **kwargs)
+
+
+def is_tf_expression(x: Any) -> bool:
+    """Check whether the input is a valid Tensorflow expression, i.e., Tensorflow Tensor, Variable, or Operation."""
+    return isinstance(x, (tf.Tensor, tf.Variable, tf.Operation))
+
+
+def shape_to_list(shape: Iterable[tf.Dimension]) -> List[Union[int, None]]:
+    """Convert a Tensorflow shape to a list of ints. Retained for backwards compatibility -- use TensorShape.as_list() in new code."""
+    return [dim.value for dim in shape]
+
+
+def flatten(x: TfExpressionEx) -> TfExpression:
+    """Shortcut function for flattening a tensor."""
+    with tf.name_scope("Flatten"):
+        return tf.reshape(x, [-1])
+
+
+def log2(x: TfExpressionEx) -> TfExpression:
+    """Logarithm in base 2."""
+    with tf.name_scope("Log2"):
+        return tf.log(x) * np.float32(1.0 / np.log(2.0))
+
+
+def exp2(x: TfExpressionEx) -> TfExpression:
+    """Exponent in base 2."""
+    with tf.name_scope("Exp2"):
+        return tf.exp(x * np.float32(np.log(2.0)))
+
+
+def erfinv(y: TfExpressionEx) -> TfExpression:
+    """Inverse of the error function."""
+    # pylint: disable=no-name-in-module
+    from tensorflow.python.ops.distributions import special_math
+    return special_math.erfinv(y)
+
+
+def lerp(a: TfExpressionEx, b: TfExpressionEx, t: TfExpressionEx) -> TfExpressionEx:
+    """Linear interpolation."""
+    with tf.name_scope("Lerp"):
+        return a + (b - a) * t
+
+
+def lerp_clip(a: TfExpressionEx, b: TfExpressionEx, t: TfExpressionEx) -> TfExpression:
+    """Linear interpolation with clip."""
+    with tf.name_scope("LerpClip"):
+        return a + (b - a) * tf.clip_by_value(t, 0.0, 1.0)
+
+
+def absolute_name_scope(scope: str) -> tf.name_scope:
+    """Forcefully enter the specified name scope, ignoring any surrounding scopes."""
+    return tf.name_scope(scope + "/")
+
+
+def absolute_variable_scope(scope: str, **kwargs) -> tf.variable_scope:
+    """Forcefully enter the specified variable scope, ignoring any surrounding scopes."""
+    return tf.variable_scope(tf.VariableScope(name=scope, **kwargs), auxiliary_name_scope=False)
+
+
+def _sanitize_tf_config(config_dict: dict = None) -> dict:
+    # Defaults.
+    cfg = dict()
+    cfg["rnd.np_random_seed"]               = None      # Random seed for NumPy. None = keep as is.
+    cfg["rnd.tf_random_seed"]               = "auto"    # Random seed for TensorFlow. 'auto' = derive from NumPy random state. None = keep as is.
+    cfg["env.TF_CPP_MIN_LOG_LEVEL"]         = "1"       # 0 = Print all available debug info from TensorFlow. 1 = Print warnings and errors, but disable debug info.
+    cfg["env.HDF5_USE_FILE_LOCKING"]        = "FALSE"   # Disable HDF5 file locking to avoid concurrency issues with network shares.
+    cfg["graph_options.place_pruned_graph"] = True      # False = Check that all ops are available on the designated device. True = Skip the check for ops that are not used.
+    cfg["gpu_options.allow_growth"]         = True      # False = Allocate all GPU memory at the beginning. True = Allocate only as much GPU memory as needed.
+
+    # Remove defaults for environment variables that are already set.
+    for key in list(cfg):
+        fields = key.split(".")
+        if fields[0] == "env":
+            assert len(fields) == 2
+            if fields[1] in os.environ:
+                del cfg[key]
+
+    # User overrides.
+    if config_dict is not None:
+        cfg.update(config_dict)
+    return cfg
+
+
+def init_tf(config_dict: dict = None) -> None:
+    """Initialize TensorFlow session using good default settings."""
+    # Skip if already initialized.
+    if tf.get_default_session() is not None:
+        return
+
+    # Setup config dict and random seeds.
+    cfg = _sanitize_tf_config(config_dict)
+    np_random_seed = cfg["rnd.np_random_seed"]
+    if np_random_seed is not None:
+        np.random.seed(np_random_seed)
+    tf_random_seed = cfg["rnd.tf_random_seed"]
+    if tf_random_seed == "auto":
+        tf_random_seed = np.random.randint(1 << 31)
+    if tf_random_seed is not None:
+        tf.set_random_seed(tf_random_seed)
+
+    # Setup environment variables.
+    for key, value in cfg.items():
+        fields = key.split(".")
+        if fields[0] == "env":
+            assert len(fields) == 2
+            os.environ[fields[1]] = str(value)
+
+    # Create default TensorFlow session.
+    create_session(cfg, force_as_default=True)
+
+
+def assert_tf_initialized():
+    """Check that TensorFlow session has been initialized."""
+    if tf.get_default_session() is None:
+        raise RuntimeError("No default TensorFlow session found. Please call dnnlib.tflib.init_tf().")
+
+
+def create_session(config_dict: dict = None, force_as_default: bool = False) -> tf.Session:
+    """Create tf.Session based on config dict."""
+    # Setup TensorFlow config proto.
+    cfg = _sanitize_tf_config(config_dict)
+    config_proto = tf.ConfigProto()
+    for key, value in cfg.items():
+        fields = key.split(".")
+        if fields[0] not in ["rnd", "env"]:
+            obj = config_proto
+            for field in fields[:-1]:
+                obj = getattr(obj, field)
+            setattr(obj, fields[-1], value)
+
+    # Create session.
+    session = tf.Session(config=config_proto)
+    if force_as_default:
+        # pylint: disable=protected-access
+        session._default_session = session.as_default()
+        session._default_session.enforce_nesting = False
+        session._default_session.__enter__()
+    return session
+
+
+def init_uninitialized_vars(target_vars: List[tf.Variable] = None) -> None:
+    """Initialize all tf.Variables that have not already been initialized.
+
+    Equivalent to the following, but more efficient and does not bloat the tf graph:
+    tf.variables_initializer(tf.report_uninitialized_variables()).run()
+    """
+    assert_tf_initialized()
+    if target_vars is None:
+        target_vars = tf.global_variables()
+
+    test_vars = []
+    test_ops = []
+
+    with tf.control_dependencies(None):  # ignore surrounding control_dependencies
+        for var in target_vars:
+            assert is_tf_expression(var)
+
+            try:
+                tf.get_default_graph().get_tensor_by_name(var.name.replace(":0", "/IsVariableInitialized:0"))
+            except KeyError:
+                # Op does not exist => variable may be uninitialized.
+                test_vars.append(var)
+
+                with absolute_name_scope(var.name.split(":")[0]):
+                    test_ops.append(tf.is_variable_initialized(var))
+
+    init_vars = [var for var, inited in zip(test_vars, run(test_ops)) if not inited]
+    run([var.initializer for var in init_vars])
+
+
+def set_vars(var_to_value_dict: dict) -> None:
+    """Set the values of given tf.Variables.
+
+    Equivalent to the following, but more efficient and does not bloat the tf graph:
+    tflib.run([tf.assign(var, value) for var, value in var_to_value_dict.items()]
+    """
+    assert_tf_initialized()
+    ops = []
+    feed_dict = {}
+
+    for var, value in var_to_value_dict.items():
+        assert is_tf_expression(var)
+
+        try:
+            setter = tf.get_default_graph().get_tensor_by_name(var.name.replace(":0", "/setter:0"))  # look for existing op
+        except KeyError:
+            with absolute_name_scope(var.name.split(":")[0]):
+                with tf.control_dependencies(None):  # ignore surrounding control_dependencies
+                    setter = tf.assign(var, tf.placeholder(var.dtype, var.shape, "new_value"), name="setter")  # create new setter
+
+        ops.append(setter)
+        feed_dict[setter.op.inputs[1]] = value
+
+    run(ops, feed_dict)
+
+
+def create_var_with_large_initial_value(initial_value: np.ndarray, *args, **kwargs):
+    """Create tf.Variable with large initial value without bloating the tf graph."""
+    assert_tf_initialized()
+    assert isinstance(initial_value, np.ndarray)
+    zeros = tf.zeros(initial_value.shape, initial_value.dtype)
+    var = tf.Variable(zeros, *args, **kwargs)
+    set_vars({var: initial_value})
+    return var
+
+
+def convert_images_from_uint8(images, drange=[-1,1], nhwc_to_nchw=False):
+    """Convert a minibatch of images from uint8 to float32 with configurable dynamic range.
+    Can be used as an input transformation for Network.run().
+    """
+    images = tf.cast(images, tf.float32)
+    if nhwc_to_nchw:
+        images = tf.transpose(images, [0, 3, 1, 2])
+    return images * ((drange[1] - drange[0]) / 255) + drange[0]
+
+
+def convert_images_to_uint8(images, drange=[-1,1], nchw_to_nhwc=False, shrink=1, uint8_cast=True):
+    """Convert a minibatch of images from float32 to uint8 with configurable dynamic range.
+    Can be used as an output transformation for Network.run().
+    """
+    images = tf.cast(images, tf.float32)
+    if shrink > 1:
+        ksize = [1, 1, shrink, shrink]
+        images = tf.nn.avg_pool(images, ksize=ksize, strides=ksize, padding="VALID", data_format="NCHW")
+    if nchw_to_nhwc:
+        images = tf.transpose(images, [0, 2, 3, 1])
+    scale = 255 / (drange[1] - drange[0])
+    images = images * scale + (0.5 - drange[0] * scale)
+    if uint8_cast:
+        images = tf.saturate_cast(images, tf.uint8)
+    return images

dnnlib/util.py

@@ -0,0 +1,472 @@
+# Copyright (c) 2020, NVIDIA CORPORATION.  All rights reserved.
+#
+# NVIDIA CORPORATION and its licensors retain all intellectual property
+# and proprietary rights in and to this software, related documentation
+# and any modifications thereto.  Any use, reproduction, disclosure or
+# distribution of this software and related documentation without an express
+# license agreement from NVIDIA CORPORATION is strictly prohibited.
+
+"""Miscellaneous utility classes and functions."""
+
+import ctypes
+import fnmatch
+import importlib
+import inspect
+import numpy as np
+import os
+import shutil
+import sys
+import types
+import io
+import pickle
+import re
+import requests
+import html
+import hashlib
+import glob
+import tempfile
+import urllib
+import urllib.request
+import uuid
+
+from distutils.util import strtobool
+from typing import Any, List, Tuple, Union
+
+
+# Util classes
+# ------------------------------------------------------------------------------------------
+
+
+class EasyDict(dict):
+    """Convenience class that behaves like a dict but allows access with the attribute syntax."""
+
+    def __getattr__(self, name: str) -> Any:
+        try:
+            return self[name]
+        except KeyError:
+            raise AttributeError(name)
+
+    def __setattr__(self, name: str, value: Any) -> None:
+        self[name] = value
+
+    def __delattr__(self, name: str) -> None:
+        del self[name]
+
+
+class Logger(object):
+    """Redirect stderr to stdout, optionally print stdout to a file, and optionally force flushing on both stdout and the file."""
+
+    def __init__(self, file_name: str = None, file_mode: str = "w", should_flush: bool = True):
+        self.file = None
+
+        if file_name is not None:
+            self.file = open(file_name, file_mode)
+
+        self.should_flush = should_flush
+        self.stdout = sys.stdout
+        self.stderr = sys.stderr
+
+        sys.stdout = self
+        sys.stderr = self
+
+    def __enter__(self) -> "Logger":
+        return self
+
+    def __exit__(self, exc_type: Any, exc_value: Any, traceback: Any) -> None:
+        self.close()
+
+    def write(self, text: str) -> None:
+        """Write text to stdout (and a file) and optionally flush."""
+        if len(text) == 0: # workaround for a bug in VSCode debugger: sys.stdout.write(''); sys.stdout.flush() => crash
+            return
+
+        if self.file is not None:
+            self.file.write(text)
+
+        self.stdout.write(text)
+
+        if self.should_flush:
+            self.flush()
+
+    def flush(self) -> None:
+        """Flush written text to both stdout and a file, if open."""
+        if self.file is not None:
+            self.file.flush()
+
+        self.stdout.flush()
+
+    def close(self) -> None:
+        """Flush, close possible files, and remove stdout/stderr mirroring."""
+        self.flush()
+
+        # if using multiple loggers, prevent closing in wrong order
+        if sys.stdout is self:
+            sys.stdout = self.stdout
+        if sys.stderr is self:
+            sys.stderr = self.stderr
+
+        if self.file is not None:
+            self.file.close()
+
+
+# Cache directories
+# ------------------------------------------------------------------------------------------
+
+_dnnlib_cache_dir = None
+
+def set_cache_dir(path: str) -> None:
+    global _dnnlib_cache_dir
+    _dnnlib_cache_dir = path
+
+def make_cache_dir_path(*paths: str) -> str:
+    if _dnnlib_cache_dir is not None:
+        return os.path.join(_dnnlib_cache_dir, *paths)
+    if 'DNNLIB_CACHE_DIR' in os.environ:
+        return os.path.join(os.environ['DNNLIB_CACHE_DIR'], *paths)
+    if 'HOME' in os.environ:
+        return os.path.join(os.environ['HOME'], '.cache', 'dnnlib', *paths)
+    if 'USERPROFILE' in os.environ:
+        return os.path.join(os.environ['USERPROFILE'], '.cache', 'dnnlib', *paths)
+    return os.path.join(tempfile.gettempdir(), '.cache', 'dnnlib', *paths)
+
+# Small util functions
+# ------------------------------------------------------------------------------------------
+
+
+def format_time(seconds: Union[int, float]) -> str:
+    """Convert the seconds to human readable string with days, hours, minutes and seconds."""
+    s = int(np.rint(seconds))
+
+    if s < 60:
+        return "{0}s".format(s)
+    elif s < 60 * 60:
+        return "{0}m {1:02}s".format(s // 60, s % 60)
+    elif s < 24 * 60 * 60:
+        return "{0}h {1:02}m {2:02}s".format(s // (60 * 60), (s // 60) % 60, s % 60)
+    else:
+        return "{0}d {1:02}h {2:02}m".format(s // (24 * 60 * 60), (s // (60 * 60)) % 24, (s // 60) % 60)
+
+
+def ask_yes_no(question: str) -> bool:
+    """Ask the user the question until the user inputs a valid answer."""
+    while True:
+        try:
+            print("{0} [y/n]".format(question))
+            return strtobool(input().lower())
+        except ValueError:
+            pass
+
+
+def tuple_product(t: Tuple) -> Any:
+    """Calculate the product of the tuple elements."""
+    result = 1
+
+    for v in t:
+        result *= v
+
+    return result
+
+
+_str_to_ctype = {
+    "uint8": ctypes.c_ubyte,
+    "uint16": ctypes.c_uint16,
+    "uint32": ctypes.c_uint32,
+    "uint64": ctypes.c_uint64,
+    "int8": ctypes.c_byte,
+    "int16": ctypes.c_int16,
+    "int32": ctypes.c_int32,
+    "int64": ctypes.c_int64,
+    "float32": ctypes.c_float,
+    "float64": ctypes.c_double
+}
+
+
+def get_dtype_and_ctype(type_obj: Any) -> Tuple[np.dtype, Any]:
+    """Given a type name string (or an object having a __name__ attribute), return matching Numpy and ctypes types that have the same size in bytes."""
+    type_str = None
+
+    if isinstance(type_obj, str):
+        type_str = type_obj
+    elif hasattr(type_obj, "__name__"):
+        type_str = type_obj.__name__
+    elif hasattr(type_obj, "name"):
+        type_str = type_obj.name
+    else:
+        raise RuntimeError("Cannot infer type name from input")
+
+    assert type_str in _str_to_ctype.keys()
+
+    my_dtype = np.dtype(type_str)
+    my_ctype = _str_to_ctype[type_str]
+
+    assert my_dtype.itemsize == ctypes.sizeof(my_ctype)
+
+    return my_dtype, my_ctype
+
+
+def is_pickleable(obj: Any) -> bool:
+    try:
+        with io.BytesIO() as stream:
+            pickle.dump(obj, stream)
+        return True
+    except:
+        return False
+
+
+# Functionality to import modules/objects by name, and call functions by name
+# ------------------------------------------------------------------------------------------
+
+def get_module_from_obj_name(obj_name: str) -> Tuple[types.ModuleType, str]:
+    """Searches for the underlying module behind the name to some python object.
+    Returns the module and the object name (original name with module part removed)."""
+
+    # allow convenience shorthands, substitute them by full names
+    obj_name = re.sub("^np.", "numpy.", obj_name)
+    obj_name = re.sub("^tf.", "tensorflow.", obj_name)
+
+    # list alternatives for (module_name, local_obj_name)
+    parts = obj_name.split(".")
+    name_pairs = [(".".join(parts[:i]), ".".join(parts[i:])) for i in range(len(parts), 0, -1)]
+
+    # try each alternative in turn
+    for module_name, local_obj_name in name_pairs:
+        try:
+            module = importlib.import_module(module_name) # may raise ImportError
+            get_obj_from_module(module, local_obj_name) # may raise AttributeError
+            return module, local_obj_name
+        except:
+            pass
+
+    # maybe some of the modules themselves contain errors?
+    for module_name, _local_obj_name in name_pairs:
+        try:
+            importlib.import_module(module_name) # may raise ImportError
+        except ImportError:
+            if not str(sys.exc_info()[1]).startswith("No module named '" + module_name + "'"):
+                raise
+
+    # maybe the requested attribute is missing?
+    for module_name, local_obj_name in name_pairs:
+        try:
+            module = importlib.import_module(module_name) # may raise ImportError
+            get_obj_from_module(module, local_obj_name) # may raise AttributeError
+        except ImportError:
+            pass
+
+    # we are out of luck, but we have no idea why
+    raise ImportError(obj_name)
+
+
+def get_obj_from_module(module: types.ModuleType, obj_name: str) -> Any:
+    """Traverses the object name and returns the last (rightmost) python object."""
+    if obj_name == '':
+        return module
+    obj = module
+    for part in obj_name.split("."):
+        obj = getattr(obj, part)
+    return obj
+
+
+def get_obj_by_name(name: str) -> Any:
+    """Finds the python object with the given name."""
+    module, obj_name = get_module_from_obj_name(name)
+    return get_obj_from_module(module, obj_name)
+
+
+def call_func_by_name(*args, func_name: str = None, **kwargs) -> Any:
+    """Finds the python object with the given name and calls it as a function."""
+    assert func_name is not None
+    func_obj = get_obj_by_name(func_name)
+    assert callable(func_obj)
+    return func_obj(*args, **kwargs)
+
+
+def construct_class_by_name(*args, class_name: str = None, **kwargs) -> Any:
+    """Finds the python class with the given name and constructs it with the given arguments."""
+    return call_func_by_name(*args, func_name=class_name, **kwargs)
+
+
+def get_module_dir_by_obj_name(obj_name: str) -> str:
+    """Get the directory path of the module containing the given object name."""
+    module, _ = get_module_from_obj_name(obj_name)
+    return os.path.dirname(inspect.getfile(module))
+
+
+def is_top_level_function(obj: Any) -> bool:
+    """Determine whether the given object is a top-level function, i.e., defined at module scope using 'def'."""
+    return callable(obj) and obj.__name__ in sys.modules[obj.__module__].__dict__
+
+
+def get_top_level_function_name(obj: Any) -> str:
+    """Return the fully-qualified name of a top-level function."""
+    assert is_top_level_function(obj)
+    module = obj.__module__
+    if module == '__main__':
+        module = os.path.splitext(os.path.basename(sys.modules[module].__file__))[0]
+    return module + "." + obj.__name__
+
+
+# File system helpers
+# ------------------------------------------------------------------------------------------
+
+def list_dir_recursively_with_ignore(dir_path: str, ignores: List[str] = None, add_base_to_relative: bool = False) -> List[Tuple[str, str]]:
+    """List all files recursively in a given directory while ignoring given file and directory names.
+    Returns list of tuples containing both absolute and relative paths."""
+    assert os.path.isdir(dir_path)
+    base_name = os.path.basename(os.path.normpath(dir_path))
+
+    if ignores is None:
+        ignores = []
+
+    result = []
+
+    for root, dirs, files in os.walk(dir_path, topdown=True):
+        for ignore_ in ignores:
+            dirs_to_remove = [d for d in dirs if fnmatch.fnmatch(d, ignore_)]
+
+            # dirs need to be edited in-place
+            for d in dirs_to_remove:
+                dirs.remove(d)
+
+            files = [f for f in files if not fnmatch.fnmatch(f, ignore_)]
+
+        absolute_paths = [os.path.join(root, f) for f in files]
+        relative_paths = [os.path.relpath(p, dir_path) for p in absolute_paths]
+
+        if add_base_to_relative:
+            relative_paths = [os.path.join(base_name, p) for p in relative_paths]
+
+        assert len(absolute_paths) == len(relative_paths)
+        result += zip(absolute_paths, relative_paths)
+
+    return result
+
+
+def copy_files_and_create_dirs(files: List[Tuple[str, str]]) -> None:
+    """Takes in a list of tuples of (src, dst) paths and copies files.
+    Will create all necessary directories."""
+    for file in files:
+        target_dir_name = os.path.dirname(file[1])
+
+        # will create all intermediate-level directories
+        if not os.path.exists(target_dir_name):
+            os.makedirs(target_dir_name)
+
+        shutil.copyfile(file[0], file[1])
+
+
+# URL helpers
+# ------------------------------------------------------------------------------------------
+
+def is_url(obj: Any, allow_file_urls: bool = False) -> bool:
+    """Determine whether the given object is a valid URL string."""
+    if not isinstance(obj, str) or not "://" in obj:
+        return False
+    if allow_file_urls and obj.startswith('file://'):
+        return True
+    try:
+        res = requests.compat.urlparse(obj)
+        if not res.scheme or not res.netloc or not "." in res.netloc:
+            return False
+        res = requests.compat.urlparse(requests.compat.urljoin(obj, "/"))
+        if not res.scheme or not res.netloc or not "." in res.netloc:
+            return False
+    except:
+        return False
+    return True
+
+
+def open_url(url: str, cache_dir: str = None, num_attempts: int = 10, verbose: bool = True, return_filename: bool = False, cache: bool = True) -> Any:
+    """Download the given URL and return a binary-mode file object to access the data."""
+    assert num_attempts >= 1
+    assert not (return_filename and (not cache))
+
+    # Doesn't look like an URL scheme so interpret it as a local filename.
+    if not re.match('^[a-z]+://', url):
+        return url if return_filename else open(url, "rb")
+
+    # Handle file URLs.  This code handles unusual file:// patterns that
+    # arise on Windows:
+    #
+    # file:///c:/foo.txt
+    #
+    # which would translate to a local '/c:/foo.txt' filename that's
+    # invalid.  Drop the forward slash for such pathnames.
+    #
+    # If you touch this code path, you should test it on both Linux and
+    # Windows.
+    #
+    # Some internet resources suggest using urllib.request.url2pathname() but
+    # but that converts forward slashes to backslashes and this causes
+    # its own set of problems.
+    if url.startswith('file://'):
+        filename = urllib.parse.urlparse(url).path
+        if re.match(r'^/[a-zA-Z]:', filename):
+            filename = filename[1:]
+        return filename if return_filename else open(filename, "rb")
+
+    assert is_url(url)
+
+    # Lookup from cache.
+    if cache_dir is None:
+        cache_dir = make_cache_dir_path('downloads')
+
+    url_md5 = hashlib.md5(url.encode("utf-8")).hexdigest()
+    if cache:
+        cache_files = glob.glob(os.path.join(cache_dir, url_md5 + "_*"))
+        if len(cache_files) == 1:
+            filename = cache_files[0]
+            return filename if return_filename else open(filename, "rb")
+
+    # Download.
+    url_name = None
+    url_data = None
+    with requests.Session() as session:
+        if verbose:
+            print("Downloading %s ..." % url, end="", flush=True)
+        for attempts_left in reversed(range(num_attempts)):
+            try:
+                with session.get(url) as res:
+                    res.raise_for_status()
+                    if len(res.content) == 0:
+                        raise IOError("No data received")
+
+                    if len(res.content) < 8192:
+                        content_str = res.content.decode("utf-8")
+                        if "download_warning" in res.headers.get("Set-Cookie", ""):
+                            links = [html.unescape(link) for link in content_str.split('"') if "export=download" in link]
+                            if len(links) == 1:
+                                url = requests.compat.urljoin(url, links[0])
+                                raise IOError("Google Drive virus checker nag")
+                        if "Google Drive - Quota exceeded" in content_str:
+                            raise IOError("Google Drive download quota exceeded -- please try again later")
+
+                    match = re.search(r'filename="([^"]*)"', res.headers.get("Content-Disposition", ""))
+                    url_name = match[1] if match else url
+                    url_data = res.content
+                    if verbose:
+                        print(" done")
+                    break
+            except:
+                if not attempts_left:
+                    if verbose:
+                        print(" failed")
+                    raise
+                if verbose:
+                    print(".", end="", flush=True)
+
+    # Save to cache.
+    if cache:
+        safe_name = re.sub(r"[^0-9a-zA-Z-._]", "_", url_name)
+        cache_file = os.path.join(cache_dir, url_md5 + "_" + safe_name)
+        temp_file = os.path.join(cache_dir, "tmp_" + uuid.uuid4().hex + "_" + url_md5 + "_" + safe_name)
+        os.makedirs(cache_dir, exist_ok=True)
+        with open(temp_file, "wb") as f:
+            f.write(url_data)
+        os.replace(temp_file, cache_file) # atomic
+        if return_filename:
+            return cache_file
+
+    # Return data as file object.
+    assert not return_filename
+    return io.BytesIO(url_data)

gallery/gallery.md

@@ -0,0 +1,15 @@
+## Gallery
+
+### Mosaics 
+![alt text](gl-mosaics1.png)
+![alt text](gl-mosaics2.png)
+![alt text](gl-mosaics3.png)
+![alt text](gl-mosaics4.png)
+![alt text](gl-mosaics5.png)
+![alt text](gl-mosaics6.png)
+![alt text](gl-mosaics7.png)
+![alt text](gl-mosaics8.png)
+![alt text](gl-mosaics9.png)
+![alt text](gl-mosaics10.png)
+
+

generate.py

@@ -0,0 +1,700 @@
+# Copyright (c) 2020, NVIDIA CORPORATION.  All rights reserved.
+#
+# NVIDIA CORPORATION and its licensors retain all intellectual property
+# and proprietary rights in and to this software, related documentation
+# and any modifications thereto.  Any use, reproduction, disclosure or
+# distribution of this software and related documentation without an express
+# license agreement from NVIDIA CORPORATION is strictly prohibited.
+
+
+"""Generate images using pretrained network pickle."""
+
+import argparse
+import sys
+import os
+import subprocess
+import pickle
+import re
+
+import scipy
+import numpy as np
+import PIL.Image
+
+import dnnlib
+import dnnlib.tflib as tflib
+
+os.environ['PYGAME_HIDE_SUPPORT_PROMPT'] = "hide"
+import moviepy.editor
+from opensimplex import OpenSimplex
+
+import warnings # mostly numpy warnings for me
+warnings.filterwarnings('ignore', category=FutureWarning)
+warnings.filterwarnings('ignore', category=DeprecationWarning)
+
+#----------------------------------------------------------------------------
+
+def create_image_grid(images, grid_size=None):
+    '''
+    Args:
+        images (np.array): images to place on the grid
+        grid_size (tuple(int, int)): size of grid (grid_w, grid_h)
+    Returns:
+        grid (np.array): image grid of size grid_size
+    '''
+    # Some sanity check:
+    assert images.ndim == 3 or images.ndim == 4
+    num, img_h, img_w = images.shape[0], images.shape[1], images.shape[2]
+    if grid_size is not None:
+        grid_w, grid_h = tuple(grid_size)
+    else:
+        grid_w = max(int(np.ceil(np.sqrt(num))), 1)
+        grid_h = max((num - 1) // grid_w + 1, 1)
+    # Get the grid
+    grid = np.zeros(
+        [grid_h * img_h, grid_w * img_w] + list(images.shape[-1:]), dtype=images.dtype
+    )
+    for idx in range(num):
+        x = (idx % grid_w) * img_w
+        y = (idx // grid_w) * img_h
+        grid[y : y + img_h, x : x + img_w, ...] = images[idx]
+    return grid
+
+#----------------------------------------------------------------------------
+
+def generate_images(network_pkl, seeds, truncation_psi, outdir, class_idx=None, dlatents_npz=None, grid=False):
+    tflib.init_tf()
+    print('Loading networks from "%s"...' % network_pkl)
+    with dnnlib.util.open_url(network_pkl) as fp:
+        _G, _D, Gs = pickle.load(fp)
+
+    os.makedirs(outdir, exist_ok=True)
+
+    # Render images for a given dlatent vector.
+    if dlatents_npz is not None:
+        print(f'Generating images from dlatents file "{dlatents_npz}"')
+        dlatents = np.load(dlatents_npz)['dlatents']
+        max_l = 2 * int(np.log2(Gs.output_shape[-1]) - 1)  # max_l=18 for 1024x1024 models
+        if dlatents.shape[1:] != (max_l, 512):  # [N, max_l, 512]
+            actual_size = int(2**(dlatents.shape[1]//2+1))
+            print(f'''Mismatch of loaded dlatents and network! dlatents was created with network of size: {actual_size}\n
+                   {network_pkl} is of size {Gs.output_shape[-1]}''')
+            sys.exit(1)
+        imgs = Gs.components.synthesis.run(dlatents, output_transform=dict(func=tflib.convert_images_to_uint8, nchw_to_nhwc=True))
+        for i, img in enumerate(imgs):
+            fname = f'{outdir}/dlatent{i:02d}.png'
+            print (f'Saved {fname}')
+            PIL.Image.fromarray(img, 'RGB').save(fname)
+        return
+
+    # Render images for dlatents initialized from random seeds.
+    Gs_kwargs = {
+        'output_transform': dict(func=tflib.convert_images_to_uint8, nchw_to_nhwc=True),
+        'randomize_noise': False
+    }
+    if truncation_psi is not None:
+        Gs_kwargs['truncation_psi'] = truncation_psi
+
+    noise_vars = [var for name, var in Gs.components.synthesis.vars.items() if name.startswith('noise')]
+    label = np.zeros([1] + Gs.input_shapes[1][1:])
+    if class_idx is not None:
+        label[:, class_idx] = 1
+
+    images = []
+    for seed_idx, seed in enumerate(seeds):
+        print('Generating image for seed %d (%d/%d) ...' % (seed, seed_idx, len(seeds)))
+        rnd = np.random.RandomState(seed)
+        z = rnd.randn(1, *Gs.input_shape[1:]) # [minibatch, component]
+        tflib.set_vars({var: rnd.randn(*var.shape.as_list()) for var in noise_vars}) # [height, width]
+        image = Gs.run(z, label, **Gs_kwargs) # [minibatch, height, width, channel]
+        images.append(image[0])
+        PIL.Image.fromarray(image[0], 'RGB').save(f'{outdir}/seed{seed:04d}.png')
+
+    # If user wants to save a grid of the generated images
+    if grid:
+        print('Generating image grid...')
+        PIL.Image.fromarray(create_image_grid(np.array(images)), 'RGB').save(f'{outdir}/grid.png')
+
+#----------------------------------------------------------------------------
+
+def truncation_traversal(network_pkl,npys,outdir,class_idx=None, seed=[0],start=-1.0,stop=1.0,increment=0.1,framerate=24):
+    tflib.init_tf()
+    print('Loading networks from "%s"...' % network_pkl)
+    with dnnlib.util.open_url(network_pkl) as fp:
+        _G, _D, Gs = pickle.load(fp)
+
+    os.makedirs(outdir, exist_ok=True)
+
+    Gs_kwargs = {
+        'output_transform': dict(func=tflib.convert_images_to_uint8, nchw_to_nhwc=True),
+        'randomize_noise': False
+    }
+
+    noise_vars = [var for name, var in Gs.components.synthesis.vars.items() if name.startswith('noise')]
+    label = np.zeros([1] + Gs.input_shapes[1][1:])
+    if class_idx is not None:
+        label[:, class_idx] = 1
+
+    count = 1
+    trunc = start
+
+    images = []
+    while trunc <= stop:
+        Gs_kwargs['truncation_psi'] = trunc
+        print('Generating truncation %0.2f' % trunc)
+
+        rnd = np.random.RandomState(seed)
+        z = rnd.randn(1, *Gs.input_shape[1:]) # [minibatch, component]
+        tflib.set_vars({var: rnd.randn(*var.shape.as_list()) for var in noise_vars}) # [height, width]
+        image = Gs.run(z, label, **Gs_kwargs) # [minibatch, height, width, channel]
+        images.append(image[0])
+        PIL.Image.fromarray(image[0], 'RGB').save(f'{outdir}/frame{count:05d}.png')
+
+        trunc+=increment
+        count+=1
+
+    cmd="ffmpeg -y -r {} -i {}/frame%05d.png -vcodec libx264 -pix_fmt yuv420p {}/truncation-traversal-seed{}-start{}-stop{}.mp4".format(framerate,outdir,outdir,seed[0],start,stop)
+    subprocess.call(cmd, shell=True)
+
+#----------------------------------------------------------------------------
+
+def valmap(value, istart, istop, ostart, ostop):
+  return ostart + (ostop - ostart) * ((value - istart) / (istop - istart))
+
+class OSN():
+  min=-1
+  max= 1
+
+  def __init__(self,seed,diameter):
+    self.tmp = OpenSimplex(seed)
+    self.d = diameter
+    self.x = 0
+    self.y = 0
+
+  def get_val(self,angle):
+    self.xoff = valmap(np.cos(angle), -1, 1, self.x, self.x + self.d);
+    self.yoff = valmap(np.sin(angle), -1, 1, self.y, self.y + self.d);
+    return self.tmp.noise2d(self.xoff,self.yoff)
+
+def get_noiseloop(endpoints, nf, d, start_seed):
+    features = []
+    zs = []
+    for i in range(512):
+      features.append(OSN(i+start_seed,d))
+
+    inc = (np.pi*2)/nf
+    for f in range(nf):
+      z = np.random.randn(1, 512)
+      for i in range(512):
+        z[0,i] = features[i].get_val(inc*f)
+      zs.append(z)
+
+    return zs
+
+def line_interpolate(zs, steps):
+   out = []
+   for i in range(len(zs)-1):
+    for index in range(steps):
+     fraction = index/float(steps)
+     out.append(zs[i+1]*fraction + zs[i]*(1-fraction))
+   return out
+
+def generate_zs_from_seeds(seeds,Gs):
+    zs = []
+    for seed_idx, seed in enumerate(seeds):
+        rnd = np.random.RandomState(seed)
+        z = rnd.randn(1, *Gs.input_shape[1:]) # [minibatch, component]
+        zs.append(z)
+    return zs
+
+def convertZtoW(latent, truncation_psi=0.7, truncation_cutoff=9):
+    dlatent = Gs.components.mapping.run(latent, None) # [seed, layer, component]
+    dlatent_avg = Gs.get_var('dlatent_avg') # [component]
+    dlatent = dlatent_avg + (dlatent - dlatent_avg) * truncation_psi
+
+    return dlatent
+
+def generate_latent_images(zs, truncation_psi, outdir, save_npy,prefix,vidname,framerate):
+    Gs_kwargs = {
+        'output_transform': dict(func=tflib.convert_images_to_uint8, nchw_to_nhwc=True),
+        'randomize_noise': False
+    }
+
+    if not isinstance(truncation_psi, list):
+        truncation_psi = [truncation_psi] * len(zs)
+
+    for z_idx, z in enumerate(zs):
+        if isinstance(z,list):
+          z = np.array(z).reshape(1,512)
+        elif isinstance(z,np.ndarray):
+          z.reshape(1,512)
+        print('Generating image for step %d/%d ...' % (z_idx, len(zs)))
+        Gs_kwargs['truncation_psi'] = truncation_psi[z_idx]
+        noise_rnd = np.random.RandomState(1) # fix noise
+        tflib.set_vars({var: noise_rnd.randn(*var.shape.as_list()) for var in noise_vars}) # [height, width]
+        images = Gs.run(z, None, **Gs_kwargs) # [minibatch, height, width, channel]
+        PIL.Image.fromarray(images[0], 'RGB').save(f'{outdir}/frames/{prefix}{z_idx:05d}.png')
+        if save_npy:
+          np.save(dnnlib.make_run_dir_path('%s%05d.npy' % (prefix,z_idx)), z)
+
+    cmd="ffmpeg -y -r {} -i {}/frames/{}%05d.png -vcodec libx264 -pix_fmt yuv420p {}/walk-{}-{}fps.mp4".format(framerate,outdir,prefix,outdir,vidname,framerate)
+    subprocess.call(cmd, shell=True)
+
+def generate_images_in_w_space(ws, truncation_psi,outdir,save_npy,prefix,vidname,framerate):
+
+    Gs_kwargs = {
+        'output_transform': dict(func=tflib.convert_images_to_uint8, nchw_to_nhwc=True),
+        'randomize_noise': False,
+        'truncation_psi': truncation_psi
+    }
+
+    for w_idx, w in enumerate(ws):
+        print('Generating image for step %d/%d ...' % (w_idx, len(ws)))
+        noise_rnd = np.random.RandomState(1) # fix noise
+        tflib.set_vars({var: noise_rnd.randn(*var.shape.as_list()) for var in noise_vars}) # [height, width]
+        images = Gs.components.synthesis.run(w, **Gs_kwargs) # [minibatch, height, width, channel]
+        PIL.Image.fromarray(images[0], 'RGB').save(f'{outdir}/frames/{prefix}{w_idx:05d}.png')
+        if save_npy:
+          np.save(dnnlib.make_run_dir_path('%s%05d.npy' % (prefix,w_idx)), w)
+
+    cmd="ffmpeg -y -r {} -i {}/frames/{}%05d.png -vcodec libx264 -pix_fmt yuv420p {}/walk-{}-{}fps.mp4".format(framerate,outdir,prefix,outdir,vidname,framerate)
+    subprocess.call(cmd, shell=True)
+
+def generate_latent_walk(network_pkl, truncation_psi, outdir, walk_type, frames, seeds, npys, save_vector, diameter=2.0, start_seed=0, framerate=24 ):
+    global _G, _D, Gs, noise_vars
+    tflib.init_tf()
+
+    print('Loading networks from "%s"...' % network_pkl)
+    with dnnlib.util.open_url(network_pkl) as fp:
+        _G, _D, Gs = pickle.load(fp)
+
+    os.makedirs(outdir, exist_ok=True)
+    os.makedirs(outdir+"/frames", exist_ok=True)
+
+    # Render images for dlatents initialized from random seeds.
+    Gs_kwargs = {
+        'output_transform': dict(func=tflib.convert_images_to_uint8, nchw_to_nhwc=True),
+        'randomize_noise': False,
+        'truncation_psi': truncation_psi
+    }
+
+    noise_vars = [var for name, var in Gs.components.synthesis.vars.items() if name.startswith('noise')]
+    zs = []
+    ws =[]
+
+
+    # npys specified, let's work with these instead of seeds
+    # npys must be saved as W's (arrays of 18x512)
+    if npys and (len(npys) > 0):
+        ws = npys
+
+
+    wt = walk_type.split('-')
+
+    if wt[0] == 'line':
+        if seeds and (len(seeds) > 0):
+            zs = generate_zs_from_seeds(seeds,Gs)
+
+        if ws == []:
+            number_of_steps = int(frames/(len(zs)-1))+1
+        else:
+            number_of_steps = int(frames/(len(ws)-1))+1
+
+        if (len(wt)>1 and wt[1] == 'w'):
+          if ws == []:
+            for i in range(len(zs)):
+              ws.append(convertZtoW(zs[i]))
+
+          points = line_interpolate(ws,number_of_steps)
+          zpoints = line_interpolate(zs,number_of_steps)
+
+        else:
+          points = line_interpolate(zs,number_of_steps)
+
+
+    # from Gene Kogan
+    elif wt[0] == 'bspline':
+        # bspline in w doesnt work yet
+        # if (len(walk_type)>1 and walk_type[1] == 'w'):
+        #   ws = []
+        #   for i in range(len(zs)):
+        #     ws.append(convertZtoW(zs[i]))
+
+        #   print(ws[0].shape)
+        #   w = []
+        #   for i in range(len(ws)):
+        #     w.append(np.asarray(ws[i]).reshape(512,18))
+        #   points = get_latent_interpolation_bspline(ws,frames,3, 20, shuffle=False)
+        # else:
+          z = []
+          for i in range(len(zs)):
+            z.append(np.asarray(zs[i]).reshape(512))
+          points = get_latent_interpolation_bspline(z,frames,3, 20, shuffle=False)
+
+    # from Dan Shiffman: https://editor.p5js.org/dvs/sketches/Gb0xavYAR
+    elif wt[0] == 'noiseloop':
+        points = get_noiseloop(None,frames,diameter,start_seed)
+
+    if (wt[0] == 'line' and len(wt)>1 and wt[1] == 'w'):
+      # print(points[0][:,:,1])
+      # print(zpoints[0][:,1])
+      # ws = []
+      # for i in enumerate(len(points)):
+      #   ws.append(convertZtoW(points[i]))
+        #added for npys
+        if seeds:
+            seed_out = 'w-' + wt[0] + ('-'.join([str(seed) for seed in seeds]))
+        else:
+            seed_out = 'w-' + wt[0] + '-dlatents'
+
+        generate_images_in_w_space(points, truncation_psi,outdir,save_vector,'frame', seed_out, framerate)
+    elif (len(wt)>1 and wt[1] == 'w'):
+      print('%s is not currently supported in w space, please change your interpolation type' % (wt[0]))
+    else:
+        if(len(wt)>1):
+            seed_out = 'z-' + wt[0] + ('-'.join([str(seed) for seed in seeds]))
+        else:
+            seed_out = 'z-' + walk_type + '-seed' +str(start_seed)
+        generate_latent_images(points, truncation_psi, outdir, save_vector,'frame', seed_out, framerate)
+
+#----------------------------------------------------------------------------
+
+def generate_neighbors(network_pkl, seeds, npys, diameter, truncation_psi, num_samples, save_vector, outdir):
+    global _G, _D, Gs, noise_vars
+    tflib.init_tf()
+    print('Loading networks from "%s"...' % network_pkl)
+    with dnnlib.util.open_url(network_pkl) as fp:
+        _G, _D, Gs = pickle.load(fp)
+
+    os.makedirs(outdir, exist_ok=True)
+
+    # Render images for dlatents initialized from random seeds.
+    Gs_kwargs = {
+        'output_transform': dict(func=tflib.convert_images_to_uint8, nchw_to_nhwc=True),
+        'randomize_noise': False,
+        'truncation_psi': truncation_psi
+    }
+
+    noise_vars = [var for name, var in Gs.components.synthesis.vars.items() if name.startswith('noise')]
+
+    for seed_idx, seed in enumerate(seeds):
+        print('Generating image for seed %d (%d/%d) ...' % (seed, seed_idx+1, len(seeds)))
+        rnd = np.random.RandomState(seed)
+
+        og_z = rnd.randn(1, *Gs.input_shape[1:]) # [minibatch, component]
+        tflib.set_vars({var: rnd.randn(*var.shape.as_list()) for var in noise_vars}) # [height, width]
+        images = Gs.run(og_z, None, **Gs_kwargs) # [minibatch, height, width, channel]
+        # PIL.Image.fromarray(images[0], 'RGB').save(dnnlib.make_run_dir_path('seed%04d.png' % seed))
+        PIL.Image.fromarray(images[0], 'RGB').save(f'{outdir}/seed{seed:05d}.png')
+
+        zs = []
+        z_prefix = 'seed%04d_neighbor' % seed
+
+        for s in range(num_samples):
+            random = np.random.uniform(-diameter,diameter,[1,512])
+#             zs.append(np.clip((og_z+random),-1,1))
+            new_z = np.clip(np.add(og_z,random),-1,1)
+            images = Gs.run(new_z, None, **Gs_kwargs) # [minibatch, height, width, channel]
+            # PIL.Image.fromarray(images[0], 'RGB').save(dnnlib.make_run_dir_path('%s%04d.png' % (z_prefix,s)))
+            PIL.Image.fromarray(images[0], 'RGB').save(f'{outdir}/{z_prefix}{s:05d}.png')
+            # generate_latent_images(zs, truncation_psi, save_vector, z_prefix)
+            if save_vector:
+                np.save(dnnlib.make_run_dir_path('%s%05d.npy' % (z_prefix,s)), new_z)
+
+
+
+#----------------------------------------------------------------------------
+
+def lerp_video(network_pkl,                # Path to pretrained model pkl file
+               seeds,                      # Random seeds
+               grid_w=None,                # Number of columns
+               grid_h=None,                # Number of rows
+               truncation_psi=1.0,         # Truncation trick
+               outdir='out',               # Output dir
+               slowdown=1,                 # Slowdown of the video (power of 2)
+               duration_sec=30.0,          # Duration of video in seconds
+               smoothing_sec=3.0,
+               mp4_fps=30,
+               mp4_codec="libx264",
+               mp4_bitrate="16M"):
+    # Sanity check regarding slowdown
+    message = 'slowdown must be a power of 2 (1, 2, 4, 8, ...) and greater than 0!'
+    assert slowdown & (slowdown - 1) == 0 and slowdown > 0, message
+    # Initialize TensorFlow and create outdir
+    tflib.init_tf()
+    os.makedirs(outdir, exist_ok=True)
+    # Total duration of video and number of frames to generate
+    num_frames = int(np.rint(duration_sec * mp4_fps))
+    total_duration = duration_sec * slowdown
+
+    print(f'Loading network from {network_pkl}...')
+    with dnnlib.util.open_url(network_pkl) as fp:
+        _G, _D, Gs = pickle.load(fp)
+
+    print("Generating latent vectors...")
+    # If there's more than one seed provided and the shape isn't specified
+    if grid_w == grid_h == None and len(seeds) >= 1:
+        # number of images according to the seeds provided
+        num = len(seeds)
+        # Get the grid width and height according to num:
+        grid_w = max(int(np.ceil(np.sqrt(num))), 1)
+        grid_h = max((num - 1) // grid_w + 1, 1)
+        grid_size = [grid_w, grid_h]
+        # [frame, image, channel, component]:
+        shape = [num_frames] + Gs.input_shape[1:]
+        # Get the latents:
+        all_latents = np.stack([np.random.RandomState(seed).randn(*shape).astype(np.float32) for seed in seeds], axis=1)
+    # If only one seed is provided and the shape is specified
+    elif None not in (grid_w, grid_h) and len(seeds) == 1:
+        # Otherwise, the user gives one seed and the grid width and height:
+        grid_size = [grid_w, grid_h]
+        # [frame, image, channel, component]:
+        shape = [num_frames, np.prod(grid_size)] + Gs.input_shape[1:]
+        # Get the latents with the random state:
+        random_state = np.random.RandomState(seeds)
+        all_latents = random_state.randn(*shape).astype(np.float32)
+    else:
+        print("Error: wrong combination of arguments! Please provide \
+                either one seed and the grid width and height, or a \
+                list of seeds to use.")
+        sys.exit(1)
+
+    all_latents = scipy.ndimage.gaussian_filter(
+        all_latents,
+        [smoothing_sec * mp4_fps] + [0] * len(Gs.input_shape),
+        mode="wrap"
+    )
+    all_latents /= np.sqrt(np.mean(np.square(all_latents)))
+    # Name of the final mp4 video
+    mp4 = f"{grid_w}x{grid_h}-lerp-{slowdown}xslowdown.mp4"
+
+    # Aux function to slowdown the video by 2x
+    def double_slowdown(latents, duration_sec, num_frames):
+        # Make an empty latent vector with double the amount of frames
+        z = np.empty(np.multiply(latents.shape, [2, 1, 1]), dtype=np.float32)
+        # Populate it
+        for i in range(len(latents)):
+            z[2*i] = latents[i]
+        # Interpolate in the odd frames
+        for i in range(1, len(z), 2):
+            # For the last frame, we loop to the first one
+            if i == len(z) - 1:
+                z[i] = (z[0] + z[i-1]) / 2
+            else:
+                z[i] = (z[i-1] + z[i+1]) / 2
+        # We also need to double the duration_sec and num_frames
+        duration_sec *= 2
+        num_frames *= 2
+        # Return the new latents, and the two previous quantities
+        return z, duration_sec, num_frames
+
+    while slowdown > 1:
+        all_latents, duration_sec, num_frames = double_slowdown(all_latents, duration_sec, num_frames)
+        slowdown //= 2
+
+    # Define the kwargs for the Generator:
+    Gs_kwargs = dnnlib.EasyDict()
+    Gs_kwargs.output_transform = dict(func=tflib.convert_images_to_uint8,
+                                      nchw_to_nhwc=True)
+    Gs_kwargs.randomize_noise = False
+    if truncation_psi is not None:
+        Gs_kwargs.truncation_psi = truncation_psi
+
+    # Aux function: Frame generation func for moviepy.
+    def make_frame(t):
+        frame_idx = int(np.clip(np.round(t * mp4_fps), 0, num_frames - 1))
+        latents = all_latents[frame_idx]
+        # Get the images (with labels = None)
+        images = Gs.run(latents, None, **Gs_kwargs)
+        # Generate the grid for this timestamp:
+        grid = create_image_grid(images, grid_size)
+        # grayscale => RGB
+        if grid.shape[2] == 1:
+            grid = grid.repeat(3, 2)
+        return grid
+
+    # Generate video using make_frame:
+    print(f'Generating interpolation video of length: {total_duration} seconds...')
+    videoclip = moviepy.editor.VideoClip(make_frame, duration=duration_sec)
+    videoclip.write_videofile(os.path.join(outdir, mp4),
+                              fps=mp4_fps,
+                              codec=mp4_codec,
+                              bitrate=mp4_bitrate)
+
+#----------------------------------------------------------------------------
+
+def _parse_num_range(s):
+    '''Accept either a comma separated list of numbers 'a,b,c' or a range 'a-c' and return as a list of ints.'''
+
+    range_re = re.compile(r'^(\d+)-(\d+)$')
+    m = range_re.match(s)
+    if m:
+        return range(int(m.group(1)), int(m.group(2))+1)
+    vals = s.split(',')
+    return [int(x) for x in vals]
+
+# My extended version of this helper function:
+def _parse_num_range_ext(s):
+    '''
+    Input:
+        s (str): Comma separated string of numbers 'a,b,c', a range 'a-c', or
+                 even a combination of both 'a,b-c', 'a-b,c', 'a,b-c,d,e-f,...'
+    Output:
+        nums (list): Ordered list of ascending ints in s, with repeating values
+                     deleted (can be modified to not do either of this)
+    '''
+    # Sanity check 0:
+    # In case there's a space between the numbers (impossible due to argparse,
+    # but hey, I am that paranoid):
+    s = s.replace(' ', '')
+    # Split w.r.t comma
+    str_list = s.split(',')
+    nums = []
+    for el in str_list:
+        if '-' in el:
+            # The range will be 'a-b', so we wish to find both a and b using re:
+            range_re = re.compile(r'^(\d+)-(\d+)$')
+            match = range_re.match(el)
+            # We get the two numbers:
+            a = int(match.group(1))
+            b = int(match.group(2))
+            # Sanity check 1: accept 'a-b' or 'b-a', with a<=b:
+            if a <= b: r = [n for n in range(a, b + 1)]
+            else: r = [n for n in range(b, a + 1)]
+            # Use extend since r will also be an array:
+            nums.extend(r)
+        else:
+            # It's a single number, so just append it:
+            nums.append(int(el))
+    # Sanity check 2: delete repeating numbers:
+    nums = list(set(nums))
+    # Return the numbers in ascending order:
+    return sorted(nums)
+
+#----------------------------------------------------------------------------
+
+def _parse_npy_files(files):
+    '''Accept a comma separated list of npy files and return a list of z vectors.'''
+
+    zs =[]
+
+    file_list = files.split(",")
+
+
+    for f in file_list:
+        # load numpy array
+        arr = np.load(f)
+        # check if it's actually npz:
+        if 'dlatents' in arr:
+            arr = arr['dlatents']
+        zs.append(arr)
+
+
+
+    return zs
+
+#----------------------------------------------------------------------------
+
+_examples = '''examples:
+
+  # Generate curated MetFaces images without truncation (Fig.10 left)
+  python %(prog)s --outdir=out --trunc=1 --seeds=85,265,297,849 \\
+      --network=https://nvlabs-fi-cdn.nvidia.com/stylegan2-ada/pretrained/metfaces.pkl
+
+  # Generate uncurated MetFaces images with truncation (Fig.12 upper left)
+  python %(prog)s --outdir=out --trunc=0.7 --seeds=600-605 \\
+      --network=https://nvlabs-fi-cdn.nvidia.com/stylegan2-ada/pretrained/metfaces.pkl
+
+  # Generate class conditional CIFAR-10 images (Fig.17 left, Car)
+  python %(prog)s --outdir=out --trunc=1 --seeds=0-35 --class=1 \\
+      --network=https://nvlabs-fi-cdn.nvidia.com/stylegan2-ada/pretrained/cifar10.pkl
+
+  # Render image from projected latent vector
+  python %(prog)s --outdir=out --dlatents=out/dlatents.npz \\
+      --network=https://nvlabs-fi-cdn.nvidia.com/stylegan2-ada/pretrained/ffhq.pkl
+'''
+
+#----------------------------------------------------------------------------
+
+def main():
+    parser = argparse.ArgumentParser(
+        description='Generate images using pretrained network pickle.',
+        epilog=_examples,
+        formatter_class=argparse.RawDescriptionHelpFormatter
+    )
+
+    subparsers = parser.add_subparsers(help='Sub-commands', dest='command')
+
+    parser_generate_images = subparsers.add_parser('generate-images', help='Generate images')
+    parser_generate_images.add_argument('--network', help='Network pickle filename', dest='network_pkl', required=True)
+    parser_generate_images.add_argument('--seeds', type=_parse_num_range, help='List of random seeds', dest='seeds', required=True)
+    parser_generate_images.add_argument('--truncation-psi', type=float, help='Truncation psi (default: %(default)s)', dest='truncation_psi', default=0.5)
+    parser_generate_images.add_argument('--class', dest='class_idx', type=int, help='Class label (default: unconditional)')
+    parser_generate_images.add_argument('--create-grid', action='store_true', help='Add flag to save the generated images in a grid', dest='grid')
+    parser_generate_images.add_argument('--outdir', help='Root directory for run results (default: %(default)s)', default='out', metavar='DIR')
+    parser_generate_images.set_defaults(func=generate_images)
+
+    parser_truncation_traversal = subparsers.add_parser('truncation-traversal', help='Generate truncation walk')
+    parser_truncation_traversal.add_argument('--network', help='Network pickle filename', dest='network_pkl', required=True)
+    parser_truncation_traversal.add_argument('--seed', type=_parse_num_range, help='Singular seed value')
+    parser_truncation_traversal.add_argument('--npys', type=_parse_npy_files, help='List of .npy files')
+    parser_truncation_traversal.add_argument('--fps', type=int, help='Starting value',default=24,dest='framerate')
+    parser_truncation_traversal.add_argument('--start', type=float, help='Starting value')
+    parser_truncation_traversal.add_argument('--stop', type=float, help='Stopping value')
+    parser_truncation_traversal.add_argument('--increment', type=float, help='Incrementing value')
+    parser_truncation_traversal.add_argument('--outdir', help='Root directory for run results (default: %(default)s)', default='out', metavar='DIR')
+    parser_truncation_traversal.set_defaults(func=truncation_traversal)
+
+    parser_generate_latent_walk = subparsers.add_parser('generate-latent-walk', help='Generate latent walk')
+    parser_generate_latent_walk.add_argument('--network', help='Network pickle filename', dest='network_pkl', required=True)
+    parser_generate_latent_walk.add_argument('--trunc', type=float, help='Truncation psi (default: %(default)s)', dest='truncation_psi', default=0.5)
+    parser_generate_latent_walk.add_argument('--walk-type', help='Type of walk (default: %(default)s)', default='line')
+    parser_generate_latent_walk.add_argument('--frames', type=int, help='Frame count (default: %(default)s', default=240)
+    parser_generate_latent_walk.add_argument('--fps', type=int, help='Starting value',default=24,dest='framerate')
+    parser_generate_latent_walk.add_argument('--seeds', type=_parse_num_range, help='List of random seeds')
+    parser_generate_latent_walk.add_argument('--npys', type=_parse_npy_files, help='List of .npy files')
+    parser_generate_latent_walk.add_argument('--save_vector', dest='save_vector', action='store_true', help='also save vector in .npy format')
+    parser_generate_latent_walk.add_argument('--diameter', type=float, help='diameter of noise loop', default=2.0)
+    parser_generate_latent_walk.add_argument('--start_seed', type=int, help='random seed to start noise loop from', default=0)
+    parser_generate_latent_walk.add_argument('--outdir', help='Root directory for run results (default: %(default)s)', default='out', metavar='DIR')
+    parser_generate_latent_walk.set_defaults(func=generate_latent_walk)
+
+    parser_generate_neighbors = subparsers.add_parser('generate-neighbors', help='Generate random neighbors of a seed')
+    parser_generate_neighbors.add_argument('--network', help='Network pickle filename', dest='network_pkl', required=True)
+    parser_generate_neighbors.add_argument('--seeds', type=_parse_num_range, help='List of random seeds')
+    parser_generate_neighbors.add_argument('--npys', type=_parse_npy_files, help='List of .npy files')
+    parser_generate_neighbors.add_argument('--diameter', type=float, help='distance around seed to sample from', default=0.1)
+    parser_generate_neighbors.add_argument('--save_vector', dest='save_vector', action='store_true', help='also save vector in .npy format')
+    parser_generate_neighbors.add_argument('--num_samples', type=int, help='How many neighbors to generate (default: %(default)s', default=25)
+    parser_generate_neighbors.add_argument('--trunc', type=float, help='Truncation psi (default: %(default)s)', dest='truncation_psi', default=0.5)
+    parser_generate_neighbors.add_argument('--outdir', help='Root directory for run results (default: %(default)s)', default='out', metavar='DIR')
+    parser_generate_neighbors.set_defaults(func=generate_neighbors)
+
+    parser_lerp_video = subparsers.add_parser('lerp-video', help='Generate interpolation video (lerp) between random vectors')
+    parser_lerp_video.add_argument('--network', help='Path to network pickle filename', dest='network_pkl', required=True)
+    parser_lerp_video.add_argument('--seeds', type=_parse_num_range_ext, help='List of random seeds', dest='seeds', required=True)
+    parser_lerp_video.add_argument('--grid-w', type=int, help='Video grid width/columns (default: %(default)s)', default=None, dest='grid_w')
+    parser_lerp_video.add_argument('--grid-h', type=int, help='Video grid height/rows (default: %(default)s)', default=None, dest='grid_h')
+    parser_lerp_video.add_argument('--trunc', type=float, help='Truncation psi (default: %(default)s)', default=1.0, dest='truncation_psi')
+    parser_lerp_video.add_argument('--slowdown', type=int, help='Slowdown the video by this amount; must be a power of 2 (default: %(default)s)', default=1, dest='slowdown')
+    parser_lerp_video.add_argument('--duration-sec', type=float, help='Duration of video (default: %(default)s)', default=30.0, dest='duration_sec')
+    parser_lerp_video.add_argument('--fps', type=int, help='FPS of generated video (default: %(default)s)', default=30, dest='mp4_fps')
+    parser_lerp_video.add_argument('--outdir', help='Root directory for run results (default: %(default)s)', default='out', metavar='DIR')
+    parser_lerp_video.set_defaults(func=lerp_video)
+
+    args = parser.parse_args()
+    kwargs = vars(args)
+    subcmd = kwargs.pop('command')
+
+    if subcmd is None:
+        print('Error: missing subcommand.  Re-run with --help for usage.')
+        sys.exit(1)
+
+    func = kwargs.pop('func')
+    func(**kwargs)
+
+#----------------------------------------------------------------------------
+
+if __name__ == "__main__":
+    main()
+
+#----------------------------------------------------------------------------

models.py

@@ -0,0 +1,142 @@
+import argparse
+import copy
+
+import warnings
+import tensorflow as tf
+tf.compat.v1.logging.set_verbosity(tf.compat.v1.logging.ERROR)
+import warnings
+warnings.filterwarnings('ignore', category=FutureWarning)
+warnings.filterwarnings('ignore', category=DeprecationWarning)
+import sys, getopt, os
+
+import numpy as np
+import dnnlib
+from dnnlib import EasyDict
+import dnnlib.tflib as tflib
+from dnnlib.tflib import tfutil
+from dnnlib.tflib.autosummary import autosummary
+
+from training import misc
+import pickle
+import argparse
+
+def create_model(config_id = 'config-f', gamma = None, height = 512, width = 512, cond = None, label_size = 0):
+    train     = EasyDict(run_func_name='training.diagnostic.create_initial_pkl') # Options for training loop.
+    G         = EasyDict(func_name='training.networks_stylegan2.G_main')       # Options for generator network.
+    D         = EasyDict(func_name='training.networks_stylegan2.D_stylegan2')  # Options for discriminator network.
+    D_loss    = EasyDict(func_name='training.loss.D_logistic_r1')              # Options for discriminator loss.
+    sched     = EasyDict()                                                     # Options for TrainingSchedule.
+    sc        = dnnlib.SubmitConfig()                                          # Options for dnnlib.submit_run().
+    tf_config = {'rnd.np_random_seed': 1000}                                   # Options for tflib.init_tf().
+
+    sched.minibatch_size_base = 192
+    sched.minibatch_gpu_base = 3
+    D_loss.gamma = 10
+    desc = 'stylegan2'
+
+    dataset_args = EasyDict() # (tfrecord_dir=dataset)
+
+    if cond:
+        desc += '-cond'; dataset_args.max_label_size = 'full' # conditioned on full label
+
+    desc += '-' + config_id
+
+    # Configs A-E: Shrink networks to match original StyleGAN.
+    if config_id != 'config-f':
+        G.fmap_base = D.fmap_base = 8 << 10
+
+    # Config E: Set gamma to 100 and override G & D architecture.
+    if config_id.startswith('config-e'):
+        D_loss.gamma = 100
+        if 'Gorig'   in config_id: G.architecture = 'orig'
+        if 'Gskip'   in config_id: G.architecture = 'skip' # (default)
+        if 'Gresnet' in config_id: G.architecture = 'resnet'
+        if 'Dorig'   in config_id: D.architecture = 'orig'
+        if 'Dskip'   in config_id: D.architecture = 'skip'
+        if 'Dresnet' in config_id: D.architecture = 'resnet' # (default)
+
+    # Configs A-D: Enable progressive growing and switch to networks that support it.
+    if config_id in ['config-a', 'config-b', 'config-c', 'config-d']:
+        sched.lod_initial_resolution = 8
+        sched.G_lrate_base = sched.D_lrate_base = 0.001
+        sched.G_lrate_dict = sched.D_lrate_dict = {128: 0.0015, 256: 0.002, 512: 0.003, 1024: 0.003}
+        sched.minibatch_size_base = 32 # (default)
+        sched.minibatch_size_dict = {8: 256, 16: 128, 32: 64, 64: 32}
+        sched.minibatch_gpu_base = 4 # (default)
+        sched.minibatch_gpu_dict = {8: 32, 16: 16, 32: 8, 64: 4}
+        G.synthesis_func = 'G_synthesis_stylegan_revised'
+        D.func_name = 'training.networks_stylegan2.D_stylegan'
+
+    # Configs A-C: Disable path length regularization.
+    if config_id in ['config-a', 'config-b', 'config-c']:
+        G_loss = EasyDict(func_name='training.loss.G_logistic_ns')
+
+    # Configs A-B: Disable lazy regularization.
+    if config_id in ['config-a', 'config-b']:
+        train.lazy_regularization = False
+
+    # Config A: Switch to original StyleGAN networks.
+    if config_id == 'config-a':
+        G = EasyDict(func_name='training.networks_stylegan.G_style')
+        D = EasyDict(func_name='training.networks_stylegan.D_basic')
+
+    if gamma is not None:
+        D_loss.gamma = gamma
+
+    G.update(resolution_h=height)
+    G.update(resolution_w=width)
+    D.update(resolution_h=height)
+    D.update(resolution_w=width)
+
+    sc.submit_target = dnnlib.SubmitTarget.DIAGNOSTIC
+    sc.local.do_not_copy_source_files = True
+    kwargs = EasyDict(train)
+    # [EDITED]
+    kwargs.update(G_args=G, D_args=D, tf_config=tf_config, config_id=config_id,
+        resolution_h=height, resolution_w=width, label_size = label_size)
+    kwargs.submit_config = copy.deepcopy(sc)
+    kwargs.submit_config.run_desc = desc
+    dnnlib.submit_diagnostic(**kwargs)
+    return f'network-initial-config-f-{height}x{width}-{label_size}.pkl'
+
+def _str_to_bool(v):
+    if isinstance(v, bool):
+        return v
+    if v.lower() in ('yes', 'true', 't', 'y', '1'):
+        return True
+    elif v.lower() in ('no', 'false', 'f', 'n', '0'):
+        return False
+    else:
+        raise argparse.ArgumentTypeError('Boolean value expected.')
+
+def _parse_comma_sep(s):
+    if s is None or s.lower() == 'none' or s == '':
+        return []
+    return s.split(',')
+
+def copy_weights(source_pkl, target_pkl, output_pkl):
+
+    tflib.init_tf()
+
+    with tf.Session():
+        with tf.device('/gpu:0'):
+
+            sourceG, sourceD, sourceGs = pickle.load(open(source_pkl, 'rb'))
+            targetG, targetD, targetGs = pickle.load(open(target_pkl, 'rb'))
+            
+            # print('Source:')
+            # sourceG.print_layers()
+            # sourceD.print_layers() 
+            # sourceGs.print_layers()
+            
+            # print('Target:')
+            # targetG.print_layers()
+            # targetD.print_layers() 
+            # targetGs.print_layers()
+            
+            targetG.copy_compatible_trainables_from(sourceG)
+            targetD.copy_compatible_trainables_from(sourceD)
+            targetGs.copy_compatible_trainables_from(sourceGs)
+
+            with open(os.path.join('./', output_pkl), 'wb') as file:
+                pickle.dump((targetG, targetD, targetGs), file, protocol=pickle.HIGHEST_PROTOCOL)

rasm.py

@@ -0,0 +1,146 @@
+
+from utils import download_url
+import argparse
+import numpy as np
+import PIL.Image
+import dnnlib
+import dnnlib.tflib as tflib
+import re
+import sys
+from io import BytesIO
+import IPython.display
+from math import ceil
+from PIL import Image, ImageDraw
+import os
+import pickle
+from utils import log_progress, imshow, create_image_grid, show_animation
+import imageio
+import glob
+import gdown 
+
+class Rasm:
+
+    def __init__(self, mode = 'calligraphy'):
+
+        if mode == 'calligraphy':
+            url = 'https://drive.google.com/uc?id=138fdURGxdkOwZq7IWvnrGLcfo5VI8O1R'
+
+        else:
+            url = 'https://drive.google.com/uc?id=13h-alXGI0hbNOJy1qbmeoroXZSPBHEG2'
+
+        output = 'model.pkl'
+        print('Downloading networks from "%s"...' %url)
+        gdown.download(url, output, quiet=False)
+        dnnlib.tflib.init_tf()
+        with dnnlib.util.open_url(output) as fp:
+            self._G, self._D, self.Gs = pickle.load(fp)
+        self.noise_vars = [var for name, var in self.Gs.components.synthesis.vars.items() if name.startswith('noise')]
+    
+    # Generates a list of images, based on a list of latent vectors (Z), and a list (or a single constant) of truncation_psi's.
+    def generate_images_in_w_space(self, dlatents, truncation_psi):
+        Gs_kwargs = dnnlib.EasyDict()
+        Gs_kwargs.output_transform = dict(func=tflib.convert_images_to_uint8, nchw_to_nhwc=True)
+        Gs_kwargs.randomize_noise = False
+        Gs_kwargs.truncation_psi = truncation_psi
+        # dlatent_avg = self.Gs.get_var('dlatent_avg') # [component]
+
+        imgs = []
+        for _, dlatent in log_progress(enumerate(dlatents), name = "Generating images"):
+            #row_dlatents = (dlatent[np.newaxis] - dlatent_avg) * np.reshape(truncation_psi, [-1, 1, 1]) + dlatent_avg
+            # dl = (dlatent-dlatent_avg)*truncation_psi   + dlatent_avg
+            row_images = self.Gs.components.synthesis.run(dlatent,  **Gs_kwargs)
+            imgs.append(PIL.Image.fromarray(row_images[0], 'RGB'))
+        return imgs       
+
+    def generate_images(self, zs, truncation_psi, class_idx = None):
+        Gs_kwargs = dnnlib.EasyDict()
+        Gs_kwargs.output_transform = dict(func=tflib.convert_images_to_uint8, nchw_to_nhwc=True)
+        Gs_kwargs.randomize_noise = False
+        if not isinstance(truncation_psi, list):
+            truncation_psi = [truncation_psi] * len(zs)
+            
+        imgs = []
+        label = np.zeros([1] + self.Gs.input_shapes[1][1:])
+        if class_idx is not None:
+            label[:, class_idx] = 1
+        else:
+            label = None
+        for z_idx, z in log_progress(enumerate(zs), size = len(zs), name = "Generating images"):
+            Gs_kwargs.truncation_psi = truncation_psi[z_idx]
+            noise_rnd = np.random.RandomState(1) # fix noise
+            tflib.set_vars({var: noise_rnd.randn(*var.shape.as_list()) for var in self.noise_vars}) # [height, width]
+            images = self.Gs.run(z, label, **Gs_kwargs) # [minibatch, height, width, channel]
+            imgs.append(PIL.Image.fromarray(images[0], 'RGB'))
+        return imgs
+    
+    def generate_from_zs(self, zs, truncation_psi = 0.5):
+        Gs_kwargs = dnnlib.EasyDict()
+        Gs_kwargs.output_transform = dict(func=tflib.convert_images_to_uint8, nchw_to_nhwc=True)
+        Gs_kwargs.randomize_noise = False
+        if not isinstance(truncation_psi, list):
+            truncation_psi = [truncation_psi] * len(zs)
+            
+        for z_idx, z in log_progress(enumerate(zs), size = len(zs), name = "Generating images"):
+            Gs_kwargs.truncation_psi = truncation_psi[z_idx]
+            noise_rnd = np.random.RandomState(1) # fix noise
+            tflib.set_vars({var: noise_rnd.randn(*var.shape.as_list()) for var in self.noise_vars}) # [height, width]
+            images = self.Gs.run(z, None, **Gs_kwargs) # [minibatch, height, width, channel]
+            img = PIL.Image.fromarray(images[0], 'RGB')
+            imshow(img)
+
+    def generate_random_zs(self, size):
+        seeds = np.random.randint(2**32, size=size)
+        zs = []
+        for _, seed in enumerate(seeds):
+            rnd = np.random.RandomState(seed)
+            z = rnd.randn(1, *self.Gs.input_shape[1:]) # [minibatch, component]
+            zs.append(z)
+        return zs
+
+
+    def generate_zs_from_seeds(self, seeds):
+        zs = []
+        for _, seed in enumerate(seeds):
+            rnd = np.random.RandomState(seed)
+            z = rnd.randn(1, *self.Gs.input_shape[1:]) # [minibatch, component]
+            zs.append(z)
+        return zs
+
+    # Generates a list of images, based on a list of seed for latent vectors (Z), and a list (or a single constant) of truncation_psi's.
+    def generate_images_from_seeds(self, seeds, truncation_psi):
+        return imshow(self.generate_images(self.generate_zs_from_seeds(seeds), truncation_psi)[0])
+
+    def generate_randomly(self, truncation_psi = 0.5):
+        return self.generate_images_from_seeds(np.random.randint(4294967295, size=1), truncation_psi=truncation_psi)
+
+    def generate_grid(self, truncation_psi = 0.7): 
+      seeds = np.random.randint((2**32 - 1), size=9)
+      return create_image_grid(self.generate_images(self.generate_zs_from_seeds(seeds), truncation_psi), 0.7 , 3)
+    
+    def generate_animation(self, size = 9, steps = 10, trunc_psi = 0.5):
+      seeds = list(np.random.randint((2**32) - 1, size=size))
+      seeds = seeds + [seeds[0]]
+      zs = self.generate_zs_from_seeds(seeds)
+
+      imgs = self.generate_images(self.interpolate(zs, steps = steps), trunc_psi)
+      movie_name = 'animation.mp4'
+      with imageio.get_writer(movie_name, mode='I') as writer:
+        for image in log_progress(list(imgs), name = "Creating animation"):
+            writer.append_data(np.array(image))
+      return show_animation(movie_name)
+
+    def convertZtoW(self, latent, truncation_psi=0.7, truncation_cutoff=9):
+        dlatent = self.Gs.components.mapping.run(latent, None) # [seed, layer, component]
+        dlatent_avg = self.Gs.get_var('dlatent_avg') # [component]
+        for i in range(truncation_cutoff):
+            dlatent[0][i] = (dlatent[0][i]-dlatent_avg)*truncation_psi + dlatent_avg
+            
+        return dlatent
+
+    def interpolate(self, zs, steps = 10):
+        out = []
+        for i in range(len(zs)-1):
+            for index in range(steps):
+                fraction = index/float(steps) 
+                out.append(zs[i+1]*fraction + zs[i]*(1-fraction))
+        return out

requirements.txt

@@ -0,0 +1,32 @@
+absl-py==0.7.0
+astor==0.7.1
+certifi==2018.11.29
+chardet==3.0.4
+Click==7.0
+Flask==1.0.2
+Flask-Cors==3.0.7
+gast==0.2.2
+gevent==1.4.0
+greenlet==0.4.15
+grpcio==1.19.0
+h5py==2.9.0
+idna==2.8
+itsdangerous==1.1.0
+Jinja2==2.10
+Keras-Applications==1.0.7
+Keras-Preprocessing==1.0.9
+Markdown==3.0.1
+MarkupSafe==1.1.1
+mock==2.0.0
+numpy==1.16.2
+pbr==5.1.2
+Pillow==5.4.1
+protobuf==3.6.1
+requests==2.21.0
+six==1.12.0
+tensorflow-gpu==1.15.0
+termcolor==1.1.0
+urllib3==1.24.1
+Werkzeug==0.14.1
+wget==3.2
+runway-python

utils.py

@@ -0,0 +1,165 @@
+import glob
+import os 
+from PIL import Image
+import urllib.request
+from tqdm import tqdm
+
+import numpy as np
+import PIL.Image
+import sys
+from io import BytesIO
+import IPython.display
+import numpy as np
+from math import ceil
+from PIL import Image, ImageDraw
+import os
+
+from IPython.display import HTML
+from base64 import b64encode
+import imageio
+
+def show_animation(movie_name):
+  mp4 = open(movie_name,'rb').read()
+  data_url = "data:video/mp4;base64," + b64encode(mp4).decode()
+  return HTML("""
+  <video width=400 controls>
+        <source src="%s" type="video/mp4">
+  </video>
+  """ % data_url)
+
+def imshow(a, format='png', jpeg_fallback=True):
+        a = np.asarray(a, dtype=np.uint8)
+        str_file = BytesIO()
+        PIL.Image.fromarray(a).save(str_file, format)
+        im_data = str_file.getvalue()
+        try:
+            disp = IPython.display.display(IPython.display.Image(im_data))
+        except IOError:
+            if jpeg_fallback and format != 'jpeg':
+                print ('Warning: image was too large to display in format "{}"; '
+                        'trying jpeg instead.').format(format)
+                return imshow(a, format='jpeg')
+            else:
+                raise
+        return disp
+
+        
+def clamp(x, minimum, maximum):
+    return max(minimum, min(x, maximum))
+
+def create_image_grid(images, scale=0.25, rows=1):
+    w,h = images[0].size
+    w = int(w*scale)
+    h = int(h*scale)
+    height = rows*h
+    cols = ceil(len(images) / rows)
+    width = cols*w
+    canvas = PIL.Image.new('RGBA', (width,height), 'white')
+    for i,img in enumerate(images):
+        img = img.resize((w,h), PIL.Image.ANTIALIAS)
+        canvas.paste(img, (w*(i % cols), h*(i // cols))) 
+    return canvas
+
+def find_latest_pkl(path):
+  curr_best = 0
+  latest_pkl = ''
+  for pkl in glob.glob(f'{path}/*.pkl'):
+    ckpt_number = int(pkl.split('-')[-1][:-4])
+    if curr_best < ckpt_number:
+      curr_best = ckpt_number
+      latest_pkl = pkl
+  return latest_pkl
+
+def resize(path, dim = (512, 512)):
+  dirs = os.listdir(path)
+  out_path = f'{path}/{dim[0]}x{dim[1]}'
+  os.makedirs(out_path, exist_ok=True)
+  for item in log_progress(dirs):
+    img_path = f'{path}/{item}'
+    if os.path.isfile(img_path):
+        im = Image.open(img_path)
+        imResize = im.resize(dim, Image.ANTIALIAS).convert('RGB')
+        imResize.save(f'{out_path}/{item}', 'JPEG', quality=90)
+  return out_path
+
+def resize_dirs(path, out_dir, dim = (512, 512)):
+  sub_dirs = os.listdir(path)
+  for sub_dir in sub_dirs:
+    out_path = f'{out_dir}/{sub_dir}'
+    os.makedirs(out_path, exist_ok=True)
+    for item in log_progress(os.listdir(f'{path}/{sub_dir}/')[:10]):
+        img_path = f'{path}/{sub_dir}/{item}'
+        if os.path.isfile(img_path):
+            im = Image.open(img_path)
+            imResize = im.resize(dim, Image.ANTIALIAS).convert('RGB')
+            imResize.save(f'{out_path}/{item}', 'JPEG', quality=90)
+  return out_dir
+
+class DownloadProgressBar(tqdm):
+    def update_to(self, b=1, bsize=1, tsize=None):
+        if tsize is not None:
+            self.total = tsize
+        self.update(b * bsize - self.n)
+
+# https://stackoverflow.com/a/53877507
+def download_url(url, output_path):
+    with DownloadProgressBar(unit='B', unit_scale=True,
+                             miniters=1, desc=url.split('/')[-1]) as t:
+        urllib.request.urlretrieve(url, filename=output_path, reporthook=t.update_to)
+
+# Taken from https://github.com/alexanderkuk/log-progress
+def log_progress(sequence, every=1, size=None, name='Items'):
+    from ipywidgets import IntProgress, HTML, VBox
+    from IPython.display import display
+
+    is_iterator = False
+    if size is None:
+        try:
+            size = len(sequence)
+        except TypeError:
+            is_iterator = True
+    if size is not None:
+        if every is None:
+            if size <= 200:
+                every = 1
+            else:
+                every = int(size / 200)     # every 0.5%
+    else:
+        assert every is not None, 'sequence is iterator, set every'
+
+    if is_iterator:
+        progress = IntProgress(min=0, max=1, value=1)
+        progress.bar_style = 'info'
+    else:
+        progress = IntProgress(min=0, max=size, value=0)
+    label = HTML()
+    box = VBox(children=[label, progress])
+    display(box)
+
+    index = 0
+    try:
+        for index, record in enumerate(sequence, 1):
+            if index == 1 or index % every == 0:
+                if is_iterator:
+                    label.value = '{name}: {index} / ?'.format(
+                        name=name,
+                        index=index
+                    )
+                else:
+                    progress.value = index
+                    label.value = u'{name}: {index} / {size}'.format(
+                        name=name,
+                        index=index,
+                        size=size
+                    )
+            yield record
+    except:
+        progress.bar_style = 'danger'
+        raise
+    else:
+        progress.bar_style = 'success'
+        progress.value = index
+        label.value = "{name}: {index}".format(
+            name=name,
+            index=str(index or '?')
+        )