WGAN_GP.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": [],
      "authorship_tag": "ABX9TyOKuuVvuxa/PKZ2VjUvlHby",
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    },
    "widgets": {
      "application/vnd.jupyter.widget-state+json": {
        "f9bdc6b94ab14f5d8383c7fb54409003": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HBoxModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HBoxModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HBoxView",
            "box_style": "",
            "children": [
              "IPY_MODEL_b214dd356fcd49b1ab683b1750d4e9a0",
              "IPY_MODEL_ac34f33cb35942368f26656196f757c4",
              "IPY_MODEL_cd9b035469ff4c989136d30fa669bf5c"
            ],
            "layout": "IPY_MODEL_186dd44a4df14e6782962ae79d833426"
          }
        },
        "b214dd356fcd49b1ab683b1750d4e9a0": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HTMLModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HTMLModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HTMLView",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_886e4617d36947e4954a35a4076fc225",
            "placeholder": "​",
            "style": "IPY_MODEL_bb768bc8f6af48dfa233dad8a7f9af3e",
            "value": "100%"
          }
        },
        "ac34f33cb35942368f26656196f757c4": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "FloatProgressModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "FloatProgressModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "ProgressView",
            "bar_style": "success",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_0f193b67cf194d69bd1e4b50a743f05b",
            "max": 9912422,
            "min": 0,
            "orientation": "horizontal",
            "style": "IPY_MODEL_7effed533e8c4beba8860b6d49bdbc8f",
            "value": 9912422
          }
        },
        "cd9b035469ff4c989136d30fa669bf5c": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HTMLModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HTMLModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HTMLView",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_49d8c1cf3379442b97171e1dcf63e0ac",
            "placeholder": "​",
            "style": "IPY_MODEL_6d85927cebd144a0a7327168a488d3e6",
            "value": " 9912422/9912422 [00:00&lt;00:00, 21752697.42it/s]"
          }
        },
        "186dd44a4df14e6782962ae79d833426": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "886e4617d36947e4954a35a4076fc225": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "bb768bc8f6af48dfa233dad8a7f9af3e": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "DescriptionStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "DescriptionStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "description_width": ""
          }
        },
        "0f193b67cf194d69bd1e4b50a743f05b": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "7effed533e8c4beba8860b6d49bdbc8f": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "ProgressStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "ProgressStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "bar_color": null,
            "description_width": ""
          }
        },
        "49d8c1cf3379442b97171e1dcf63e0ac": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "6d85927cebd144a0a7327168a488d3e6": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "DescriptionStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "DescriptionStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "description_width": ""
          }
        },
        "0b8992d775b14889a617ddff6efe1362": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HBoxModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HBoxModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HBoxView",
            "box_style": "",
            "children": [
              "IPY_MODEL_be62ad0778ba4be5bf17c7bc5f181c9a",
              "IPY_MODEL_3af65e68e4a54e46af18207756b30095",
              "IPY_MODEL_64a70b0bd2e04dde8b1365c1d105ef2a"
            ],
            "layout": "IPY_MODEL_2dd8d57f2d574c4cb5a89adf0c25037b"
          }
        },
        "be62ad0778ba4be5bf17c7bc5f181c9a": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HTMLModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HTMLModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HTMLView",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_fc28876ba0da487fb78f7d1d279d6934",
            "placeholder": "​",
            "style": "IPY_MODEL_ff87c106bb2c45f5a28a8af79a259730",
            "value": "100%"
          }
        },
        "3af65e68e4a54e46af18207756b30095": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "FloatProgressModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "FloatProgressModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "ProgressView",
            "bar_style": "success",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_0bc9c686e46c4e73a8b6622ae59846b3",
            "max": 28881,
            "min": 0,
            "orientation": "horizontal",
            "style": "IPY_MODEL_1fce092a7796478fa4b86693feabb339",
            "value": 28881
          }
        },
        "64a70b0bd2e04dde8b1365c1d105ef2a": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HTMLModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HTMLModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HTMLView",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_744b6f6e58dc473bb1b16740217d8709",
            "placeholder": "​",
            "style": "IPY_MODEL_d9315352a5e143499fa0cd31276ef4f9",
            "value": " 28881/28881 [00:00&lt;00:00, 523891.20it/s]"
          }
        },
        "2dd8d57f2d574c4cb5a89adf0c25037b": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "fc28876ba0da487fb78f7d1d279d6934": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "ff87c106bb2c45f5a28a8af79a259730": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "DescriptionStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "DescriptionStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "description_width": ""
          }
        },
        "0bc9c686e46c4e73a8b6622ae59846b3": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "1fce092a7796478fa4b86693feabb339": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "ProgressStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "ProgressStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "bar_color": null,
            "description_width": ""
          }
        },
        "744b6f6e58dc473bb1b16740217d8709": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "d9315352a5e143499fa0cd31276ef4f9": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "DescriptionStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "DescriptionStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "description_width": ""
          }
        },
        "b31ce68c359f42a7b12622acbaf031da": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HBoxModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HBoxModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HBoxView",
            "box_style": "",
            "children": [
              "IPY_MODEL_1c21d222feb044c5917f0eca4f85cb8e",
              "IPY_MODEL_dff6c6c5ad204abeabb2106b6e83fb20",
              "IPY_MODEL_cf804a98df054874b915eb21da9ca536"
            ],
            "layout": "IPY_MODEL_5815ff5cf1bd4c3287ea857088054627"
          }
        },
        "1c21d222feb044c5917f0eca4f85cb8e": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HTMLModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HTMLModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HTMLView",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_ed55a18fe7d7414da586b21c62ce74a6",
            "placeholder": "​",
            "style": "IPY_MODEL_e79c4a4353ac431fbf6f69bf88a65cb4",
            "value": "100%"
          }
        },
        "dff6c6c5ad204abeabb2106b6e83fb20": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "FloatProgressModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "FloatProgressModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "ProgressView",
            "bar_style": "success",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_1a9cb7b02268475caefb554d7c80de1d",
            "max": 1648877,
            "min": 0,
            "orientation": "horizontal",
            "style": "IPY_MODEL_33ce54d2d70e49eabb57cf9e559c17b0",
            "value": 1648877
          }
        },
        "cf804a98df054874b915eb21da9ca536": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HTMLModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HTMLModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HTMLView",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_7faaecddc5ce4f28b21392271f01dbaa",
            "placeholder": "​",
            "style": "IPY_MODEL_756f6b248d7e47dcae8fb8bd8d3c9f80",
            "value": " 1648877/1648877 [00:00&lt;00:00, 14398949.37it/s]"
          }
        },
        "5815ff5cf1bd4c3287ea857088054627": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "ed55a18fe7d7414da586b21c62ce74a6": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "e79c4a4353ac431fbf6f69bf88a65cb4": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "DescriptionStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "DescriptionStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "description_width": ""
          }
        },
        "1a9cb7b02268475caefb554d7c80de1d": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "33ce54d2d70e49eabb57cf9e559c17b0": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "ProgressStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "ProgressStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "bar_color": null,
            "description_width": ""
          }
        },
        "7faaecddc5ce4f28b21392271f01dbaa": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "756f6b248d7e47dcae8fb8bd8d3c9f80": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "DescriptionStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "DescriptionStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "description_width": ""
          }
        },
        "9bf1ec46cf694cf8bd3f453a441ac130": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HBoxModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HBoxModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HBoxView",
            "box_style": "",
            "children": [
              "IPY_MODEL_966ddd3b8dc2450092eac2bbc213d000",
              "IPY_MODEL_74c8baaf41794b73a3d4635279d940a9",
              "IPY_MODEL_bfee2d5dafad4a6eb7ac1424b02d941e"
            ],
            "layout": "IPY_MODEL_62e23952e4ce403c8f867aaedbcaa43b"
          }
        },
        "966ddd3b8dc2450092eac2bbc213d000": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HTMLModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HTMLModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HTMLView",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_f867d1a6388b49ee9662126dcea68061",
            "placeholder": "​",
            "style": "IPY_MODEL_79d9268b3a654cc5a7fe8a4125384658",
            "value": "100%"
          }
        },
        "74c8baaf41794b73a3d4635279d940a9": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "FloatProgressModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "FloatProgressModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "ProgressView",
            "bar_style": "success",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_21b6492411074f6bb0088655bb062217",
            "max": 4542,
            "min": 0,
            "orientation": "horizontal",
            "style": "IPY_MODEL_2789c61c4694478c907fdf83703f3d81",
            "value": 4542
          }
        },
        "bfee2d5dafad4a6eb7ac1424b02d941e": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HTMLModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HTMLModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HTMLView",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_a4054af378c34659b8ba0bf0a81f2e6b",
            "placeholder": "​",
            "style": "IPY_MODEL_0360b616956444df8977675e5391cbb9",
            "value": " 4542/4542 [00:00&lt;00:00, 86500.52it/s]"
          }
        },
        "62e23952e4ce403c8f867aaedbcaa43b": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "f867d1a6388b49ee9662126dcea68061": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "79d9268b3a654cc5a7fe8a4125384658": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "DescriptionStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "DescriptionStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "description_width": ""
          }
        },
        "21b6492411074f6bb0088655bb062217": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "2789c61c4694478c907fdf83703f3d81": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "ProgressStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "ProgressStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "bar_color": null,
            "description_width": ""
          }
        },
        "a4054af378c34659b8ba0bf0a81f2e6b": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "0360b616956444df8977675e5391cbb9": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "DescriptionStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "DescriptionStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "description_width": ""
          }
        }
      }
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/alaeddinehamroun/Working-with-GANs/blob/main/WGAN_GP.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "The W stands for Wasserstein.\n",
        "\n",
        "GP stands for Gradient Penalties: it's job is to prevent mode collapse."
      ],
      "metadata": {
        "id": "LaXKyKzyY0_A"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Generator and Critic"
      ],
      "metadata": {
        "id": "DsMxXkyzZSNY"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "In WGAN-GP, you no longer use a discriminator that classifies fake and real as 0 an 1 but rather a critic that scores images with real numbers."
      ],
      "metadata": {
        "id": "OdrEElGnZZU1"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "#### Packages and Visualizations"
      ],
      "metadata": {
        "id": "oG-6Ua3CZtZD"
      }
    },
    {
      "cell_type": "code",
      "execution_count": 1,
      "metadata": {
        "id": "vH4TerezYSie"
      },
      "outputs": [],
      "source": [
        "import torch\n",
        "from torch import nn\n",
        "from tqdm.auto import tqdm\n",
        "from torchvision import transforms\n",
        "from torchvision.datasets import MNIST\n",
        "from torchvision.utils import make_grid\n",
        "from torch.utils.data import DataLoader\n",
        "import matplotlib.pyplot as plt\n",
        "\n",
        "def show_tensor_images(image_tensor, num_images=25, size=(1, 28, 28)):\n",
        "  '''\n",
        "  Function for visualizing images: Given a tensor of images, number of images, and\n",
        "  size per image, plots and prints the images in an uniform gird.\n",
        "  '''\n",
        "  image_tensor = (image_tensor + 1) / 2\n",
        "  image_unflat = image_tensor.detach().cpu()\n",
        "  image_grid = make_grid(image_unflat[:num_images], nrow=5)\n",
        "  plt.imshow(image_grid.permute(1, 2, 0).squeeze())\n",
        "  plt.show()\n",
        "\n",
        "def make_grad_hook():\n",
        "  '''\n",
        "  Function to keep track of gradients for visualization purposes,\n",
        "  which fills the grads list when using model.apply(grad_hook).\n",
        "  '''\n",
        "  grads = []\n",
        "  def grad_hook(m):\n",
        "    if isinstance(m, nn.Conv2d) or isinstance(m, nn.ConvTranspose2d):\n",
        "      grads.append(m.weight.grad)\n",
        "    return grads, grad_hook"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "#### Generator and Noise"
      ],
      "metadata": {
        "id": "7tLOSMzdbDmB"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "class Generator(nn.Module):\n",
        "    '''\n",
        "    Generator Class\n",
        "    Values:\n",
        "        z_dim: the dimension of the noise vector, a scalar\n",
        "        im_chan: the number of channels in the images, fitted for the dataset used, a scalar\n",
        "              (MNIST is black-and-white, so 1 channel is your default)\n",
        "        hidden_dim: the inner dimension, a scalar\n",
        "    '''\n",
        "    def __init__(self, z_dim=10, im_chan=1, hidden_dim=64):\n",
        "        super(Generator, self).__init__()\n",
        "        self.z_dim = z_dim\n",
        "        # Build the neural network\n",
        "        self.gen = nn.Sequential(\n",
        "            self.make_gen_block(z_dim, hidden_dim * 4),\n",
        "            self.make_gen_block(hidden_dim * 4, hidden_dim * 2, kernel_size=4, stride=1),\n",
        "            self.make_gen_block(hidden_dim * 2, hidden_dim),\n",
        "            self.make_gen_block(hidden_dim, im_chan, kernel_size=4, final_layer=True),\n",
        "        )\n",
        "\n",
        "    def make_gen_block(self, input_channels, output_channels, kernel_size=3, stride=2, final_layer=False):\n",
        "        '''\n",
        "        Function to return a sequence of operations corresponding to a generator block of DCGAN;\n",
        "        a transposed convolution, a batchnorm (except in the final layer), and an activation.\n",
        "        Parameters:\n",
        "            input_channels: how many channels the input feature representation has\n",
        "            output_channels: how many channels the output feature representation should have\n",
        "            kernel_size: the size of each convolutional filter, equivalent to (kernel_size, kernel_size)\n",
        "            stride: the stride of the convolution\n",
        "            final_layer: a boolean, true if it is the final layer and false otherwise \n",
        "                      (affects activation and batchnorm)\n",
        "        '''\n",
        "        if not final_layer:\n",
        "            return nn.Sequential(\n",
        "                nn.ConvTranspose2d(input_channels, output_channels, kernel_size, stride),\n",
        "                nn.BatchNorm2d(output_channels),\n",
        "                nn.ReLU(inplace=True),\n",
        "            )\n",
        "        else:\n",
        "            return nn.Sequential(\n",
        "                nn.ConvTranspose2d(input_channels, output_channels, kernel_size, stride),\n",
        "                nn.Tanh(),\n",
        "            )\n",
        "\n",
        "    def forward(self, noise):\n",
        "        '''\n",
        "        Function for completing a forward pass of the generator: Given a noise tensor,\n",
        "        returns generated images.\n",
        "        Parameters:\n",
        "            noise: a noise tensor with dimensions (n_samples, z_dim)\n",
        "        '''\n",
        "        x = noise.view(len(noise), self.z_dim, 1, 1)\n",
        "        return self.gen(x)\n",
        "\n",
        "def get_noise(n_samples, z_dim, device='cpu'):\n",
        "    '''\n",
        "    Function for creating noise vectors: Given the dimensions (n_samples, z_dim)\n",
        "    creates a tensor of that shape filled with random numbers from the normal distribution.\n",
        "    Parameters:\n",
        "      n_samples: the number of samples to generate, a scalar\n",
        "      z_dim: the dimension of the noise vector, a scalar\n",
        "      device: the device type\n",
        "    '''\n",
        "    return torch.randn(n_samples, z_dim, device=device)"
      ],
      "metadata": {
        "id": "meU29M4jbHyy"
      },
      "execution_count": 2,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "#### Critic"
      ],
      "metadata": {
        "id": "cNUdvmhabNS7"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "class Critic(nn.Module):\n",
        "    '''\n",
        "    Critic Class\n",
        "    Values:\n",
        "        im_chan: the number of channels in the images, fitted for the dataset used, a scalar\n",
        "              (MNIST is black-and-white, so 1 channel is your default)\n",
        "        hidden_dim: the inner dimension, a scalar\n",
        "    '''\n",
        "    def __init__(self, im_chan=1, hidden_dim=64):\n",
        "        super(Critic, self).__init__()\n",
        "        self.crit = nn.Sequential(\n",
        "            self.make_crit_block(im_chan, hidden_dim),\n",
        "            self.make_crit_block(hidden_dim, hidden_dim * 2),\n",
        "            self.make_crit_block(hidden_dim * 2, 1, final_layer=True),\n",
        "        )\n",
        "\n",
        "    def make_crit_block(self, input_channels, output_channels, kernel_size=4, stride=2, final_layer=False):\n",
        "        '''\n",
        "        Function to return a sequence of operations corresponding to a critic block of DCGAN;\n",
        "        a convolution, a batchnorm (except in the final layer), and an activation (except in the final layer).\n",
        "        Parameters:\n",
        "            input_channels: how many channels the input feature representation has\n",
        "            output_channels: how many channels the output feature representation should have\n",
        "            kernel_size: the size of each convolutional filter, equivalent to (kernel_size, kernel_size)\n",
        "            stride: the stride of the convolution\n",
        "            final_layer: a boolean, true if it is the final layer and false otherwise \n",
        "                      (affects activation and batchnorm)\n",
        "        '''\n",
        "        if not final_layer:\n",
        "            return nn.Sequential(\n",
        "                nn.Conv2d(input_channels, output_channels, kernel_size, stride),\n",
        "                nn.BatchNorm2d(output_channels),\n",
        "                nn.LeakyReLU(0.2, inplace=True),\n",
        "            )\n",
        "        else:\n",
        "            return nn.Sequential(\n",
        "                nn.Conv2d(input_channels, output_channels, kernel_size, stride),\n",
        "            )\n",
        "\n",
        "    def forward(self, image):\n",
        "        '''\n",
        "        Function for completing a forward pass of the critic: Given an image tensor, \n",
        "        returns a 1-dimension tensor representing fake/real.\n",
        "        Parameters:\n",
        "            image: a flattened image tensor with dimension (im_chan)\n",
        "        '''\n",
        "        crit_pred = self.crit(image)\n",
        "        return crit_pred.view(len(crit_pred), -1)"
      ],
      "metadata": {
        "id": "OTSy4qpNbPGE"
      },
      "execution_count": 3,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Training Initializations"
      ],
      "metadata": {
        "id": "17PB_zk_bRV8"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "n_epochs = 100\n",
        "z_dim = 64\n",
        "display_step = 500\n",
        "batch_size = 128\n",
        "lr = 0.0002\n",
        "beta_1 = 0.5\n",
        "beta_2 = 0.999\n",
        "c_lambda = 10\n",
        "crit_repeats = 5\n",
        "device = 'cuda'\n",
        "\n",
        "transform = transforms.Compose([\n",
        "    transforms.ToTensor(),\n",
        "    transforms.Normalize((0.5,), (0.5,)),\n",
        "])\n",
        "\n",
        "dataloader = DataLoader(\n",
        "    MNIST('.', download=True, transform=transform),\n",
        "    batch_size=batch_size,\n",
        "    shuffle=True)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 443,
          "referenced_widgets": [
            "f9bdc6b94ab14f5d8383c7fb54409003",
            "b214dd356fcd49b1ab683b1750d4e9a0",
            "ac34f33cb35942368f26656196f757c4",
            "cd9b035469ff4c989136d30fa669bf5c",
            "186dd44a4df14e6782962ae79d833426",
            "886e4617d36947e4954a35a4076fc225",
            "bb768bc8f6af48dfa233dad8a7f9af3e",
            "0f193b67cf194d69bd1e4b50a743f05b",
            "7effed533e8c4beba8860b6d49bdbc8f",
            "49d8c1cf3379442b97171e1dcf63e0ac",
            "6d85927cebd144a0a7327168a488d3e6",
            "0b8992d775b14889a617ddff6efe1362",
            "be62ad0778ba4be5bf17c7bc5f181c9a",
            "3af65e68e4a54e46af18207756b30095",
            "64a70b0bd2e04dde8b1365c1d105ef2a",
            "2dd8d57f2d574c4cb5a89adf0c25037b",
            "fc28876ba0da487fb78f7d1d279d6934",
            "ff87c106bb2c45f5a28a8af79a259730",
            "0bc9c686e46c4e73a8b6622ae59846b3",
            "1fce092a7796478fa4b86693feabb339",
            "744b6f6e58dc473bb1b16740217d8709",
            "d9315352a5e143499fa0cd31276ef4f9",
            "b31ce68c359f42a7b12622acbaf031da",
            "1c21d222feb044c5917f0eca4f85cb8e",
            "dff6c6c5ad204abeabb2106b6e83fb20",
            "cf804a98df054874b915eb21da9ca536",
            "5815ff5cf1bd4c3287ea857088054627",
            "ed55a18fe7d7414da586b21c62ce74a6",
            "e79c4a4353ac431fbf6f69bf88a65cb4",
            "1a9cb7b02268475caefb554d7c80de1d",
            "33ce54d2d70e49eabb57cf9e559c17b0",
            "7faaecddc5ce4f28b21392271f01dbaa",
            "756f6b248d7e47dcae8fb8bd8d3c9f80",
            "9bf1ec46cf694cf8bd3f453a441ac130",
            "966ddd3b8dc2450092eac2bbc213d000",
            "74c8baaf41794b73a3d4635279d940a9",
            "bfee2d5dafad4a6eb7ac1424b02d941e",
            "62e23952e4ce403c8f867aaedbcaa43b",
            "f867d1a6388b49ee9662126dcea68061",
            "79d9268b3a654cc5a7fe8a4125384658",
            "21b6492411074f6bb0088655bb062217",
            "2789c61c4694478c907fdf83703f3d81",
            "a4054af378c34659b8ba0bf0a81f2e6b",
            "0360b616956444df8977675e5391cbb9"
          ]
        },
        "id": "1pQpu_ksbWGm",
        "outputId": "de2a631b-01ea-4c59-e43b-c5b9611cbaf1"
      },
      "execution_count": 4,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Downloading http://yann.lecun.com/exdb/mnist/train-images-idx3-ubyte.gz\n",
            "Downloading http://yann.lecun.com/exdb/mnist/train-images-idx3-ubyte.gz to ./MNIST/raw/train-images-idx3-ubyte.gz\n"
          ]
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "  0%|          | 0/9912422 [00:00<?, ?it/s]"
            ],
            "application/vnd.jupyter.widget-view+json": {
              "version_major": 2,
              "version_minor": 0,
              "model_id": "f9bdc6b94ab14f5d8383c7fb54409003"
            }
          },
          "metadata": {}
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Extracting ./MNIST/raw/train-images-idx3-ubyte.gz to ./MNIST/raw\n",
            "\n",
            "Downloading http://yann.lecun.com/exdb/mnist/train-labels-idx1-ubyte.gz\n",
            "Downloading http://yann.lecun.com/exdb/mnist/train-labels-idx1-ubyte.gz to ./MNIST/raw/train-labels-idx1-ubyte.gz\n"
          ]
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "  0%|          | 0/28881 [00:00<?, ?it/s]"
            ],
            "application/vnd.jupyter.widget-view+json": {
              "version_major": 2,
              "version_minor": 0,
              "model_id": "0b8992d775b14889a617ddff6efe1362"
            }
          },
          "metadata": {}
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Extracting ./MNIST/raw/train-labels-idx1-ubyte.gz to ./MNIST/raw\n",
            "\n",
            "Downloading http://yann.lecun.com/exdb/mnist/t10k-images-idx3-ubyte.gz\n",
            "Downloading http://yann.lecun.com/exdb/mnist/t10k-images-idx3-ubyte.gz to ./MNIST/raw/t10k-images-idx3-ubyte.gz\n"
          ]
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "  0%|          | 0/1648877 [00:00<?, ?it/s]"
            ],
            "application/vnd.jupyter.widget-view+json": {
              "version_major": 2,
              "version_minor": 0,
              "model_id": "b31ce68c359f42a7b12622acbaf031da"
            }
          },
          "metadata": {}
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Extracting ./MNIST/raw/t10k-images-idx3-ubyte.gz to ./MNIST/raw\n",
            "\n",
            "Downloading http://yann.lecun.com/exdb/mnist/t10k-labels-idx1-ubyte.gz\n",
            "Downloading http://yann.lecun.com/exdb/mnist/t10k-labels-idx1-ubyte.gz to ./MNIST/raw/t10k-labels-idx1-ubyte.gz\n"
          ]
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "  0%|          | 0/4542 [00:00<?, ?it/s]"
            ],
            "application/vnd.jupyter.widget-view+json": {
              "version_major": 2,
              "version_minor": 0,
              "model_id": "9bf1ec46cf694cf8bd3f453a441ac130"
            }
          },
          "metadata": {}
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Extracting ./MNIST/raw/t10k-labels-idx1-ubyte.gz to ./MNIST/raw\n",
            "\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "gen = Generator(z_dim).to(device)\n",
        "gen_opt = torch.optim.Adam(gen.parameters(), lr=lr, betas=(beta_1, beta_2))\n",
        "crit = Critic().to(device) \n",
        "crit_opt = torch.optim.Adam(crit.parameters(), lr=lr, betas=(beta_1, beta_2))\n",
        "\n",
        "def weights_init(m):\n",
        "    if isinstance(m, nn.Conv2d) or isinstance(m, nn.ConvTranspose2d):\n",
        "        torch.nn.init.normal_(m.weight, 0.0, 0.02)\n",
        "    if isinstance(m, nn.BatchNorm2d):\n",
        "        torch.nn.init.normal_(m.weight, 0.0, 0.02)\n",
        "        torch.nn.init.constant_(m.bias, 0)\n",
        "gen = gen.apply(weights_init)\n",
        "crit = crit.apply(weights_init)\n"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 381
        },
        "id": "vi8LKi-KbZwh",
        "outputId": "059010c6-5244-4d87-975f-7afb7d82d4e0"
      },
      "execution_count": 5,
      "outputs": [
        {
          "output_type": "error",
          "ename": "RuntimeError",
          "evalue": "ignored",
          "traceback": [
            "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
            "\u001b[0;31mRuntimeError\u001b[0m                              Traceback (most recent call last)",
            "\u001b[0;32m<ipython-input-5-fcdedc729d7f>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[0;32m----> 1\u001b[0;31m \u001b[0mgen\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mGenerator\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mz_dim\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mto\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mdevice\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m      2\u001b[0m \u001b[0mgen_opt\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mtorch\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0moptim\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mAdam\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mgen\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mparameters\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mlr\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mlr\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mbetas\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mbeta_1\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mbeta_2\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      3\u001b[0m \u001b[0mcrit\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mCritic\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mto\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mdevice\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      4\u001b[0m \u001b[0mcrit_opt\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mtorch\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0moptim\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mAdam\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mcrit\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mparameters\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mlr\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mlr\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mbetas\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mbeta_1\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mbeta_2\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      5\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m/usr/local/lib/python3.9/dist-packages/torch/nn/modules/module.py\u001b[0m in \u001b[0;36mto\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m    987\u001b[0m             \u001b[0;32mreturn\u001b[0m \u001b[0mt\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mto\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mdevice\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mdtype\u001b[0m \u001b[0;32mif\u001b[0m \u001b[0mt\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mis_floating_point\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;32mor\u001b[0m \u001b[0mt\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mis_complex\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;32melse\u001b[0m \u001b[0;32mNone\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mnon_blocking\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    988\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 989\u001b[0;31m         \u001b[0;32mreturn\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_apply\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mconvert\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    990\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    991\u001b[0m     def register_backward_hook(\n",
            "\u001b[0;32m/usr/local/lib/python3.9/dist-packages/torch/nn/modules/module.py\u001b[0m in \u001b[0;36m_apply\u001b[0;34m(self, fn)\u001b[0m\n\u001b[1;32m    639\u001b[0m     \u001b[0;32mdef\u001b[0m \u001b[0m_apply\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mfn\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    640\u001b[0m         \u001b[0;32mfor\u001b[0m \u001b[0mmodule\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mchildren\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 641\u001b[0;31m             \u001b[0mmodule\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_apply\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mfn\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    642\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    643\u001b[0m         \u001b[0;32mdef\u001b[0m \u001b[0mcompute_should_use_set_data\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mtensor\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mtensor_applied\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m/usr/local/lib/python3.9/dist-packages/torch/nn/modules/module.py\u001b[0m in \u001b[0;36m_apply\u001b[0;34m(self, fn)\u001b[0m\n\u001b[1;32m    639\u001b[0m     \u001b[0;32mdef\u001b[0m \u001b[0m_apply\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mfn\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    640\u001b[0m         \u001b[0;32mfor\u001b[0m \u001b[0mmodule\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mchildren\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 641\u001b[0;31m             \u001b[0mmodule\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_apply\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mfn\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    642\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    643\u001b[0m         \u001b[0;32mdef\u001b[0m \u001b[0mcompute_should_use_set_data\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mtensor\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mtensor_applied\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m/usr/local/lib/python3.9/dist-packages/torch/nn/modules/module.py\u001b[0m in \u001b[0;36m_apply\u001b[0;34m(self, fn)\u001b[0m\n\u001b[1;32m    639\u001b[0m     \u001b[0;32mdef\u001b[0m \u001b[0m_apply\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mfn\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    640\u001b[0m         \u001b[0;32mfor\u001b[0m \u001b[0mmodule\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mchildren\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 641\u001b[0;31m             \u001b[0mmodule\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_apply\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mfn\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    642\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    643\u001b[0m         \u001b[0;32mdef\u001b[0m \u001b[0mcompute_should_use_set_data\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mtensor\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mtensor_applied\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m/usr/local/lib/python3.9/dist-packages/torch/nn/modules/module.py\u001b[0m in \u001b[0;36m_apply\u001b[0;34m(self, fn)\u001b[0m\n\u001b[1;32m    662\u001b[0m             \u001b[0;31m# `with torch.no_grad():`\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    663\u001b[0m             \u001b[0;32mwith\u001b[0m \u001b[0mtorch\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mno_grad\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 664\u001b[0;31m                 \u001b[0mparam_applied\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mfn\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mparam\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    665\u001b[0m             \u001b[0mshould_use_set_data\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mcompute_should_use_set_data\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mparam\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mparam_applied\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    666\u001b[0m             \u001b[0;32mif\u001b[0m \u001b[0mshould_use_set_data\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m/usr/local/lib/python3.9/dist-packages/torch/nn/modules/module.py\u001b[0m in \u001b[0;36mconvert\u001b[0;34m(t)\u001b[0m\n\u001b[1;32m    985\u001b[0m                 return t.to(device, dtype if t.is_floating_point() or t.is_complex() else None,\n\u001b[1;32m    986\u001b[0m                             non_blocking, memory_format=convert_to_format)\n\u001b[0;32m--> 987\u001b[0;31m             \u001b[0;32mreturn\u001b[0m \u001b[0mt\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mto\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mdevice\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mdtype\u001b[0m \u001b[0;32mif\u001b[0m \u001b[0mt\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mis_floating_point\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;32mor\u001b[0m \u001b[0mt\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mis_complex\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;32melse\u001b[0m \u001b[0;32mNone\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mnon_blocking\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    988\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    989\u001b[0m         \u001b[0;32mreturn\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_apply\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mconvert\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m/usr/local/lib/python3.9/dist-packages/torch/cuda/__init__.py\u001b[0m in \u001b[0;36m_lazy_init\u001b[0;34m()\u001b[0m\n\u001b[1;32m    227\u001b[0m         \u001b[0;32mif\u001b[0m \u001b[0;34m'CUDA_MODULE_LOADING'\u001b[0m \u001b[0;32mnot\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mos\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0menviron\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    228\u001b[0m             \u001b[0mos\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0menviron\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;34m'CUDA_MODULE_LOADING'\u001b[0m\u001b[0;34m]\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;34m'LAZY'\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 229\u001b[0;31m         \u001b[0mtorch\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_C\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_cuda_init\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    230\u001b[0m         \u001b[0;31m# Some of the queued calls may reentrantly call _lazy_init();\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    231\u001b[0m         \u001b[0;31m# we need to just return without initializing in that case.\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;31mRuntimeError\u001b[0m: No CUDA GPUs are available"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Gradient Penalty"
      ],
      "metadata": {
        "id": "rBMC0aUObcPT"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "def get_gradient(crit, real, fake, epsilon):\n",
        "    '''\n",
        "    Return the gradient of the critic's scores with respect to mixes of real and fake images.\n",
        "    Parameters:\n",
        "        crit: the critic model\n",
        "        real: a batch of real images\n",
        "        fake: a batch of fake images\n",
        "        epsilon: a vector of the uniformly random proportions of real/fake per mixed image\n",
        "    Returns:\n",
        "        gradient: the gradient of the critic's scores, with respect to the mixed image\n",
        "    '''\n",
        "    # Mix the images together\n",
        "    mixed_images = real * epsilon + fake * (1 - epsilon)\n",
        "\n",
        "    # Calculate the critic's scores on the mixed images\n",
        "    mixed_scores = crit(mixed_images)\n",
        "    \n",
        "    # Take the gradient of the scores with respect to the images\n",
        "    gradient = torch.autograd.grad(\n",
        "        # Note: You need to take the gradient of outputs with respect to inputs.\n",
        "        # This documentation may be useful, but it should not be necessary:\n",
        "        # https://pytorch.org/docs/stable/autograd.html#torch.autograd.grad\n",
        "        inputs=mixed_images,\n",
        "        outputs=mixed_scores,\n",
        "        # These other parameters have to do with the pytorch autograd engine works\n",
        "        grad_outputs=torch.ones_like(mixed_scores), \n",
        "        create_graph=True,\n",
        "        retain_graph=True,\n",
        "    )[0]\n",
        "    return gradient"
      ],
      "metadata": {
        "id": "ULZEa3DXbpE3"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "def gradient_penalty(gradient):\n",
        "    '''\n",
        "    Return the gradient penalty, given a gradient.\n",
        "    Given a batch of image gradients, you calculate the magnitude of each image's gradient\n",
        "    and penalize the mean quadratic distance of each magnitude to 1.\n",
        "    Parameters:\n",
        "        gradient: the gradient of the critic's scores, with respect to the mixed image\n",
        "    Returns:\n",
        "        penalty: the gradient penalty\n",
        "    '''\n",
        "    # Flatten the gradients so that each row captures one image\n",
        "    gradient = gradient.view(len(gradient), -1)\n",
        "\n",
        "    # Calculate the magnitude of every row\n",
        "    gradient_norm = gradient.norm(2, dim=1)\n",
        "    \n",
        "    # Penalize the mean squared distance of the gradient norms from 1\n",
        "    penalty = torch.mean((gradient_norm - 1)**2)\n",
        "    \n",
        "    return penalty"
      ],
      "metadata": {
        "id": "SYDf2DzYbtJV"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Losses"
      ],
      "metadata": {
        "id": "WufdNweAbx7j"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "def get_gen_loss(crit_fake_pred):\n",
        "    '''\n",
        "    Return the loss of a generator given the critic's scores of the generator's fake images.\n",
        "    Parameters:\n",
        "        crit_fake_pred: the critic's scores of the fake images\n",
        "    Returns:\n",
        "        gen_loss: a scalar loss value for the current batch of the generator\n",
        "    '''\n",
        "    gen_loss = -1. * torch.mean(crit_fake_pred)\n",
        "    return gen_loss"
      ],
      "metadata": {
        "id": "I_uVVofrb0bX"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "def get_crit_loss(crit_fake_pred, crit_real_pred, gp, c_lambda):\n",
        "    '''\n",
        "    Return the loss of a critic given the critic's scores for fake and real images,\n",
        "    the gradient penalty, and gradient penalty weight.\n",
        "    Parameters:\n",
        "        crit_fake_pred: the critic's scores of the fake images\n",
        "        crit_real_pred: the critic's scores of the real images\n",
        "        gp: the unweighted gradient penalty\n",
        "        c_lambda: the current weight of the gradient penalty \n",
        "    Returns:\n",
        "        crit_loss: a scalar for the critic's loss, accounting for the relevant factors\n",
        "    '''\n",
        "\n",
        "    crit_loss = crit_loss = torch.mean(crit_fake_pred) - torch.mean(crit_real_pred) + c_lambda * gp\n",
        "\n",
        "    return crit_loss"
      ],
      "metadata": {
        "id": "5tSFrkuFb3SP"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Putting It All Together\n",
        "\n",
        "Before you put everything together, there are a few things to note.\n",
        "1.   Even on GPU, the **training will run more slowly** than previous labs because the gradient penalty requires you to compute the gradient of a gradient -- this means potentially a few minutes per epoch! For best results, run this for as long as you can while on GPU.\n",
        "2.   One important difference from earlier versions is that you will **update the critic multiple times** every time you update the generator This helps prevent the generator from overpowering the critic. Sometimes, you might see the reverse, with the generator updated more times than the critic. This depends on architectural (e.g. the depth and width of the network) and algorithmic choices (e.g. which loss you're using). \n",
        "3.   WGAN-GP isn't necessarily meant to improve overall performance of a GAN, but just **increases stability** and avoids mode collapse. In general, a WGAN will be able to train in a much more stable way than the vanilla DCGAN from last assignment, though it will generally run a bit slower. You should also be able to train your model for more epochs without it collapsing.\n"
      ],
      "metadata": {
        "id": "HZ_OCeSAb6Kx"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "import matplotlib.pyplot as plt\n",
        "\n",
        "cur_step = 0\n",
        "generator_losses = []\n",
        "critic_losses = []\n",
        "for epoch in range(n_epochs):\n",
        "    # Dataloader returns the batches\n",
        "    for real, _ in tqdm(dataloader):\n",
        "        cur_batch_size = len(real)\n",
        "        real = real.to(device)\n",
        "\n",
        "        mean_iteration_critic_loss = 0\n",
        "        for _ in range(crit_repeats):\n",
        "            ### Update critic ###\n",
        "            crit_opt.zero_grad()\n",
        "            fake_noise = get_noise(cur_batch_size, z_dim, device=device)\n",
        "            fake = gen(fake_noise)\n",
        "            crit_fake_pred = crit(fake.detach())\n",
        "            crit_real_pred = crit(real)\n",
        "\n",
        "            epsilon = torch.rand(len(real), 1, 1, 1, device=device, requires_grad=True)\n",
        "            gradient = get_gradient(crit, real, fake.detach(), epsilon)\n",
        "            gp = gradient_penalty(gradient)\n",
        "            crit_loss = get_crit_loss(crit_fake_pred, crit_real_pred, gp, c_lambda)\n",
        "\n",
        "            # Keep track of the average critic loss in this batch\n",
        "            mean_iteration_critic_loss += crit_loss.item() / crit_repeats\n",
        "            # Update gradients\n",
        "            crit_loss.backward(retain_graph=True)\n",
        "            # Update optimizer\n",
        "            crit_opt.step()\n",
        "        critic_losses += [mean_iteration_critic_loss]\n",
        "\n",
        "        ### Update generator ###\n",
        "        gen_opt.zero_grad()\n",
        "        fake_noise_2 = get_noise(cur_batch_size, z_dim, device=device)\n",
        "        fake_2 = gen(fake_noise_2)\n",
        "        crit_fake_pred = crit(fake_2)\n",
        "        \n",
        "        gen_loss = get_gen_loss(crit_fake_pred)\n",
        "        gen_loss.backward()\n",
        "\n",
        "        # Update the weights\n",
        "        gen_opt.step()\n",
        "\n",
        "        # Keep track of the average generator loss\n",
        "        generator_losses += [gen_loss.item()]\n",
        "\n",
        "        ### Visualization code ###\n",
        "        if cur_step % display_step == 0 and cur_step > 0:\n",
        "            gen_mean = sum(generator_losses[-display_step:]) / display_step\n",
        "            crit_mean = sum(critic_losses[-display_step:]) / display_step\n",
        "            print(f\"Epoch {epoch}, step {cur_step}: Generator loss: {gen_mean}, critic loss: {crit_mean}\")\n",
        "            show_tensor_images(fake)\n",
        "            show_tensor_images(real)\n",
        "            step_bins = 20\n",
        "            num_examples = (len(generator_losses) // step_bins) * step_bins\n",
        "            plt.plot(\n",
        "                range(num_examples // step_bins), \n",
        "                torch.Tensor(generator_losses[:num_examples]).view(-1, step_bins).mean(1),\n",
        "                label=\"Generator Loss\"\n",
        "            )\n",
        "            plt.plot(\n",
        "                range(num_examples // step_bins), \n",
        "                torch.Tensor(critic_losses[:num_examples]).view(-1, step_bins).mean(1),\n",
        "                label=\"Critic Loss\"\n",
        "            )\n",
        "            plt.legend()\n",
        "            plt.show()\n",
        "\n",
        "        cur_step += 1\n"
      ],
      "metadata": {
        "id": "TzAHUq1ZcKXc"
      },
      "execution_count": null,
      "outputs": []
    }
  ]
}