{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "Vision Transformer AugReg",
      "provenance": [],
      "collapsed_sections": [],
      "toc_visible": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    },
    "accelerator": "GPU",
    "widgets": {
      "application/vnd.jupyter.widget-state+json": {
        "225ebb8a26264dca90b312edb1e73fb6": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HBoxModel",
          "state": {
            "_view_name": "HBoxView",
            "_dom_classes": [],
            "_model_name": "HBoxModel",
            "_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_7d3d39e266ab4a81983fbf829795f3a7",
            "_model_module": "@jupyter-widgets/controls",
            "children": [
              "IPY_MODEL_c529bc24ac1943dd8628873d1648c275",
              "IPY_MODEL_d4e82b87ab4a485c9e15f3fd6e77a969"
            ]
          }
        },
        "7d3d39e266ab4a81983fbf829795f3a7": {
          "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
          }
        },
        "c529bc24ac1943dd8628873d1648c275": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "FloatProgressModel",
          "state": {
            "_view_name": "ProgressView",
            "style": "IPY_MODEL_7bdc1b98ea574b8e99fc1bd3bbd2b456",
            "_dom_classes": [],
            "description": "Dl Completed...: 100%",
            "_model_name": "FloatProgressModel",
            "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_63c6ec28cbad46359e38d1f921b15ffa"
          }
        },
        "d4e82b87ab4a485c9e15f3fd6e77a969": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HTMLModel",
          "state": {
            "_view_name": "HTMLView",
            "style": "IPY_MODEL_d9d5245d91fe42999d8b8693674880f7",
            "_dom_classes": [],
            "description": "",
            "_model_name": "HTMLModel",
            "placeholder": "​",
            "_view_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "value": " 2/2 [00:47\u0026lt;00:00, 23.84s/ url]",
            "_view_count": null,
            "_view_module_version": "1.5.0",
            "description_tooltip": null,
            "_model_module": "@jupyter-widgets/controls",
            "layout": "IPY_MODEL_113db5d71db84749805a9394e89b9e0b"
          }
        },
        "7bdc1b98ea574b8e99fc1bd3bbd2b456": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "ProgressStyleModel",
          "state": {
            "_view_name": "StyleView",
            "_model_name": "ProgressStyleModel",
            "description_width": "initial",
            "_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"
          }
        },
        "63c6ec28cbad46359e38d1f921b15ffa": {
          "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
          }
        },
        "d9d5245d91fe42999d8b8693674880f7": {
          "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"
          }
        },
        "113db5d71db84749805a9394e89b9e0b": {
          "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
          }
        },
        "a5a294a2b18f4c078b04b8126be1d84b": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HBoxModel",
          "state": {
            "_view_name": "HBoxView",
            "_dom_classes": [],
            "_model_name": "HBoxModel",
            "_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_c2ce9e71f07a4be59374466f7a8f33ff",
            "_model_module": "@jupyter-widgets/controls",
            "children": [
              "IPY_MODEL_608ff4fab08743d8ade2cf99ac8dcd69",
              "IPY_MODEL_107c6e554e9a4a2a913f1684761d3736"
            ]
          }
        },
        "c2ce9e71f07a4be59374466f7a8f33ff": {
          "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
          }
        },
        "608ff4fab08743d8ade2cf99ac8dcd69": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "FloatProgressModel",
          "state": {
            "_view_name": "ProgressView",
            "style": "IPY_MODEL_0e51cd83a88140728b937542605c7e62",
            "_dom_classes": [],
            "description": "Dl Size...: 100%",
            "_model_name": "FloatProgressModel",
            "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_a6f978d87f114e4ea5fd99ca45c67ba1"
          }
        },
        "107c6e554e9a4a2a913f1684761d3736": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HTMLModel",
          "state": {
            "_view_name": "HTMLView",
            "style": "IPY_MODEL_c9598f8547494601afc86dbdcfc88516",
            "_dom_classes": [],
            "description": "",
            "_model_name": "HTMLModel",
            "placeholder": "​",
            "_view_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "value": " 773/773 [00:47\u0026lt;00:00, 16.23 MiB/s]",
            "_view_count": null,
            "_view_module_version": "1.5.0",
            "description_tooltip": null,
            "_model_module": "@jupyter-widgets/controls",
            "layout": "IPY_MODEL_c625f7758d544bb6937f265f209cd2c8"
          }
        },
        "0e51cd83a88140728b937542605c7e62": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "ProgressStyleModel",
          "state": {
            "_view_name": "StyleView",
            "_model_name": "ProgressStyleModel",
            "description_width": "initial",
            "_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"
          }
        },
        "a6f978d87f114e4ea5fd99ca45c67ba1": {
          "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
          }
        },
        "c9598f8547494601afc86dbdcfc88516": {
          "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"
          }
        },
        "c625f7758d544bb6937f265f209cd2c8": {
          "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
          }
        },
        "132048250a42471785d742ecfa679a7b": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HBoxModel",
          "state": {
            "_view_name": "HBoxView",
            "_dom_classes": [],
            "_model_name": "HBoxModel",
            "_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_a9d7f91c071345ce9012a39adc3037b5",
            "_model_module": "@jupyter-widgets/controls",
            "children": [
              "IPY_MODEL_79577fafaa0c46cc85e057d50e1f640f",
              "IPY_MODEL_a0e3deff575141949907cb8647a769be"
            ]
          }
        },
        "a9d7f91c071345ce9012a39adc3037b5": {
          "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
          }
        },
        "79577fafaa0c46cc85e057d50e1f640f": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "FloatProgressModel",
          "state": {
            "_view_name": "ProgressView",
            "style": "IPY_MODEL_770dd8577d454c2e82774043376809c1",
            "_dom_classes": [],
            "description": "Extraction completed...: 100%",
            "_model_name": "FloatProgressModel",
            "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_a73ff8ef581e44d8a8503f4cf25df02f"
          }
        },
        "a0e3deff575141949907cb8647a769be": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HTMLModel",
          "state": {
            "_view_name": "HTMLView",
            "style": "IPY_MODEL_a54f014c7dff445f846b97281652c14f",
            "_dom_classes": [],
            "description": "",
            "_model_name": "HTMLModel",
            "placeholder": "​",
            "_view_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "value": " 2/2 [00:47\u0026lt;00:00, 23.79s/ file]",
            "_view_count": null,
            "_view_module_version": "1.5.0",
            "description_tooltip": null,
            "_model_module": "@jupyter-widgets/controls",
            "layout": "IPY_MODEL_b3f5112e065d4cb884c514a8e4097b38"
          }
        },
        "770dd8577d454c2e82774043376809c1": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "ProgressStyleModel",
          "state": {
            "_view_name": "StyleView",
            "_model_name": "ProgressStyleModel",
            "description_width": "initial",
            "_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"
          }
        },
        "a73ff8ef581e44d8a8503f4cf25df02f": {
          "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
          }
        },
        "a54f014c7dff445f846b97281652c14f": {
          "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"
          }
        },
        "b3f5112e065d4cb884c514a8e4097b38": {
          "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
          }
        },
        "43a06b56ab8742cabec36cd19b51945c": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HBoxModel",
          "state": {
            "_view_name": "HBoxView",
            "_dom_classes": [],
            "_model_name": "HBoxModel",
            "_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_b0e5b9a1cd5c4b918300b2006daacbdb",
            "_model_module": "@jupyter-widgets/controls",
            "children": [
              "IPY_MODEL_acc7c1e9e5ec470bb7d22a49fbab75dc",
              "IPY_MODEL_5f70500c82aa4c189b85c4765047305f"
            ]
          }
        },
        "b0e5b9a1cd5c4b918300b2006daacbdb": {
          "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
          }
        },
        "acc7c1e9e5ec470bb7d22a49fbab75dc": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "FloatProgressModel",
          "state": {
            "_view_name": "ProgressView",
            "style": "IPY_MODEL_df9bdd9cfdb74967a5cd83fbcbc3d57e",
            "_dom_classes": [],
            "description": "",
            "_model_name": "FloatProgressModel",
            "bar_style": "info",
            "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_8529b470a86640a1819098d4b809d7c8"
          }
        },
        "5f70500c82aa4c189b85c4765047305f": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HTMLModel",
          "state": {
            "_view_name": "HTMLView",
            "style": "IPY_MODEL_8f0adf499e434dc2863bdf6537c08454",
            "_dom_classes": [],
            "description": "",
            "_model_name": "HTMLModel",
            "placeholder": "​",
            "_view_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "value": " 3680/0 [00:02\u0026lt;00:00, 1339.10 examples/s]",
            "_view_count": null,
            "_view_module_version": "1.5.0",
            "description_tooltip": null,
            "_model_module": "@jupyter-widgets/controls",
            "layout": "IPY_MODEL_e5cafaad07cd45dfb2329b4fa3b18ad4"
          }
        },
        "df9bdd9cfdb74967a5cd83fbcbc3d57e": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "ProgressStyleModel",
          "state": {
            "_view_name": "StyleView",
            "_model_name": "ProgressStyleModel",
            "description_width": "initial",
            "_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"
          }
        },
        "8529b470a86640a1819098d4b809d7c8": {
          "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
          }
        },
        "8f0adf499e434dc2863bdf6537c08454": {
          "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"
          }
        },
        "e5cafaad07cd45dfb2329b4fa3b18ad4": {
          "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
          }
        },
        "f0da817110a947d2b0b54a43c7879ec6": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HBoxModel",
          "state": {
            "_view_name": "HBoxView",
            "_dom_classes": [],
            "_model_name": "HBoxModel",
            "_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_13b28f5de68b428bba06b648c201c5e2",
            "_model_module": "@jupyter-widgets/controls",
            "children": [
              "IPY_MODEL_5a54036091c84c6582ccf8ea545f7f4c",
              "IPY_MODEL_ee0a499dc13b4a6eb36627a242d6e92a"
            ]
          }
        },
        "13b28f5de68b428bba06b648c201c5e2": {
          "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
          }
        },
        "5a54036091c84c6582ccf8ea545f7f4c": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "FloatProgressModel",
          "state": {
            "_view_name": "ProgressView",
            "style": "IPY_MODEL_febe24498bc0445ab747c120c6abe261",
            "_dom_classes": [],
            "description": " 89%",
            "_model_name": "FloatProgressModel",
            "bar_style": "danger",
            "max": 3680,
            "_view_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "value": 3278,
            "_view_count": null,
            "_view_module_version": "1.5.0",
            "orientation": "horizontal",
            "min": 0,
            "description_tooltip": null,
            "_model_module": "@jupyter-widgets/controls",
            "layout": "IPY_MODEL_56f06ed9e69d40f9b5f39e1550fd2eb1"
          }
        },
        "ee0a499dc13b4a6eb36627a242d6e92a": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HTMLModel",
          "state": {
            "_view_name": "HTMLView",
            "style": "IPY_MODEL_7d7deac790b4445cb170356ae88b36e9",
            "_dom_classes": [],
            "description": "",
            "_model_name": "HTMLModel",
            "placeholder": "​",
            "_view_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "value": " 3278/3680 [00:00\u0026lt;00:00, 4546.66 examples/s]",
            "_view_count": null,
            "_view_module_version": "1.5.0",
            "description_tooltip": null,
            "_model_module": "@jupyter-widgets/controls",
            "layout": "IPY_MODEL_747b920f1510475dad43d49ba782be4d"
          }
        },
        "febe24498bc0445ab747c120c6abe261": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "ProgressStyleModel",
          "state": {
            "_view_name": "StyleView",
            "_model_name": "ProgressStyleModel",
            "description_width": "initial",
            "_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"
          }
        },
        "56f06ed9e69d40f9b5f39e1550fd2eb1": {
          "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
          }
        },
        "7d7deac790b4445cb170356ae88b36e9": {
          "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"
          }
        },
        "747b920f1510475dad43d49ba782be4d": {
          "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
          }
        },
        "3f995dd7906341eb806cb4cfc52c78eb": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HBoxModel",
          "state": {
            "_view_name": "HBoxView",
            "_dom_classes": [],
            "_model_name": "HBoxModel",
            "_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_80048ea0d80c4bf5a9172380e948799d",
            "_model_module": "@jupyter-widgets/controls",
            "children": [
              "IPY_MODEL_716975e14c8f4b59ade02dc968d58b5e",
              "IPY_MODEL_2acfe7938b3543bf8b9542556f5a23b8"
            ]
          }
        },
        "80048ea0d80c4bf5a9172380e948799d": {
          "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
          }
        },
        "716975e14c8f4b59ade02dc968d58b5e": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "FloatProgressModel",
          "state": {
            "_view_name": "ProgressView",
            "style": "IPY_MODEL_817ad1768fa14aa0b548e259d2162944",
            "_dom_classes": [],
            "description": "",
            "_model_name": "FloatProgressModel",
            "bar_style": "info",
            "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_413ba3075e7d4f54989b8de944aa0759"
          }
        },
        "2acfe7938b3543bf8b9542556f5a23b8": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HTMLModel",
          "state": {
            "_view_name": "HTMLView",
            "style": "IPY_MODEL_2a3f564c7c884ff1936d1c6db79ac40b",
            "_dom_classes": [],
            "description": "",
            "_model_name": "HTMLModel",
            "placeholder": "​",
            "_view_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "value": " 3669/0 [00:02\u0026lt;00:00, 1314.42 examples/s]",
            "_view_count": null,
            "_view_module_version": "1.5.0",
            "description_tooltip": null,
            "_model_module": "@jupyter-widgets/controls",
            "layout": "IPY_MODEL_3b863d829bbc4c1ea618ed5dfb3f0538"
          }
        },
        "817ad1768fa14aa0b548e259d2162944": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "ProgressStyleModel",
          "state": {
            "_view_name": "StyleView",
            "_model_name": "ProgressStyleModel",
            "description_width": "initial",
            "_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"
          }
        },
        "413ba3075e7d4f54989b8de944aa0759": {
          "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
          }
        },
        "2a3f564c7c884ff1936d1c6db79ac40b": {
          "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"
          }
        },
        "3b863d829bbc4c1ea618ed5dfb3f0538": {
          "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
          }
        },
        "ee27f121919d4345acd966ad98518b3a": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HBoxModel",
          "state": {
            "_view_name": "HBoxView",
            "_dom_classes": [],
            "_model_name": "HBoxModel",
            "_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_fa485b969a6c4893bbb42904fed24eeb",
            "_model_module": "@jupyter-widgets/controls",
            "children": [
              "IPY_MODEL_7d4bff3e31c4431e91054562d51f80e8",
              "IPY_MODEL_902db9dfbf8543a4956409a22c2efa80"
            ]
          }
        },
        "fa485b969a6c4893bbb42904fed24eeb": {
          "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
          }
        },
        "7d4bff3e31c4431e91054562d51f80e8": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "FloatProgressModel",
          "state": {
            "_view_name": "ProgressView",
            "style": "IPY_MODEL_74a0420bbcc64cd0a662f8e27a23cf5f",
            "_dom_classes": [],
            "description": " 96%",
            "_model_name": "FloatProgressModel",
            "bar_style": "danger",
            "max": 3669,
            "_view_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "value": 3515,
            "_view_count": null,
            "_view_module_version": "1.5.0",
            "orientation": "horizontal",
            "min": 0,
            "description_tooltip": null,
            "_model_module": "@jupyter-widgets/controls",
            "layout": "IPY_MODEL_9ff4e5d0a6cd4953a9b04cdd115de440"
          }
        },
        "902db9dfbf8543a4956409a22c2efa80": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HTMLModel",
          "state": {
            "_view_name": "HTMLView",
            "style": "IPY_MODEL_661c19ff44f048eaa8b21ee203319c68",
            "_dom_classes": [],
            "description": "",
            "_model_name": "HTMLModel",
            "placeholder": "​",
            "_view_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "value": " 3515/3669 [00:01\u0026lt;00:00, 3097.05 examples/s]",
            "_view_count": null,
            "_view_module_version": "1.5.0",
            "description_tooltip": null,
            "_model_module": "@jupyter-widgets/controls",
            "layout": "IPY_MODEL_e9bc91cf5f67406689c7fe90800f729f"
          }
        },
        "74a0420bbcc64cd0a662f8e27a23cf5f": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "ProgressStyleModel",
          "state": {
            "_view_name": "StyleView",
            "_model_name": "ProgressStyleModel",
            "description_width": "initial",
            "_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"
          }
        },
        "9ff4e5d0a6cd4953a9b04cdd115de440": {
          "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
          }
        },
        "661c19ff44f048eaa8b21ee203319c68": {
          "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"
          }
        },
        "e9bc91cf5f67406689c7fe90800f729f": {
          "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
          }
        },
        "5ed291cf3157447684e4c732e15713b8": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HBoxModel",
          "state": {
            "_view_name": "HBoxView",
            "_dom_classes": [],
            "_model_name": "HBoxModel",
            "_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_06ea24f8073b4185965fb4b1494e0f19",
            "_model_module": "@jupyter-widgets/controls",
            "children": [
              "IPY_MODEL_6b6f717b6613435ca113a68a1e1f4126",
              "IPY_MODEL_a9deea48888c423cbb5b3e310b106147"
            ]
          }
        },
        "06ea24f8073b4185965fb4b1494e0f19": {
          "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
          }
        },
        "6b6f717b6613435ca113a68a1e1f4126": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "FloatProgressModel",
          "state": {
            "_view_name": "ProgressView",
            "style": "IPY_MODEL_e2388f624f314ae68d3c6a84127ed400",
            "_dom_classes": [],
            "description": "Dl Completed...: 100%",
            "_model_name": "FloatProgressModel",
            "bar_style": "success",
            "max": 5,
            "_view_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "value": 5,
            "_view_count": null,
            "_view_module_version": "1.5.0",
            "orientation": "horizontal",
            "min": 0,
            "description_tooltip": null,
            "_model_module": "@jupyter-widgets/controls",
            "layout": "IPY_MODEL_d51511f4aa504dc2a56c089a167b1c6a"
          }
        },
        "a9deea48888c423cbb5b3e310b106147": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HTMLModel",
          "state": {
            "_view_name": "HTMLView",
            "style": "IPY_MODEL_05ec97fde7084f5c9cf54b707bf0edf1",
            "_dom_classes": [],
            "description": "",
            "_model_name": "HTMLModel",
            "placeholder": "​",
            "_view_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "value": " 5/5 [00:02\u0026lt;00:00,  1.73 file/s]",
            "_view_count": null,
            "_view_module_version": "1.5.0",
            "description_tooltip": null,
            "_model_module": "@jupyter-widgets/controls",
            "layout": "IPY_MODEL_61ee19bf89fb4955bda1e759dfd99cb7"
          }
        },
        "e2388f624f314ae68d3c6a84127ed400": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "ProgressStyleModel",
          "state": {
            "_view_name": "StyleView",
            "_model_name": "ProgressStyleModel",
            "description_width": "initial",
            "_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"
          }
        },
        "d51511f4aa504dc2a56c089a167b1c6a": {
          "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
          }
        },
        "05ec97fde7084f5c9cf54b707bf0edf1": {
          "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"
          }
        },
        "61ee19bf89fb4955bda1e759dfd99cb7": {
          "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": "markdown",
      "metadata": {
        "id": "sXhZm0kpPpH6"
      },
      "source": [
        "##### Copyright 2021 Google LLC."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "cellView": "form",
        "id": "KfmzfvFxPuk7"
      },
      "source": [
        "#@title Licensed under the Apache License, Version 2.0 (the \"License\");\n",
        "# Licensed under the Apache License, Version 2.0 (the \"License\");\n",
        "# you may not use this file except in compliance with the License.\n",
        "# You may obtain a copy of the License at\n",
        "#\n",
        "# https://www.apache.org/licenses/LICENSE-2.0\n",
        "#\n",
        "# Unless required by applicable law or agreed to in writing, software\n",
        "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
        "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
        "# See the License for the specific language governing permissions and\n",
        "# limitations under the License."
      ],
      "execution_count": 1,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "iOVCm4CnP1Do"
      },
      "source": [
        "\u003ca href=\"https://colab.research.google.com/github/google-research/vision_transformer/blob/master/vit_jax_augreg.ipynb\" target=\"_parent\"\u003e\u003cimg src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/\u003e\u003c/a\u003e"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "8dkWVp5qbziu"
      },
      "source": [
        "## How to train your ViT? Data, Augmentation, and Regularization in Vision Transformers"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ofwi7yvlx5hj"
      },
      "source": [
        "Model repository published with the paper\n",
        "\n",
        "[**How to train your ViT? Data, Augmentation, and Regularization in Vision\n",
        "Transformers**](https://arxiv.org/abs/2106.10270)\n",
        "\n",
        "This Colab shows how to\n",
        "[find checkpoints](#scrollTo=F4SLGDtFxlsC)\n",
        "in the repository, how to\n",
        "[select and load a model](#scrollTo=wh_SLkQtQ6K4)\n",
        "form the repository and use it for inference\n",
        "([also with PyTorch](#scrollTo=1nMyWmDycpAo)),\n",
        "and how to\n",
        "[fine-tune on a dataset](#scrollTo=iAruT3YOxqB6).\n",
        "\n",
        "For more details, please refer to the repository:\n",
        "\n",
        "https://github.com/google-research/vision_transformer/\n",
        "\n",
        "Note that this Colab directly uses the unmodified code from the repository. If\n",
        "you want to modify the modules and persist your changes, you can do all that\n",
        "using free GPUs and TPUs without leaving the Colab environment - see\n",
        "\n",
        "https://colab.research.google.com/github/google-research/vision_transformer/blob/master/vit_jax.ipynb"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "LPNkX_dfCNog"
      },
      "source": [
        "### Imports"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "_ihyJdwbCVmW",
        "outputId": "998b714c-ebaa-4ece-aaf1-533ef08b4d57"
      },
      "source": [
        "# Fetch vision_transformer repository.\n",
        "![ -d vision_transformer ] || git clone --depth=1 https://github.com/google-research/vision_transformer\n"
      ],
      "execution_count": 2,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Cloning into 'vision_transformer'...\n",
            "remote: Enumerating objects: 37, done.\u001b[K\n",
            "remote: Counting objects: 100% (37/37), done.\u001b[K\n",
            "remote: Compressing objects: 100% (34/34), done.\u001b[K\n",
            "remote: Total 37 (delta 4), reused 15 (delta 1), pack-reused 0\u001b[K\n",
            "Unpacking objects: 100% (37/37), done.\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "__7nV7c6C4Zn",
        "outputId": "cd868560-3aae-4887-e735-3a0a9d1ddef2"
      },
      "source": [
        "# Install dependencies.\n",
        "!pip install -qr vision_transformer/vit_jax/requirements.txt"
      ],
      "execution_count": 3,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "\u001b[?25l\r\u001b[K     |█████▋                          | 10 kB 29.7 MB/s eta 0:00:01\r\u001b[K     |███████████▎                    | 20 kB 21.8 MB/s eta 0:00:01\r\u001b[K     |█████████████████               | 30 kB 16.6 MB/s eta 0:00:01\r\u001b[K     |██████████████████████▋         | 40 kB 15.2 MB/s eta 0:00:01\r\u001b[K     |████████████████████████████▎   | 51 kB 7.9 MB/s eta 0:00:01\r\u001b[K     |████████████████████████████████| 57 kB 3.4 MB/s \n",
            "\u001b[K     |████████████████████████████████| 76 kB 5.0 MB/s \n",
            "\u001b[K     |████████████████████████████████| 179 kB 10.1 MB/s \n",
            "\u001b[K     |████████████████████████████████| 88 kB 5.8 MB/s \n",
            "\u001b[K     |████████████████████████████████| 168.3 MB 36 kB/s \n",
            "\u001b[?25h"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "R9VrgMtUC2w-"
      },
      "source": [
        "# Import files from repository.\n",
        "\n",
        "import sys\n",
        "if './vision_transformer' not in sys.path:\n",
        "  sys.path.append('./vision_transformer')\n",
        "\n",
        "%load_ext autoreload\n",
        "%autoreload 2\n",
        "\n",
        "from vit_jax import checkpoint\n",
        "from vit_jax import models\n",
        "from vit_jax import train\n",
        "from vit_jax.configs import augreg as augreg_config\n",
        "from vit_jax.configs import models as models_config"
      ],
      "execution_count": 4,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "MpYDX1idbVxq",
        "outputId": "45a024e9-9b96-4a21-b0de-2d6e4c913f29"
      },
      "source": [
        "# Connect to TPUs if runtime type is of type TPU.\n",
        "\n",
        "import os\n",
        "if 'google.colab' in str(get_ipython()) and 'COLAB_TPU_ADDR' in os.environ:\n",
        "  import jax\n",
        "  import jax.tools.colab_tpu\n",
        "  jax.tools.colab_tpu.setup_tpu()\n",
        "  print('Connected to TPU.')\n",
        "else:\n",
        "  # Otherwise print information about GPU.\n",
        "  !nvidia-smi"
      ],
      "execution_count": 5,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Thu Jul 29 06:07:36 2021       \n",
            "+-----------------------------------------------------------------------------+\n",
            "| NVIDIA-SMI 470.42.01    Driver Version: 460.32.03    CUDA Version: 11.2     |\n",
            "|-------------------------------+----------------------+----------------------+\n",
            "| GPU  Name        Persistence-M| Bus-Id        Disp.A | Volatile Uncorr. ECC |\n",
            "| Fan  Temp  Perf  Pwr:Usage/Cap|         Memory-Usage | GPU-Util  Compute M. |\n",
            "|                               |                      |               MIG M. |\n",
            "|===============================+======================+======================|\n",
            "|   0  Tesla T4            Off  | 00000000:00:04.0 Off |                    0 |\n",
            "| N/A   45C    P8    10W /  70W |      0MiB / 15109MiB |      0%      Default |\n",
            "|                               |                      |                  N/A |\n",
            "+-------------------------------+----------------------+----------------------+\n",
            "                                                                               \n",
            "+-----------------------------------------------------------------------------+\n",
            "| Processes:                                                                  |\n",
            "|  GPU   GI   CI        PID   Type   Process name                  GPU Memory |\n",
            "|        ID   ID                                                   Usage      |\n",
            "|=============================================================================|\n",
            "|  No running processes found                                                 |\n",
            "+-----------------------------------------------------------------------------+\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "6xTeXue8DfwL"
      },
      "source": [
        "# Some more imports used in this Colab.\n",
        "\n",
        "import glob\n",
        "import os\n",
        "import random\n",
        "import shutil\n",
        "import time\n",
        "\n",
        "from absl import logging\n",
        "import pandas as pd\n",
        "import seaborn as sns\n",
        "import tensorflow as tf\n",
        "import tensorflow_datasets as tfds\n",
        "from matplotlib import pyplot as plt\n",
        "\n",
        "pd.options.display.max_colwidth = None\n",
        "logging.set_verbosity(logging.INFO)  # Shows logs during training."
      ],
      "execution_count": 6,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "F4SLGDtFxlsC"
      },
      "source": [
        "### Explore checkpoints\n",
        "\n",
        "This section contains shows how to use the `index.csv` table for model\n",
        "selection.\n",
        "\n",
        "See\n",
        "[`vit_jax.checkpoint.get_augreg_df()`](https://github.com/google-research/vision_transformer/blob/ed1491238f5ff6099cca81087c575a215281ed14/vit_jax/checkpoint.py#L181-L228)\n",
        "for a detailed description of the individual columns"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "yy-cuGxyD6Xw"
      },
      "source": [
        "# Load master table from Cloud.\n",
        "with tf.io.gfile.GFile('gs://vit_models/augreg/index.csv') as f:\n",
        "  df = pd.read_csv(f)"
      ],
      "execution_count": 7,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "j8iAESfDKCQS",
        "outputId": "b92ef316-7df4-464c-9448-25d18dbde999"
      },
      "source": [
        "# This is a pretty large table with lots of columns:\n",
        "print(f'loaded {len(df):,} rows')\n",
        "df.columns"
      ],
      "execution_count": 8,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "loaded 51,509 rows\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "Index(['name', 'ds', 'epochs', 'lr', 'aug', 'wd', 'do', 'sd', 'best_val',\n",
              "       'final_val', 'final_test', 'adapt_ds', 'adapt_lr', 'adapt_steps',\n",
              "       'adapt_resolution', 'adapt_final_val', 'adapt_final_test', 'params',\n",
              "       'infer_samples_per_sec', 'filename', 'adapt_filename'],\n",
              "      dtype='object')"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 8
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "Wi4CwizkKIx_",
        "outputId": "22288fad-82c7-4c10-dcd6-763eb8f39547"
      },
      "source": [
        "# Number of distinct checkpoints\n",
        "len(tf.io.gfile.glob('gs://vit_models/augreg/*.npz'))"
      ],
      "execution_count": 9,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "52268"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 9
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "FOM3HkkbKPdk",
        "outputId": "f54635ac-47f7-44b4-f2c1-6d352b16d12a"
      },
      "source": [
        "# Any column prefixed with \"adapt_\" pertains to the fine-tuned checkpoints.\n",
        "# Any column without that prefix pertains to the pre-trained checkpoints.\n",
        "len(set(df.filename)), len(set(df.adapt_filename))"
      ],
      "execution_count": 10,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "(759, 51509)"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 10
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "h3EyLd3S4XDS",
        "outputId": "6569412c-44da-4f8a-abd8-4c3b1318844b",
        "colab": {
          "base_uri": "https://localhost:8080/"
        }
      },
      "source": [
        "df.name.unique()"
      ],
      "execution_count": 15,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "array(['Ti/16', 'S/32', 'B/16', 'L/16', 'R50+L/32', 'R26+S/32', 'S/16',\n",
              "       'B/32', 'R+Ti/16', 'B/8'], dtype=object)"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 15
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 359
        },
        "id": "w38Gjj-SGLVu",
        "outputId": "67776302-9f11-4cb0-8615-7f193b473f7d"
      },
      "source": [
        "# Upstream AugReg parameters (section 3.3):\n",
        "(\n",
        "df.groupby(['ds', 'name', 'wd', 'do', 'sd', 'aug']).filename\n",
        "  .count().unstack().unstack().unstack()\n",
        "  .dropna(1, 'all').fillna(0).astype(int)\n",
        "  .iloc[:7]  # Just show beginning of a long table.\n",
        ")"
      ],
      "execution_count": 17,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/html": [
              "\u003cdiv\u003e\n",
              "\u003cstyle scoped\u003e\n",
              "    .dataframe tbody tr th:only-of-type {\n",
              "        vertical-align: middle;\n",
              "    }\n",
              "\n",
              "    .dataframe tbody tr th {\n",
              "        vertical-align: top;\n",
              "    }\n",
              "\n",
              "    .dataframe thead tr th {\n",
              "        text-align: left;\n",
              "    }\n",
              "\n",
              "    .dataframe thead tr:last-of-type th {\n",
              "        text-align: right;\n",
              "    }\n",
              "\u003c/style\u003e\n",
              "\u003ctable border=\"1\" class=\"dataframe\"\u003e\n",
              "  \u003cthead\u003e\n",
              "    \u003ctr\u003e\n",
              "      \u003cth\u003e\u003c/th\u003e\n",
              "      \u003cth\u003e\u003c/th\u003e\n",
              "      \u003cth\u003eaug\u003c/th\u003e\n",
              "      \u003cth colspan=\"2\" halign=\"left\"\u003elight0\u003c/th\u003e\n",
              "      \u003cth colspan=\"2\" halign=\"left\"\u003elight1\u003c/th\u003e\n",
              "      \u003cth colspan=\"2\" halign=\"left\"\u003emedium1\u003c/th\u003e\n",
              "      \u003cth colspan=\"2\" halign=\"left\"\u003emedium2\u003c/th\u003e\n",
              "      \u003cth colspan=\"2\" halign=\"left\"\u003enone\u003c/th\u003e\n",
              "      \u003cth colspan=\"2\" halign=\"left\"\u003estrong1\u003c/th\u003e\n",
              "      \u003cth colspan=\"2\" halign=\"left\"\u003estrong2\u003c/th\u003e\n",
              "    \u003c/tr\u003e\n",
              "    \u003ctr\u003e\n",
              "      \u003cth\u003e\u003c/th\u003e\n",
              "      \u003cth\u003e\u003c/th\u003e\n",
              "      \u003cth\u003esd\u003c/th\u003e\n",
              "      \u003cth\u003e0.0\u003c/th\u003e\n",
              "      \u003cth\u003e0.1\u003c/th\u003e\n",
              "      \u003cth\u003e0.0\u003c/th\u003e\n",
              "      \u003cth\u003e0.1\u003c/th\u003e\n",
              "      \u003cth\u003e0.0\u003c/th\u003e\n",
              "      \u003cth\u003e0.1\u003c/th\u003e\n",
              "      \u003cth\u003e0.0\u003c/th\u003e\n",
              "      \u003cth\u003e0.1\u003c/th\u003e\n",
              "      \u003cth\u003e0.0\u003c/th\u003e\n",
              "      \u003cth\u003e0.1\u003c/th\u003e\n",
              "      \u003cth\u003e0.0\u003c/th\u003e\n",
              "      \u003cth\u003e0.1\u003c/th\u003e\n",
              "      \u003cth\u003e0.0\u003c/th\u003e\n",
              "      \u003cth\u003e0.1\u003c/th\u003e\n",
              "    \u003c/tr\u003e\n",
              "    \u003ctr\u003e\n",
              "      \u003cth\u003e\u003c/th\u003e\n",
              "      \u003cth\u003e\u003c/th\u003e\n",
              "      \u003cth\u003edo\u003c/th\u003e\n",
              "      \u003cth\u003e0.0\u003c/th\u003e\n",
              "      \u003cth\u003e0.1\u003c/th\u003e\n",
              "      \u003cth\u003e0.0\u003c/th\u003e\n",
              "      \u003cth\u003e0.1\u003c/th\u003e\n",
              "      \u003cth\u003e0.0\u003c/th\u003e\n",
              "      \u003cth\u003e0.1\u003c/th\u003e\n",
              "      \u003cth\u003e0.0\u003c/th\u003e\n",
              "      \u003cth\u003e0.1\u003c/th\u003e\n",
              "      \u003cth\u003e0.0\u003c/th\u003e\n",
              "      \u003cth\u003e0.1\u003c/th\u003e\n",
              "      \u003cth\u003e0.0\u003c/th\u003e\n",
              "      \u003cth\u003e0.1\u003c/th\u003e\n",
              "      \u003cth\u003e0.0\u003c/th\u003e\n",
              "      \u003cth\u003e0.1\u003c/th\u003e\n",
              "    \u003c/tr\u003e\n",
              "    \u003ctr\u003e\n",
              "      \u003cth\u003eds\u003c/th\u003e\n",
              "      \u003cth\u003ename\u003c/th\u003e\n",
              "      \u003cth\u003ewd\u003c/th\u003e\n",
              "      \u003cth\u003e\u003c/th\u003e\n",
              "      \u003cth\u003e\u003c/th\u003e\n",
              "      \u003cth\u003e\u003c/th\u003e\n",
              "      \u003cth\u003e\u003c/th\u003e\n",
              "      \u003cth\u003e\u003c/th\u003e\n",
              "      \u003cth\u003e\u003c/th\u003e\n",
              "      \u003cth\u003e\u003c/th\u003e\n",
              "      \u003cth\u003e\u003c/th\u003e\n",
              "      \u003cth\u003e\u003c/th\u003e\n",
              "      \u003cth\u003e\u003c/th\u003e\n",
              "      \u003cth\u003e\u003c/th\u003e\n",
              "      \u003cth\u003e\u003c/th\u003e\n",
              "      \u003cth\u003e\u003c/th\u003e\n",
              "      \u003cth\u003e\u003c/th\u003e\n",
              "    \u003c/tr\u003e\n",
              "  \u003c/thead\u003e\n",
              "  \u003ctbody\u003e\n",
              "    \u003ctr\u003e\n",
              "      \u003cth rowspan=\"7\" valign=\"top\"\u003ei1k\u003c/th\u003e\n",
              "      \u003cth rowspan=\"2\" valign=\"top\"\u003eB/16\u003c/th\u003e\n",
              "      \u003cth\u003e0.03\u003c/th\u003e\n",
              "      \u003ctd\u003e68\u003c/td\u003e\n",
              "      \u003ctd\u003e68\u003c/td\u003e\n",
              "      \u003ctd\u003e68\u003c/td\u003e\n",
              "      \u003ctd\u003e68\u003c/td\u003e\n",
              "      \u003ctd\u003e68\u003c/td\u003e\n",
              "      \u003ctd\u003e68\u003c/td\u003e\n",
              "      \u003ctd\u003e68\u003c/td\u003e\n",
              "      \u003ctd\u003e68\u003c/td\u003e\n",
              "      \u003ctd\u003e68\u003c/td\u003e\n",
              "      \u003ctd\u003e68\u003c/td\u003e\n",
              "      \u003ctd\u003e68\u003c/td\u003e\n",
              "      \u003ctd\u003e68\u003c/td\u003e\n",
              "      \u003ctd\u003e68\u003c/td\u003e\n",
              "      \u003ctd\u003e68\u003c/td\u003e\n",
              "    \u003c/tr\u003e\n",
              "    \u003ctr\u003e\n",
              "      \u003cth\u003e0.10\u003c/th\u003e\n",
              "      \u003ctd\u003e68\u003c/td\u003e\n",
              "      \u003ctd\u003e68\u003c/td\u003e\n",
              "      \u003ctd\u003e68\u003c/td\u003e\n",
              "      \u003ctd\u003e68\u003c/td\u003e\n",
              "      \u003ctd\u003e68\u003c/td\u003e\n",
              "      \u003ctd\u003e68\u003c/td\u003e\n",
              "      \u003ctd\u003e68\u003c/td\u003e\n",
              "      \u003ctd\u003e68\u003c/td\u003e\n",
              "      \u003ctd\u003e68\u003c/td\u003e\n",
              "      \u003ctd\u003e68\u003c/td\u003e\n",
              "      \u003ctd\u003e68\u003c/td\u003e\n",
              "      \u003ctd\u003e68\u003c/td\u003e\n",
              "      \u003ctd\u003e68\u003c/td\u003e\n",
              "      \u003ctd\u003e68\u003c/td\u003e\n",
              "    \u003c/tr\u003e\n",
              "    \u003ctr\u003e\n",
              "      \u003cth rowspan=\"2\" valign=\"top\"\u003eB/32\u003c/th\u003e\n",
              "      \u003cth\u003e0.03\u003c/th\u003e\n",
              "      \u003ctd\u003e68\u003c/td\u003e\n",
              "      \u003ctd\u003e68\u003c/td\u003e\n",
              "      \u003ctd\u003e68\u003c/td\u003e\n",
              "      \u003ctd\u003e68\u003c/td\u003e\n",
              "      \u003ctd\u003e68\u003c/td\u003e\n",
              "      \u003ctd\u003e68\u003c/td\u003e\n",
              "      \u003ctd\u003e68\u003c/td\u003e\n",
              "      \u003ctd\u003e68\u003c/td\u003e\n",
              "      \u003ctd\u003e68\u003c/td\u003e\n",
              "      \u003ctd\u003e68\u003c/td\u003e\n",
              "      \u003ctd\u003e68\u003c/td\u003e\n",
              "      \u003ctd\u003e68\u003c/td\u003e\n",
              "      \u003ctd\u003e68\u003c/td\u003e\n",
              "      \u003ctd\u003e68\u003c/td\u003e\n",
              "    \u003c/tr\u003e\n",
              "    \u003ctr\u003e\n",
              "      \u003cth\u003e0.10\u003c/th\u003e\n",
              "      \u003ctd\u003e68\u003c/td\u003e\n",
              "      \u003ctd\u003e68\u003c/td\u003e\n",
              "      \u003ctd\u003e68\u003c/td\u003e\n",
              "      \u003ctd\u003e68\u003c/td\u003e\n",
              "      \u003ctd\u003e68\u003c/td\u003e\n",
              "      \u003ctd\u003e68\u003c/td\u003e\n",
              "      \u003ctd\u003e68\u003c/td\u003e\n",
              "      \u003ctd\u003e68\u003c/td\u003e\n",
              "      \u003ctd\u003e68\u003c/td\u003e\n",
              "      \u003ctd\u003e68\u003c/td\u003e\n",
              "      \u003ctd\u003e68\u003c/td\u003e\n",
              "      \u003ctd\u003e68\u003c/td\u003e\n",
              "      \u003ctd\u003e68\u003c/td\u003e\n",
              "      \u003ctd\u003e68\u003c/td\u003e\n",
              "    \u003c/tr\u003e\n",
              "    \u003ctr\u003e\n",
              "      \u003cth rowspan=\"2\" valign=\"top\"\u003eL/16\u003c/th\u003e\n",
              "      \u003cth\u003e0.03\u003c/th\u003e\n",
              "      \u003ctd\u003e68\u003c/td\u003e\n",
              "      \u003ctd\u003e68\u003c/td\u003e\n",
              "      \u003ctd\u003e68\u003c/td\u003e\n",
              "      \u003ctd\u003e68\u003c/td\u003e\n",
              "      \u003ctd\u003e68\u003c/td\u003e\n",
              "      \u003ctd\u003e68\u003c/td\u003e\n",
              "      \u003ctd\u003e68\u003c/td\u003e\n",
              "      \u003ctd\u003e68\u003c/td\u003e\n",
              "      \u003ctd\u003e68\u003c/td\u003e\n",
              "      \u003ctd\u003e68\u003c/td\u003e\n",
              "      \u003ctd\u003e68\u003c/td\u003e\n",
              "      \u003ctd\u003e68\u003c/td\u003e\n",
              "      \u003ctd\u003e68\u003c/td\u003e\n",
              "      \u003ctd\u003e68\u003c/td\u003e\n",
              "    \u003c/tr\u003e\n",
              "    \u003ctr\u003e\n",
              "      \u003cth\u003e0.10\u003c/th\u003e\n",
              "      \u003ctd\u003e68\u003c/td\u003e\n",
              "      \u003ctd\u003e68\u003c/td\u003e\n",
              "      \u003ctd\u003e68\u003c/td\u003e\n",
              "      \u003ctd\u003e68\u003c/td\u003e\n",
              "      \u003ctd\u003e68\u003c/td\u003e\n",
              "      \u003ctd\u003e68\u003c/td\u003e\n",
              "      \u003ctd\u003e68\u003c/td\u003e\n",
              "      \u003ctd\u003e68\u003c/td\u003e\n",
              "      \u003ctd\u003e68\u003c/td\u003e\n",
              "      \u003ctd\u003e68\u003c/td\u003e\n",
              "      \u003ctd\u003e68\u003c/td\u003e\n",
              "      \u003ctd\u003e68\u003c/td\u003e\n",
              "      \u003ctd\u003e68\u003c/td\u003e\n",
              "      \u003ctd\u003e68\u003c/td\u003e\n",
              "    \u003c/tr\u003e\n",
              "    \u003ctr\u003e\n",
              "      \u003cth\u003eR+Ti/16\u003c/th\u003e\n",
              "      \u003cth\u003e0.03\u003c/th\u003e\n",
              "      \u003ctd\u003e68\u003c/td\u003e\n",
              "      \u003ctd\u003e68\u003c/td\u003e\n",
              "      \u003ctd\u003e68\u003c/td\u003e\n",
              "      \u003ctd\u003e68\u003c/td\u003e\n",
              "      \u003ctd\u003e68\u003c/td\u003e\n",
              "      \u003ctd\u003e68\u003c/td\u003e\n",
              "      \u003ctd\u003e68\u003c/td\u003e\n",
              "      \u003ctd\u003e68\u003c/td\u003e\n",
              "      \u003ctd\u003e68\u003c/td\u003e\n",
              "      \u003ctd\u003e68\u003c/td\u003e\n",
              "      \u003ctd\u003e68\u003c/td\u003e\n",
              "      \u003ctd\u003e68\u003c/td\u003e\n",
              "      \u003ctd\u003e68\u003c/td\u003e\n",
              "      \u003ctd\u003e68\u003c/td\u003e\n",
              "    \u003c/tr\u003e\n",
              "  \u003c/tbody\u003e\n",
              "\u003c/table\u003e\n",
              "\u003c/div\u003e"
            ],
            "text/plain": [
              "aug              light0     light1     medium1  ... none strong1     strong2    \n",
              "sd                  0.0 0.1    0.0 0.1     0.0  ...  0.1     0.0 0.1     0.0 0.1\n",
              "do                  0.0 0.1    0.0 0.1     0.0  ...  0.1     0.0 0.1     0.0 0.1\n",
              "ds  name    wd                                  ...                             \n",
              "i1k B/16    0.03     68  68     68  68      68  ...   68      68  68      68  68\n",
              "            0.10     68  68     68  68      68  ...   68      68  68      68  68\n",
              "    B/32    0.03     68  68     68  68      68  ...   68      68  68      68  68\n",
              "            0.10     68  68     68  68      68  ...   68      68  68      68  68\n",
              "    L/16    0.03     68  68     68  68      68  ...   68      68  68      68  68\n",
              "            0.10     68  68     68  68      68  ...   68      68  68      68  68\n",
              "    R+Ti/16 0.03     68  68     68  68      68  ...   68      68  68      68  68\n",
              "\n",
              "[7 rows x 14 columns]"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 17
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 421
        },
        "id": "PjI5jX36GeFr",
        "outputId": "27c96ac3-d61d-4974-f614-cc86456a673b"
      },
      "source": [
        "# Downstream parameters (table 4)\n",
        "# (Imbalance in 224 vs. 384 is due to recently added B/8 checkpoints)\n",
        "(\n",
        "df.groupby(['adapt_resolution', 'adapt_ds', 'adapt_lr', 'adapt_steps']).filename\n",
        "  .count().astype(str).unstack().unstack()\n",
        "  .dropna(1, 'all').fillna('')\n",
        ")"
      ],
      "execution_count": 18,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/html": [
              "\u003cdiv\u003e\n",
              "\u003cstyle scoped\u003e\n",
              "    .dataframe tbody tr th:only-of-type {\n",
              "        vertical-align: middle;\n",
              "    }\n",
              "\n",
              "    .dataframe tbody tr th {\n",
              "        vertical-align: top;\n",
              "    }\n",
              "\n",
              "    .dataframe thead tr th {\n",
              "        text-align: left;\n",
              "    }\n",
              "\n",
              "    .dataframe thead tr:last-of-type th {\n",
              "        text-align: right;\n",
              "    }\n",
              "\u003c/style\u003e\n",
              "\u003ctable border=\"1\" class=\"dataframe\"\u003e\n",
              "  \u003cthead\u003e\n",
              "    \u003ctr\u003e\n",
              "      \u003cth\u003e\u003c/th\u003e\n",
              "      \u003cth\u003eadapt_steps\u003c/th\u003e\n",
              "      \u003cth colspan=\"4\" halign=\"left\"\u003e500\u003c/th\u003e\n",
              "      \u003cth colspan=\"4\" halign=\"left\"\u003e2500\u003c/th\u003e\n",
              "      \u003cth colspan=\"4\" halign=\"left\"\u003e10000\u003c/th\u003e\n",
              "      \u003cth colspan=\"2\" halign=\"left\"\u003e20000\u003c/th\u003e\n",
              "    \u003c/tr\u003e\n",
              "    \u003ctr\u003e\n",
              "      \u003cth\u003e\u003c/th\u003e\n",
              "      \u003cth\u003eadapt_lr\u003c/th\u003e\n",
              "      \u003cth\u003e0.001\u003c/th\u003e\n",
              "      \u003cth\u003e0.003\u003c/th\u003e\n",
              "      \u003cth\u003e0.010\u003c/th\u003e\n",
              "      \u003cth\u003e0.030\u003c/th\u003e\n",
              "      \u003cth\u003e0.001\u003c/th\u003e\n",
              "      \u003cth\u003e0.003\u003c/th\u003e\n",
              "      \u003cth\u003e0.010\u003c/th\u003e\n",
              "      \u003cth\u003e0.030\u003c/th\u003e\n",
              "      \u003cth\u003e0.001\u003c/th\u003e\n",
              "      \u003cth\u003e0.003\u003c/th\u003e\n",
              "      \u003cth\u003e0.010\u003c/th\u003e\n",
              "      \u003cth\u003e0.030\u003c/th\u003e\n",
              "      \u003cth\u003e0.010\u003c/th\u003e\n",
              "      \u003cth\u003e0.030\u003c/th\u003e\n",
              "    \u003c/tr\u003e\n",
              "    \u003ctr\u003e\n",
              "      \u003cth\u003eadapt_resolution\u003c/th\u003e\n",
              "      \u003cth\u003eadapt_ds\u003c/th\u003e\n",
              "      \u003cth\u003e\u003c/th\u003e\n",
              "      \u003cth\u003e\u003c/th\u003e\n",
              "      \u003cth\u003e\u003c/th\u003e\n",
              "      \u003cth\u003e\u003c/th\u003e\n",
              "      \u003cth\u003e\u003c/th\u003e\n",
              "      \u003cth\u003e\u003c/th\u003e\n",
              "      \u003cth\u003e\u003c/th\u003e\n",
              "      \u003cth\u003e\u003c/th\u003e\n",
              "      \u003cth\u003e\u003c/th\u003e\n",
              "      \u003cth\u003e\u003c/th\u003e\n",
              "      \u003cth\u003e\u003c/th\u003e\n",
              "      \u003cth\u003e\u003c/th\u003e\n",
              "      \u003cth\u003e\u003c/th\u003e\n",
              "      \u003cth\u003e\u003c/th\u003e\n",
              "    \u003c/tr\u003e\n",
              "  \u003c/thead\u003e\n",
              "  \u003ctbody\u003e\n",
              "    \u003ctr\u003e\n",
              "      \u003cth rowspan=\"5\" valign=\"top\"\u003e224\u003c/th\u003e\n",
              "      \u003cth\u003ecifar100\u003c/th\u003e\n",
              "      \u003ctd\u003e\u003c/td\u003e\n",
              "      \u003ctd\u003e\u003c/td\u003e\n",
              "      \u003ctd\u003e\u003c/td\u003e\n",
              "      \u003ctd\u003e\u003c/td\u003e\n",
              "      \u003ctd\u003e759\u003c/td\u003e\n",
              "      \u003ctd\u003e759\u003c/td\u003e\n",
              "      \u003ctd\u003e759\u003c/td\u003e\n",
              "      \u003ctd\u003e759\u003c/td\u003e\n",
              "      \u003ctd\u003e759\u003c/td\u003e\n",
              "      \u003ctd\u003e759\u003c/td\u003e\n",
              "      \u003ctd\u003e759\u003c/td\u003e\n",
              "      \u003ctd\u003e759\u003c/td\u003e\n",
              "      \u003ctd\u003e\u003c/td\u003e\n",
              "      \u003ctd\u003e\u003c/td\u003e\n",
              "    \u003c/tr\u003e\n",
              "    \u003ctr\u003e\n",
              "      \u003cth\u003eimagenet2012\u003c/th\u003e\n",
              "      \u003ctd\u003e\u003c/td\u003e\n",
              "      \u003ctd\u003e\u003c/td\u003e\n",
              "      \u003ctd\u003e\u003c/td\u003e\n",
              "      \u003ctd\u003e\u003c/td\u003e\n",
              "      \u003ctd\u003e\u003c/td\u003e\n",
              "      \u003ctd\u003e\u003c/td\u003e\n",
              "      \u003ctd\u003e\u003c/td\u003e\n",
              "      \u003ctd\u003e\u003c/td\u003e\n",
              "      \u003ctd\u003e\u003c/td\u003e\n",
              "      \u003ctd\u003e\u003c/td\u003e\n",
              "      \u003ctd\u003e\u003c/td\u003e\n",
              "      \u003ctd\u003e\u003c/td\u003e\n",
              "      \u003ctd\u003e759\u003c/td\u003e\n",
              "      \u003ctd\u003e759\u003c/td\u003e\n",
              "    \u003c/tr\u003e\n",
              "    \u003ctr\u003e\n",
              "      \u003cth\u003ekitti\u003c/th\u003e\n",
              "      \u003ctd\u003e759\u003c/td\u003e\n",
              "      \u003ctd\u003e759\u003c/td\u003e\n",
              "      \u003ctd\u003e759\u003c/td\u003e\n",
              "      \u003ctd\u003e759\u003c/td\u003e\n",
              "      \u003ctd\u003e759\u003c/td\u003e\n",
              "      \u003ctd\u003e759\u003c/td\u003e\n",
              "      \u003ctd\u003e759\u003c/td\u003e\n",
              "      \u003ctd\u003e759\u003c/td\u003e\n",
              "      \u003ctd\u003e\u003c/td\u003e\n",
              "      \u003ctd\u003e\u003c/td\u003e\n",
              "      \u003ctd\u003e\u003c/td\u003e\n",
              "      \u003ctd\u003e\u003c/td\u003e\n",
              "      \u003ctd\u003e\u003c/td\u003e\n",
              "      \u003ctd\u003e\u003c/td\u003e\n",
              "    \u003c/tr\u003e\n",
              "    \u003ctr\u003e\n",
              "      \u003cth\u003eoxford_iiit_pet\u003c/th\u003e\n",
              "      \u003ctd\u003e759\u003c/td\u003e\n",
              "      \u003ctd\u003e759\u003c/td\u003e\n",
              "      \u003ctd\u003e759\u003c/td\u003e\n",
              "      \u003ctd\u003e759\u003c/td\u003e\n",
              "      \u003ctd\u003e759\u003c/td\u003e\n",
              "      \u003ctd\u003e759\u003c/td\u003e\n",
              "      \u003ctd\u003e759\u003c/td\u003e\n",
              "      \u003ctd\u003e759\u003c/td\u003e\n",
              "      \u003ctd\u003e\u003c/td\u003e\n",
              "      \u003ctd\u003e\u003c/td\u003e\n",
              "      \u003ctd\u003e\u003c/td\u003e\n",
              "      \u003ctd\u003e\u003c/td\u003e\n",
              "      \u003ctd\u003e\u003c/td\u003e\n",
              "      \u003ctd\u003e\u003c/td\u003e\n",
              "    \u003c/tr\u003e\n",
              "    \u003ctr\u003e\n",
              "      \u003cth\u003eresisc45\u003c/th\u003e\n",
              "      \u003ctd\u003e\u003c/td\u003e\n",
              "      \u003ctd\u003e\u003c/td\u003e\n",
              "      \u003ctd\u003e\u003c/td\u003e\n",
              "      \u003ctd\u003e\u003c/td\u003e\n",
              "      \u003ctd\u003e759\u003c/td\u003e\n",
              "      \u003ctd\u003e759\u003c/td\u003e\n",
              "      \u003ctd\u003e759\u003c/td\u003e\n",
              "      \u003ctd\u003e759\u003c/td\u003e\n",
              "      \u003ctd\u003e759\u003c/td\u003e\n",
              "      \u003ctd\u003e759\u003c/td\u003e\n",
              "      \u003ctd\u003e759\u003c/td\u003e\n",
              "      \u003ctd\u003e759\u003c/td\u003e\n",
              "      \u003ctd\u003e\u003c/td\u003e\n",
              "      \u003ctd\u003e\u003c/td\u003e\n",
              "    \u003c/tr\u003e\n",
              "    \u003ctr\u003e\n",
              "      \u003cth rowspan=\"5\" valign=\"top\"\u003e384\u003c/th\u003e\n",
              "      \u003cth\u003ecifar100\u003c/th\u003e\n",
              "      \u003ctd\u003e\u003c/td\u003e\n",
              "      \u003ctd\u003e\u003c/td\u003e\n",
              "      \u003ctd\u003e\u003c/td\u003e\n",
              "      \u003ctd\u003e\u003c/td\u003e\n",
              "      \u003ctd\u003e756\u003c/td\u003e\n",
              "      \u003ctd\u003e756\u003c/td\u003e\n",
              "      \u003ctd\u003e756\u003c/td\u003e\n",
              "      \u003ctd\u003e756\u003c/td\u003e\n",
              "      \u003ctd\u003e756\u003c/td\u003e\n",
              "      \u003ctd\u003e756\u003c/td\u003e\n",
              "      \u003ctd\u003e756\u003c/td\u003e\n",
              "      \u003ctd\u003e756\u003c/td\u003e\n",
              "      \u003ctd\u003e\u003c/td\u003e\n",
              "      \u003ctd\u003e\u003c/td\u003e\n",
              "    \u003c/tr\u003e\n",
              "    \u003ctr\u003e\n",
              "      \u003cth\u003eimagenet2012\u003c/th\u003e\n",
              "      \u003ctd\u003e\u003c/td\u003e\n",
              "      \u003ctd\u003e\u003c/td\u003e\n",
              "      \u003ctd\u003e\u003c/td\u003e\n",
              "      \u003ctd\u003e\u003c/td\u003e\n",
              "      \u003ctd\u003e\u003c/td\u003e\n",
              "      \u003ctd\u003e\u003c/td\u003e\n",
              "      \u003ctd\u003e\u003c/td\u003e\n",
              "      \u003ctd\u003e\u003c/td\u003e\n",
              "      \u003ctd\u003e\u003c/td\u003e\n",
              "      \u003ctd\u003e\u003c/td\u003e\n",
              "      \u003ctd\u003e\u003c/td\u003e\n",
              "      \u003ctd\u003e\u003c/td\u003e\n",
              "      \u003ctd\u003e756\u003c/td\u003e\n",
              "      \u003ctd\u003e755\u003c/td\u003e\n",
              "    \u003c/tr\u003e\n",
              "    \u003ctr\u003e\n",
              "      \u003cth\u003ekitti\u003c/th\u003e\n",
              "      \u003ctd\u003e756\u003c/td\u003e\n",
              "      \u003ctd\u003e756\u003c/td\u003e\n",
              "      \u003ctd\u003e756\u003c/td\u003e\n",
              "      \u003ctd\u003e756\u003c/td\u003e\n",
              "      \u003ctd\u003e756\u003c/td\u003e\n",
              "      \u003ctd\u003e756\u003c/td\u003e\n",
              "      \u003ctd\u003e756\u003c/td\u003e\n",
              "      \u003ctd\u003e756\u003c/td\u003e\n",
              "      \u003ctd\u003e\u003c/td\u003e\n",
              "      \u003ctd\u003e\u003c/td\u003e\n",
              "      \u003ctd\u003e\u003c/td\u003e\n",
              "      \u003ctd\u003e\u003c/td\u003e\n",
              "      \u003ctd\u003e\u003c/td\u003e\n",
              "      \u003ctd\u003e\u003c/td\u003e\n",
              "    \u003c/tr\u003e\n",
              "    \u003ctr\u003e\n",
              "      \u003cth\u003eoxford_iiit_pet\u003c/th\u003e\n",
              "      \u003ctd\u003e756\u003c/td\u003e\n",
              "      \u003ctd\u003e756\u003c/td\u003e\n",
              "      \u003ctd\u003e756\u003c/td\u003e\n",
              "      \u003ctd\u003e756\u003c/td\u003e\n",
              "      \u003ctd\u003e756\u003c/td\u003e\n",
              "      \u003ctd\u003e756\u003c/td\u003e\n",
              "      \u003ctd\u003e756\u003c/td\u003e\n",
              "      \u003ctd\u003e756\u003c/td\u003e\n",
              "      \u003ctd\u003e\u003c/td\u003e\n",
              "      \u003ctd\u003e\u003c/td\u003e\n",
              "      \u003ctd\u003e\u003c/td\u003e\n",
              "      \u003ctd\u003e\u003c/td\u003e\n",
              "      \u003ctd\u003e\u003c/td\u003e\n",
              "      \u003ctd\u003e\u003c/td\u003e\n",
              "    \u003c/tr\u003e\n",
              "    \u003ctr\u003e\n",
              "      \u003cth\u003eresisc45\u003c/th\u003e\n",
              "      \u003ctd\u003e\u003c/td\u003e\n",
              "      \u003ctd\u003e\u003c/td\u003e\n",
              "      \u003ctd\u003e\u003c/td\u003e\n",
              "      \u003ctd\u003e\u003c/td\u003e\n",
              "      \u003ctd\u003e756\u003c/td\u003e\n",
              "      \u003ctd\u003e756\u003c/td\u003e\n",
              "      \u003ctd\u003e756\u003c/td\u003e\n",
              "      \u003ctd\u003e756\u003c/td\u003e\n",
              "      \u003ctd\u003e756\u003c/td\u003e\n",
              "      \u003ctd\u003e756\u003c/td\u003e\n",
              "      \u003ctd\u003e756\u003c/td\u003e\n",
              "      \u003ctd\u003e756\u003c/td\u003e\n",
              "      \u003ctd\u003e\u003c/td\u003e\n",
              "      \u003ctd\u003e\u003c/td\u003e\n",
              "    \u003c/tr\u003e\n",
              "  \u003c/tbody\u003e\n",
              "\u003c/table\u003e\n",
              "\u003c/div\u003e"
            ],
            "text/plain": [
              "adapt_steps                      500                ... 10000 20000      \n",
              "adapt_lr                         0.001 0.003 0.010  ... 0.030 0.010 0.030\n",
              "adapt_resolution adapt_ds                           ...                  \n",
              "224              cifar100                           ...   759            \n",
              "                 imagenet2012                       ...         759   759\n",
              "                 kitti             759   759   759  ...                  \n",
              "                 oxford_iiit_pet   759   759   759  ...                  \n",
              "                 resisc45                           ...   759            \n",
              "384              cifar100                           ...   756            \n",
              "                 imagenet2012                       ...         756   755\n",
              "                 kitti             756   756   756  ...                  \n",
              "                 oxford_iiit_pet   756   756   756  ...                  \n",
              "                 resisc45                           ...   756            \n",
              "\n",
              "[10 rows x 14 columns]"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 18
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "8Gbt7cdhEHg7",
        "outputId": "1d07cd9e-3b5f-4cdb-c38b-15bb3755843a"
      },
      "source": [
        "# Let's first select the \"best checkpoint\" for every model. We show in the\n",
        "# paper (section 4.5) that one can get a good performance by simply choosing the\n",
        "# best model by final pre-train validation accuracy (\"final-val\" column).\n",
        "# Pre-training with imagenet21k 300 epochs (ds==\"i21k\") gives the best\n",
        "# performance in almost all cases (figure 6, table 5).\n",
        "best_filenames = set(\n",
        "    df.query('ds==\"i21k\"')\n",
        "    .groupby('name')\n",
        "    .apply(lambda df: df.sort_values('final_val').iloc[-1])\n",
        "    .filename\n",
        ")\n",
        "\n",
        "# Select all finetunes from these models.\n",
        "best_df = df.loc[df.filename.apply(lambda filename: filename in best_filenames)]\n",
        "\n",
        "# Note: 9 * 68 == 612\n",
        "len(best_filenames), len(best_df)"
      ],
      "execution_count": 19,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "(10, 646)"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 19
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "poiez-EcZmtI",
        "outputId": "c6ab958a-3bc8-4d56-df11-fdc6a5f0ea83"
      },
      "source": [
        "best_df.columns"
      ],
      "execution_count": 20,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "Index(['name', 'ds', 'epochs', 'lr', 'aug', 'wd', 'do', 'sd', 'best_val',\n",
              "       'final_val', 'final_test', 'adapt_ds', 'adapt_lr', 'adapt_steps',\n",
              "       'adapt_resolution', 'adapt_final_val', 'adapt_final_test', 'params',\n",
              "       'infer_samples_per_sec', 'filename', 'adapt_filename'],\n",
              "      dtype='object')"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 20
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 730
        },
        "id": "MI_XTt1vX930",
        "outputId": "70747a26-5c5a-47b5-c269-12ef28848301"
      },
      "source": [
        "# Note that this dataframe contains the models from the \"i21k_300\" column of\n",
        "# table 3:\n",
        "best_df.query('adapt_ds==\"imagenet2012\"').groupby('name').apply(\n",
        "    lambda df: df.sort_values('adapt_final_val').iloc[-1]\n",
        ")[[\n",
        "   # Columns from upstream\n",
        "   'name', 'ds', 'filename',\n",
        "   # Columns from downstream\n",
        "   'adapt_resolution', 'infer_samples_per_sec','adapt_ds', 'adapt_final_test', 'adapt_filename',\n",
        "]].sort_values('infer_samples_per_sec')"
      ],
      "execution_count": 21,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/html": [
              "\u003cdiv\u003e\n",
              "\u003cstyle scoped\u003e\n",
              "    .dataframe tbody tr th:only-of-type {\n",
              "        vertical-align: middle;\n",
              "    }\n",
              "\n",
              "    .dataframe tbody tr th {\n",
              "        vertical-align: top;\n",
              "    }\n",
              "\n",
              "    .dataframe thead th {\n",
              "        text-align: right;\n",
              "    }\n",
              "\u003c/style\u003e\n",
              "\u003ctable border=\"1\" class=\"dataframe\"\u003e\n",
              "  \u003cthead\u003e\n",
              "    \u003ctr style=\"text-align: right;\"\u003e\n",
              "      \u003cth\u003e\u003c/th\u003e\n",
              "      \u003cth\u003ename\u003c/th\u003e\n",
              "      \u003cth\u003eds\u003c/th\u003e\n",
              "      \u003cth\u003efilename\u003c/th\u003e\n",
              "      \u003cth\u003eadapt_resolution\u003c/th\u003e\n",
              "      \u003cth\u003einfer_samples_per_sec\u003c/th\u003e\n",
              "      \u003cth\u003eadapt_ds\u003c/th\u003e\n",
              "      \u003cth\u003eadapt_final_test\u003c/th\u003e\n",
              "      \u003cth\u003eadapt_filename\u003c/th\u003e\n",
              "    \u003c/tr\u003e\n",
              "    \u003ctr\u003e\n",
              "      \u003cth\u003ename\u003c/th\u003e\n",
              "      \u003cth\u003e\u003c/th\u003e\n",
              "      \u003cth\u003e\u003c/th\u003e\n",
              "      \u003cth\u003e\u003c/th\u003e\n",
              "      \u003cth\u003e\u003c/th\u003e\n",
              "      \u003cth\u003e\u003c/th\u003e\n",
              "      \u003cth\u003e\u003c/th\u003e\n",
              "      \u003cth\u003e\u003c/th\u003e\n",
              "      \u003cth\u003e\u003c/th\u003e\n",
              "    \u003c/tr\u003e\n",
              "  \u003c/thead\u003e\n",
              "  \u003ctbody\u003e\n",
              "    \u003ctr\u003e\n",
              "      \u003cth\u003eL/16\u003c/th\u003e\n",
              "      \u003ctd\u003eL/16\u003c/td\u003e\n",
              "      \u003ctd\u003ei21k\u003c/td\u003e\n",
              "      \u003ctd\u003eL_16-i21k-300ep-lr_0.001-aug_strong1-wd_0.1-do_0.0-sd_0.0\u003c/td\u003e\n",
              "      \u003ctd\u003e384\u003c/td\u003e\n",
              "      \u003ctd\u003e49.87\u003c/td\u003e\n",
              "      \u003ctd\u003eimagenet2012\u003c/td\u003e\n",
              "      \u003ctd\u003e0.85588\u003c/td\u003e\n",
              "      \u003ctd\u003eL_16-i21k-300ep-lr_0.001-aug_strong1-wd_0.1-do_0.0-sd_0.0--imagenet2012-steps_20k-lr_0.01-res_384\u003c/td\u003e\n",
              "    \u003c/tr\u003e\n",
              "    \u003ctr\u003e\n",
              "      \u003cth\u003eB/16\u003c/th\u003e\n",
              "      \u003ctd\u003eB/16\u003c/td\u003e\n",
              "      \u003ctd\u003ei21k\u003c/td\u003e\n",
              "      \u003ctd\u003eB_16-i21k-300ep-lr_0.001-aug_medium1-wd_0.1-do_0.0-sd_0.0\u003c/td\u003e\n",
              "      \u003ctd\u003e384\u003c/td\u003e\n",
              "      \u003ctd\u003e137.92\u003c/td\u003e\n",
              "      \u003ctd\u003eimagenet2012\u003c/td\u003e\n",
              "      \u003ctd\u003e0.85486\u003c/td\u003e\n",
              "      \u003ctd\u003eB_16-i21k-300ep-lr_0.001-aug_medium1-wd_0.1-do_0.0-sd_0.0--imagenet2012-steps_20k-lr_0.03-res_384\u003c/td\u003e\n",
              "    \u003c/tr\u003e\n",
              "    \u003ctr\u003e\n",
              "      \u003cth\u003eS/16\u003c/th\u003e\n",
              "      \u003ctd\u003eS/16\u003c/td\u003e\n",
              "      \u003ctd\u003ei21k\u003c/td\u003e\n",
              "      \u003ctd\u003eS_16-i21k-300ep-lr_0.001-aug_light1-wd_0.03-do_0.0-sd_0.0\u003c/td\u003e\n",
              "      \u003ctd\u003e384\u003c/td\u003e\n",
              "      \u003ctd\u003e300.12\u003c/td\u003e\n",
              "      \u003ctd\u003eimagenet2012\u003c/td\u003e\n",
              "      \u003ctd\u003e0.83728\u003c/td\u003e\n",
              "      \u003ctd\u003eS_16-i21k-300ep-lr_0.001-aug_light1-wd_0.03-do_0.0-sd_0.0--imagenet2012-steps_20k-lr_0.03-res_384\u003c/td\u003e\n",
              "    \u003c/tr\u003e\n",
              "    \u003ctr\u003e\n",
              "      \u003cth\u003eR50+L/32\u003c/th\u003e\n",
              "      \u003ctd\u003eR50+L/32\u003c/td\u003e\n",
              "      \u003ctd\u003ei21k\u003c/td\u003e\n",
              "      \u003ctd\u003eR50_L_32-i21k-300ep-lr_0.001-aug_medium1-wd_0.1-do_0.1-sd_0.1\u003c/td\u003e\n",
              "      \u003ctd\u003e384\u003c/td\u003e\n",
              "      \u003ctd\u003e326.73\u003c/td\u003e\n",
              "      \u003ctd\u003eimagenet2012\u003c/td\u003e\n",
              "      \u003ctd\u003e0.85992\u003c/td\u003e\n",
              "      \u003ctd\u003eR50_L_32-i21k-300ep-lr_0.001-aug_medium1-wd_0.1-do_0.1-sd_0.1--imagenet2012-steps_20k-lr_0.01-res_384\u003c/td\u003e\n",
              "    \u003c/tr\u003e\n",
              "    \u003ctr\u003e\n",
              "      \u003cth\u003eR26+S/32\u003c/th\u003e\n",
              "      \u003ctd\u003eR26+S/32\u003c/td\u003e\n",
              "      \u003ctd\u003ei21k\u003c/td\u003e\n",
              "      \u003ctd\u003eR26_S_32-i21k-300ep-lr_0.001-aug_light1-wd_0.1-do_0.0-sd_0.0\u003c/td\u003e\n",
              "      \u003ctd\u003e384\u003c/td\u003e\n",
              "      \u003ctd\u003e560.40\u003c/td\u003e\n",
              "      \u003ctd\u003eimagenet2012\u003c/td\u003e\n",
              "      \u003ctd\u003e0.83846\u003c/td\u003e\n",
              "      \u003ctd\u003eR26_S_32-i21k-300ep-lr_0.001-aug_light1-wd_0.1-do_0.0-sd_0.0--imagenet2012-steps_20k-lr_0.01-res_384\u003c/td\u003e\n",
              "    \u003c/tr\u003e\n",
              "    \u003ctr\u003e\n",
              "      \u003cth\u003eTi/16\u003c/th\u003e\n",
              "      \u003ctd\u003eTi/16\u003c/td\u003e\n",
              "      \u003ctd\u003ei21k\u003c/td\u003e\n",
              "      \u003ctd\u003eTi_16-i21k-300ep-lr_0.001-aug_none-wd_0.03-do_0.0-sd_0.0\u003c/td\u003e\n",
              "      \u003ctd\u003e384\u003c/td\u003e\n",
              "      \u003ctd\u003e609.58\u003c/td\u003e\n",
              "      \u003ctd\u003eimagenet2012\u003c/td\u003e\n",
              "      \u003ctd\u003e0.78220\u003c/td\u003e\n",
              "      \u003ctd\u003eTi_16-i21k-300ep-lr_0.001-aug_none-wd_0.03-do_0.0-sd_0.0--imagenet2012-steps_20k-lr_0.03-res_384\u003c/td\u003e\n",
              "    \u003c/tr\u003e\n",
              "    \u003ctr\u003e\n",
              "      \u003cth\u003eB/32\u003c/th\u003e\n",
              "      \u003ctd\u003eB/32\u003c/td\u003e\n",
              "      \u003ctd\u003ei21k\u003c/td\u003e\n",
              "      \u003ctd\u003eB_32-i21k-300ep-lr_0.001-aug_light1-wd_0.1-do_0.0-sd_0.0\u003c/td\u003e\n",
              "      \u003ctd\u003e384\u003c/td\u003e\n",
              "      \u003ctd\u003e954.94\u003c/td\u003e\n",
              "      \u003ctd\u003eimagenet2012\u003c/td\u003e\n",
              "      \u003ctd\u003e0.83586\u003c/td\u003e\n",
              "      \u003ctd\u003eB_32-i21k-300ep-lr_0.001-aug_light1-wd_0.1-do_0.0-sd_0.0--imagenet2012-steps_20k-lr_0.01-res_384\u003c/td\u003e\n",
              "    \u003c/tr\u003e\n",
              "    \u003ctr\u003e\n",
              "      \u003cth\u003eS/32\u003c/th\u003e\n",
              "      \u003ctd\u003eS/32\u003c/td\u003e\n",
              "      \u003ctd\u003ei21k\u003c/td\u003e\n",
              "      \u003ctd\u003eS_32-i21k-300ep-lr_0.001-aug_none-wd_0.1-do_0.0-sd_0.0\u003c/td\u003e\n",
              "      \u003ctd\u003e384\u003c/td\u003e\n",
              "      \u003ctd\u003e2153.94\u003c/td\u003e\n",
              "      \u003ctd\u003eimagenet2012\u003c/td\u003e\n",
              "      \u003ctd\u003e0.79584\u003c/td\u003e\n",
              "      \u003ctd\u003eS_32-i21k-300ep-lr_0.001-aug_none-wd_0.1-do_0.0-sd_0.0--imagenet2012-steps_20k-lr_0.01-res_384\u003c/td\u003e\n",
              "    \u003c/tr\u003e\n",
              "    \u003ctr\u003e\n",
              "      \u003cth\u003eR+Ti/16\u003c/th\u003e\n",
              "      \u003ctd\u003eR+Ti/16\u003c/td\u003e\n",
              "      \u003ctd\u003ei21k\u003c/td\u003e\n",
              "      \u003ctd\u003eR_Ti_16-i21k-300ep-lr_0.001-aug_none-wd_0.03-do_0.0-sd_0.0\u003c/td\u003e\n",
              "      \u003ctd\u003e384\u003c/td\u003e\n",
              "      \u003ctd\u003e2425.77\u003c/td\u003e\n",
              "      \u003ctd\u003eimagenet2012\u003c/td\u003e\n",
              "      \u003ctd\u003e0.75402\u003c/td\u003e\n",
              "      \u003ctd\u003eR_Ti_16-i21k-300ep-lr_0.001-aug_none-wd_0.03-do_0.0-sd_0.0--imagenet2012-steps_20k-lr_0.03-res_384\u003c/td\u003e\n",
              "    \u003c/tr\u003e\n",
              "    \u003ctr\u003e\n",
              "      \u003cth\u003eB/8\u003c/th\u003e\n",
              "      \u003ctd\u003eB/8\u003c/td\u003e\n",
              "      \u003ctd\u003ei21k\u003c/td\u003e\n",
              "      \u003ctd\u003eB_8-i21k-300ep-lr_0.001-aug_medium2-wd_0.1-do_0.0-sd_0.0\u003c/td\u003e\n",
              "      \u003ctd\u003e224\u003c/td\u003e\n",
              "      \u003ctd\u003eNaN\u003c/td\u003e\n",
              "      \u003ctd\u003eimagenet2012\u003c/td\u003e\n",
              "      \u003ctd\u003e0.85578\u003c/td\u003e\n",
              "      \u003ctd\u003eB_8-i21k-300ep-lr_0.001-aug_medium2-wd_0.1-do_0.0-sd_0.0--imagenet2012-steps_20k-lr_0.03-res_224\u003c/td\u003e\n",
              "    \u003c/tr\u003e\n",
              "  \u003c/tbody\u003e\n",
              "\u003c/table\u003e\n",
              "\u003c/div\u003e"
            ],
            "text/plain": [
              "              name  ...                                                                                         adapt_filename\n",
              "name                ...                                                                                                       \n",
              "L/16          L/16  ...      L_16-i21k-300ep-lr_0.001-aug_strong1-wd_0.1-do_0.0-sd_0.0--imagenet2012-steps_20k-lr_0.01-res_384\n",
              "B/16          B/16  ...      B_16-i21k-300ep-lr_0.001-aug_medium1-wd_0.1-do_0.0-sd_0.0--imagenet2012-steps_20k-lr_0.03-res_384\n",
              "S/16          S/16  ...      S_16-i21k-300ep-lr_0.001-aug_light1-wd_0.03-do_0.0-sd_0.0--imagenet2012-steps_20k-lr_0.03-res_384\n",
              "R50+L/32  R50+L/32  ...  R50_L_32-i21k-300ep-lr_0.001-aug_medium1-wd_0.1-do_0.1-sd_0.1--imagenet2012-steps_20k-lr_0.01-res_384\n",
              "R26+S/32  R26+S/32  ...   R26_S_32-i21k-300ep-lr_0.001-aug_light1-wd_0.1-do_0.0-sd_0.0--imagenet2012-steps_20k-lr_0.01-res_384\n",
              "Ti/16        Ti/16  ...       Ti_16-i21k-300ep-lr_0.001-aug_none-wd_0.03-do_0.0-sd_0.0--imagenet2012-steps_20k-lr_0.03-res_384\n",
              "B/32          B/32  ...       B_32-i21k-300ep-lr_0.001-aug_light1-wd_0.1-do_0.0-sd_0.0--imagenet2012-steps_20k-lr_0.01-res_384\n",
              "S/32          S/32  ...         S_32-i21k-300ep-lr_0.001-aug_none-wd_0.1-do_0.0-sd_0.0--imagenet2012-steps_20k-lr_0.01-res_384\n",
              "R+Ti/16    R+Ti/16  ...     R_Ti_16-i21k-300ep-lr_0.001-aug_none-wd_0.03-do_0.0-sd_0.0--imagenet2012-steps_20k-lr_0.03-res_384\n",
              "B/8            B/8  ...       B_8-i21k-300ep-lr_0.001-aug_medium2-wd_0.1-do_0.0-sd_0.0--imagenet2012-steps_20k-lr_0.03-res_224\n",
              "\n",
              "[10 rows x 8 columns]"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 21
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 388
        },
        "id": "iPYUh4V1EdHg",
        "outputId": "fbe7aa0d-5ddd-46d3-c3bb-c6006a995ec2"
      },
      "source": [
        "# Visualize the 2 (resolution) * 9 (models) * 8 (lr, steps) finetunings for a\n",
        "# single dataset (Pets37).\n",
        "# Note how larger models get better scores up to B/16 @384 even on this tiny\n",
        "# dataset, if pre-trained sufficiently.\n",
        "sns.relplot(\n",
        "    data=best_df.query('adapt_ds==\"oxford_iiit_pet\"'),\n",
        "    x='infer_samples_per_sec',\n",
        "    y='adapt_final_val',\n",
        "    hue='name',\n",
        "    style='adapt_resolution'\n",
        ")\n",
        "plt.gca().set_xscale('log');"
      ],
      "execution_count": 22,
      "outputs": [
        {
          "output_type": "display_data",
          "data": {
            "image/png": "[encoded data, 38500 characters]\n",
            "text/plain": [
              "\u003cFigure size 459.375x360 with 1 Axes\u003e"
            ]
          },
          "metadata": {
            "tags": [],
            "needs_background": "light"
          }
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 817
        },
        "id": "9i7Y9AcpFtlr",
        "outputId": "1a92afd1-ce4f-4e3a-82c9-f0f0a8a11a99"
      },
      "source": [
        "# More details for a single pre-trained checkpoint.\n",
        "best_df.query('name==\"R26+S/32\" and adapt_ds==\"oxford_iiit_pet\"')[[\n",
        "  col for col in best_df.columns if col.startswith('adapt_')\n",
        "]].sort_values('adapt_final_val')"
      ],
      "execution_count": 23,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/html": [
              "\u003cdiv\u003e\n",
              "\u003cstyle scoped\u003e\n",
              "    .dataframe tbody tr th:only-of-type {\n",
              "        vertical-align: middle;\n",
              "    }\n",
              "\n",
              "    .dataframe tbody tr th {\n",
              "        vertical-align: top;\n",
              "    }\n",
              "\n",
              "    .dataframe thead th {\n",
              "        text-align: right;\n",
              "    }\n",
              "\u003c/style\u003e\n",
              "\u003ctable border=\"1\" class=\"dataframe\"\u003e\n",
              "  \u003cthead\u003e\n",
              "    \u003ctr style=\"text-align: right;\"\u003e\n",
              "      \u003cth\u003e\u003c/th\u003e\n",
              "      \u003cth\u003eadapt_ds\u003c/th\u003e\n",
              "      \u003cth\u003eadapt_lr\u003c/th\u003e\n",
              "      \u003cth\u003eadapt_steps\u003c/th\u003e\n",
              "      \u003cth\u003eadapt_resolution\u003c/th\u003e\n",
              "      \u003cth\u003eadapt_final_val\u003c/th\u003e\n",
              "      \u003cth\u003eadapt_final_test\u003c/th\u003e\n",
              "      \u003cth\u003eadapt_filename\u003c/th\u003e\n",
              "    \u003c/tr\u003e\n",
              "  \u003c/thead\u003e\n",
              "  \u003ctbody\u003e\n",
              "    \u003ctr\u003e\n",
              "      \u003cth\u003e5730\u003c/th\u003e\n",
              "      \u003ctd\u003eoxford_iiit_pet\u003c/td\u003e\n",
              "      \u003ctd\u003e0.030\u003c/td\u003e\n",
              "      \u003ctd\u003e2500\u003c/td\u003e\n",
              "      \u003ctd\u003e224\u003c/td\u003e\n",
              "      \u003ctd\u003e0.894022\u003c/td\u003e\n",
              "      \u003ctd\u003e0.867539\u003c/td\u003e\n",
              "      \u003ctd\u003eR26_S_32-i21k-300ep-lr_0.001-aug_light1-wd_0.1-do_0.0-sd_0.0--oxford_iiit_pet-steps_2k-lr_0.03-res_224\u003c/td\u003e\n",
              "    \u003c/tr\u003e\n",
              "    \u003ctr\u003e\n",
              "      \u003cth\u003e5762\u003c/th\u003e\n",
              "      \u003ctd\u003eoxford_iiit_pet\u003c/td\u003e\n",
              "      \u003ctd\u003e0.030\u003c/td\u003e\n",
              "      \u003ctd\u003e2500\u003c/td\u003e\n",
              "      \u003ctd\u003e384\u003c/td\u003e\n",
              "      \u003ctd\u003e0.915761\u003c/td\u003e\n",
              "      \u003ctd\u003e0.894794\u003c/td\u003e\n",
              "      \u003ctd\u003eR26_S_32-i21k-300ep-lr_0.001-aug_light1-wd_0.1-do_0.0-sd_0.0--oxford_iiit_pet-steps_2k-lr_0.03-res_384\u003c/td\u003e\n",
              "    \u003c/tr\u003e\n",
              "    \u003ctr\u003e\n",
              "      \u003cth\u003e5734\u003c/th\u003e\n",
              "      \u003ctd\u003eoxford_iiit_pet\u003c/td\u003e\n",
              "      \u003ctd\u003e0.030\u003c/td\u003e\n",
              "      \u003ctd\u003e500\u003c/td\u003e\n",
              "      \u003ctd\u003e224\u003c/td\u003e\n",
              "      \u003ctd\u003e0.932065\u003c/td\u003e\n",
              "      \u003ctd\u003e0.916326\u003c/td\u003e\n",
              "      \u003ctd\u003eR26_S_32-i21k-300ep-lr_0.001-aug_light1-wd_0.1-do_0.0-sd_0.0--oxford_iiit_pet-steps_0k-lr_0.03-res_224\u003c/td\u003e\n",
              "    \u003c/tr\u003e\n",
              "    \u003ctr\u003e\n",
              "      \u003cth\u003e5731\u003c/th\u003e\n",
              "      \u003ctd\u003eoxford_iiit_pet\u003c/td\u003e\n",
              "      \u003ctd\u003e0.010\u003c/td\u003e\n",
              "      \u003ctd\u003e2500\u003c/td\u003e\n",
              "      \u003ctd\u003e224\u003c/td\u003e\n",
              "      \u003ctd\u003e0.934783\u003c/td\u003e\n",
              "      \u003ctd\u003e0.908694\u003c/td\u003e\n",
              "      \u003ctd\u003eR26_S_32-i21k-300ep-lr_0.001-aug_light1-wd_0.1-do_0.0-sd_0.0--oxford_iiit_pet-steps_2k-lr_0.01-res_224\u003c/td\u003e\n",
              "    \u003c/tr\u003e\n",
              "    \u003ctr\u003e\n",
              "      \u003cth\u003e5732\u003c/th\u003e\n",
              "      \u003ctd\u003eoxford_iiit_pet\u003c/td\u003e\n",
              "      \u003ctd\u003e0.003\u003c/td\u003e\n",
              "      \u003ctd\u003e2500\u003c/td\u003e\n",
              "      \u003ctd\u003e224\u003c/td\u003e\n",
              "      \u003ctd\u003e0.945652\u003c/td\u003e\n",
              "      \u003ctd\u003e0.926410\u003c/td\u003e\n",
              "      \u003ctd\u003eR26_S_32-i21k-300ep-lr_0.001-aug_light1-wd_0.1-do_0.0-sd_0.0--oxford_iiit_pet-steps_2k-lr_0.003-res_224\u003c/td\u003e\n",
              "    \u003c/tr\u003e\n",
              "    \u003ctr\u003e\n",
              "      \u003cth\u003e5763\u003c/th\u003e\n",
              "      \u003ctd\u003eoxford_iiit_pet\u003c/td\u003e\n",
              "      \u003ctd\u003e0.010\u003c/td\u003e\n",
              "      \u003ctd\u003e2500\u003c/td\u003e\n",
              "      \u003ctd\u003e384\u003c/td\u003e\n",
              "      \u003ctd\u003e0.945652\u003c/td\u003e\n",
              "      \u003ctd\u003e0.934315\u003c/td\u003e\n",
              "      \u003ctd\u003eR26_S_32-i21k-300ep-lr_0.001-aug_light1-wd_0.1-do_0.0-sd_0.0--oxford_iiit_pet-steps_2k-lr_0.01-res_384\u003c/td\u003e\n",
              "    \u003c/tr\u003e\n",
              "    \u003ctr\u003e\n",
              "      \u003cth\u003e5735\u003c/th\u003e\n",
              "      \u003ctd\u003eoxford_iiit_pet\u003c/td\u003e\n",
              "      \u003ctd\u003e0.010\u003c/td\u003e\n",
              "      \u003ctd\u003e500\u003c/td\u003e\n",
              "      \u003ctd\u003e224\u003c/td\u003e\n",
              "      \u003ctd\u003e0.951087\u003c/td\u003e\n",
              "      \u003ctd\u003e0.928046\u003c/td\u003e\n",
              "      \u003ctd\u003eR26_S_32-i21k-300ep-lr_0.001-aug_light1-wd_0.1-do_0.0-sd_0.0--oxford_iiit_pet-steps_0k-lr_0.01-res_224\u003c/td\u003e\n",
              "    \u003c/tr\u003e\n",
              "    \u003ctr\u003e\n",
              "      \u003cth\u003e5766\u003c/th\u003e\n",
              "      \u003ctd\u003eoxford_iiit_pet\u003c/td\u003e\n",
              "      \u003ctd\u003e0.030\u003c/td\u003e\n",
              "      \u003ctd\u003e500\u003c/td\u003e\n",
              "      \u003ctd\u003e384\u003c/td\u003e\n",
              "      \u003ctd\u003e0.953804\u003c/td\u003e\n",
              "      \u003ctd\u003e0.941128\u003c/td\u003e\n",
              "      \u003ctd\u003eR26_S_32-i21k-300ep-lr_0.001-aug_light1-wd_0.1-do_0.0-sd_0.0--oxford_iiit_pet-steps_0k-lr_0.03-res_384\u003c/td\u003e\n",
              "    \u003c/tr\u003e\n",
              "    \u003ctr\u003e\n",
              "      \u003cth\u003e5733\u003c/th\u003e\n",
              "      \u003ctd\u003eoxford_iiit_pet\u003c/td\u003e\n",
              "      \u003ctd\u003e0.001\u003c/td\u003e\n",
              "      \u003ctd\u003e2500\u003c/td\u003e\n",
              "      \u003ctd\u003e224\u003c/td\u003e\n",
              "      \u003ctd\u003e0.959239\u003c/td\u003e\n",
              "      \u003ctd\u003e0.931044\u003c/td\u003e\n",
              "      \u003ctd\u003eR26_S_32-i21k-300ep-lr_0.001-aug_light1-wd_0.1-do_0.0-sd_0.0--oxford_iiit_pet-steps_2k-lr_0.001-res_224\u003c/td\u003e\n",
              "    \u003c/tr\u003e\n",
              "    \u003ctr\u003e\n",
              "      \u003cth\u003e5737\u003c/th\u003e\n",
              "      \u003ctd\u003eoxford_iiit_pet\u003c/td\u003e\n",
              "      \u003ctd\u003e0.001\u003c/td\u003e\n",
              "      \u003ctd\u003e500\u003c/td\u003e\n",
              "      \u003ctd\u003e224\u003c/td\u003e\n",
              "      \u003ctd\u003e0.959239\u003c/td\u003e\n",
              "      \u003ctd\u003e0.923412\u003c/td\u003e\n",
              "      \u003ctd\u003eR26_S_32-i21k-300ep-lr_0.001-aug_light1-wd_0.1-do_0.0-sd_0.0--oxford_iiit_pet-steps_0k-lr_0.001-res_224\u003c/td\u003e\n",
              "    \u003c/tr\u003e\n",
              "    \u003ctr\u003e\n",
              "      \u003cth\u003e5764\u003c/th\u003e\n",
              "      \u003ctd\u003eoxford_iiit_pet\u003c/td\u003e\n",
              "      \u003ctd\u003e0.003\u003c/td\u003e\n",
              "      \u003ctd\u003e2500\u003c/td\u003e\n",
              "      \u003ctd\u003e384\u003c/td\u003e\n",
              "      \u003ctd\u003e0.961957\u003c/td\u003e\n",
              "      \u003ctd\u003e0.941401\u003c/td\u003e\n",
              "      \u003ctd\u003eR26_S_32-i21k-300ep-lr_0.001-aug_light1-wd_0.1-do_0.0-sd_0.0--oxford_iiit_pet-steps_2k-lr_0.003-res_384\u003c/td\u003e\n",
              "    \u003c/tr\u003e\n",
              "    \u003ctr\u003e\n",
              "      \u003cth\u003e5765\u003c/th\u003e\n",
              "      \u003ctd\u003eoxford_iiit_pet\u003c/td\u003e\n",
              "      \u003ctd\u003e0.001\u003c/td\u003e\n",
              "      \u003ctd\u003e2500\u003c/td\u003e\n",
              "      \u003ctd\u003e384\u003c/td\u003e\n",
              "      \u003ctd\u003e0.961957\u003c/td\u003e\n",
              "      \u003ctd\u003e0.944672\u003c/td\u003e\n",
              "      \u003ctd\u003eR26_S_32-i21k-300ep-lr_0.001-aug_light1-wd_0.1-do_0.0-sd_0.0--oxford_iiit_pet-steps_2k-lr_0.001-res_384\u003c/td\u003e\n",
              "    \u003c/tr\u003e\n",
              "    \u003ctr\u003e\n",
              "      \u003cth\u003e5767\u003c/th\u003e\n",
              "      \u003ctd\u003eoxford_iiit_pet\u003c/td\u003e\n",
              "      \u003ctd\u003e0.010\u003c/td\u003e\n",
              "      \u003ctd\u003e500\u003c/td\u003e\n",
              "      \u003ctd\u003e384\u003c/td\u003e\n",
              "      \u003ctd\u003e0.961957\u003c/td\u003e\n",
              "      \u003ctd\u003e0.946852\u003c/td\u003e\n",
              "      \u003ctd\u003eR26_S_32-i21k-300ep-lr_0.001-aug_light1-wd_0.1-do_0.0-sd_0.0--oxford_iiit_pet-steps_0k-lr_0.01-res_384\u003c/td\u003e\n",
              "    \u003c/tr\u003e\n",
              "    \u003ctr\u003e\n",
              "      \u003cth\u003e5769\u003c/th\u003e\n",
              "      \u003ctd\u003eoxford_iiit_pet\u003c/td\u003e\n",
              "      \u003ctd\u003e0.001\u003c/td\u003e\n",
              "      \u003ctd\u003e500\u003c/td\u003e\n",
              "      \u003ctd\u003e384\u003c/td\u003e\n",
              "      \u003ctd\u003e0.964674\u003c/td\u003e\n",
              "      \u003ctd\u003e0.930499\u003c/td\u003e\n",
              "      \u003ctd\u003eR26_S_32-i21k-300ep-lr_0.001-aug_light1-wd_0.1-do_0.0-sd_0.0--oxford_iiit_pet-steps_0k-lr_0.001-res_384\u003c/td\u003e\n",
              "    \u003c/tr\u003e\n",
              "    \u003ctr\u003e\n",
              "      \u003cth\u003e5736\u003c/th\u003e\n",
              "      \u003ctd\u003eoxford_iiit_pet\u003c/td\u003e\n",
              "      \u003ctd\u003e0.003\u003c/td\u003e\n",
              "      \u003ctd\u003e500\u003c/td\u003e\n",
              "      \u003ctd\u003e224\u003c/td\u003e\n",
              "      \u003ctd\u003e0.967391\u003c/td\u003e\n",
              "      \u003ctd\u003e0.931589\u003c/td\u003e\n",
              "      \u003ctd\u003eR26_S_32-i21k-300ep-lr_0.001-aug_light1-wd_0.1-do_0.0-sd_0.0--oxford_iiit_pet-steps_0k-lr_0.003-res_224\u003c/td\u003e\n",
              "    \u003c/tr\u003e\n",
              "    \u003ctr\u003e\n",
              "      \u003cth\u003e5768\u003c/th\u003e\n",
              "      \u003ctd\u003eoxford_iiit_pet\u003c/td\u003e\n",
              "      \u003ctd\u003e0.003\u003c/td\u003e\n",
              "      \u003ctd\u003e500\u003c/td\u003e\n",
              "      \u003ctd\u003e384\u003c/td\u003e\n",
              "      \u003ctd\u003e0.967391\u003c/td\u003e\n",
              "      \u003ctd\u003e0.941946\u003c/td\u003e\n",
              "      \u003ctd\u003eR26_S_32-i21k-300ep-lr_0.001-aug_light1-wd_0.1-do_0.0-sd_0.0--oxford_iiit_pet-steps_0k-lr_0.003-res_384\u003c/td\u003e\n",
              "    \u003c/tr\u003e\n",
              "  \u003c/tbody\u003e\n",
              "\u003c/table\u003e\n",
              "\u003c/div\u003e"
            ],
            "text/plain": [
              "             adapt_ds  ...                                                                                           adapt_filename\n",
              "5730  oxford_iiit_pet  ...   R26_S_32-i21k-300ep-lr_0.001-aug_light1-wd_0.1-do_0.0-sd_0.0--oxford_iiit_pet-steps_2k-lr_0.03-res_224\n",
              "5762  oxford_iiit_pet  ...   R26_S_32-i21k-300ep-lr_0.001-aug_light1-wd_0.1-do_0.0-sd_0.0--oxford_iiit_pet-steps_2k-lr_0.03-res_384\n",
              "5734  oxford_iiit_pet  ...   R26_S_32-i21k-300ep-lr_0.001-aug_light1-wd_0.1-do_0.0-sd_0.0--oxford_iiit_pet-steps_0k-lr_0.03-res_224\n",
              "5731  oxford_iiit_pet  ...   R26_S_32-i21k-300ep-lr_0.001-aug_light1-wd_0.1-do_0.0-sd_0.0--oxford_iiit_pet-steps_2k-lr_0.01-res_224\n",
              "5732  oxford_iiit_pet  ...  R26_S_32-i21k-300ep-lr_0.001-aug_light1-wd_0.1-do_0.0-sd_0.0--oxford_iiit_pet-steps_2k-lr_0.003-res_224\n",
              "5763  oxford_iiit_pet  ...   R26_S_32-i21k-300ep-lr_0.001-aug_light1-wd_0.1-do_0.0-sd_0.0--oxford_iiit_pet-steps_2k-lr_0.01-res_384\n",
              "5735  oxford_iiit_pet  ...   R26_S_32-i21k-300ep-lr_0.001-aug_light1-wd_0.1-do_0.0-sd_0.0--oxford_iiit_pet-steps_0k-lr_0.01-res_224\n",
              "5766  oxford_iiit_pet  ...   R26_S_32-i21k-300ep-lr_0.001-aug_light1-wd_0.1-do_0.0-sd_0.0--oxford_iiit_pet-steps_0k-lr_0.03-res_384\n",
              "5733  oxford_iiit_pet  ...  R26_S_32-i21k-300ep-lr_0.001-aug_light1-wd_0.1-do_0.0-sd_0.0--oxford_iiit_pet-steps_2k-lr_0.001-res_224\n",
              "5737  oxford_iiit_pet  ...  R26_S_32-i21k-300ep-lr_0.001-aug_light1-wd_0.1-do_0.0-sd_0.0--oxford_iiit_pet-steps_0k-lr_0.001-res_224\n",
              "5764  oxford_iiit_pet  ...  R26_S_32-i21k-300ep-lr_0.001-aug_light1-wd_0.1-do_0.0-sd_0.0--oxford_iiit_pet-steps_2k-lr_0.003-res_384\n",
              "5765  oxford_iiit_pet  ...  R26_S_32-i21k-300ep-lr_0.001-aug_light1-wd_0.1-do_0.0-sd_0.0--oxford_iiit_pet-steps_2k-lr_0.001-res_384\n",
              "5767  oxford_iiit_pet  ...   R26_S_32-i21k-300ep-lr_0.001-aug_light1-wd_0.1-do_0.0-sd_0.0--oxford_iiit_pet-steps_0k-lr_0.01-res_384\n",
              "5769  oxford_iiit_pet  ...  R26_S_32-i21k-300ep-lr_0.001-aug_light1-wd_0.1-do_0.0-sd_0.0--oxford_iiit_pet-steps_0k-lr_0.001-res_384\n",
              "5736  oxford_iiit_pet  ...  R26_S_32-i21k-300ep-lr_0.001-aug_light1-wd_0.1-do_0.0-sd_0.0--oxford_iiit_pet-steps_0k-lr_0.003-res_224\n",
              "5768  oxford_iiit_pet  ...  R26_S_32-i21k-300ep-lr_0.001-aug_light1-wd_0.1-do_0.0-sd_0.0--oxford_iiit_pet-steps_0k-lr_0.003-res_384\n",
              "\n",
              "[16 rows x 7 columns]"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 23
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "wh_SLkQtQ6K4"
      },
      "source": [
        "### Load a checkpoint"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "cwWBMJSpQ7EZ",
        "outputId": "c7e94c93-0e17-4209-bedf-c6a394c4a8cf"
      },
      "source": [
        "# Select a value from \"adapt_filename\" above that is a fine-tuned checkpoint.\n",
        "filename = 'R26_S_32-i21k-300ep-lr_0.001-aug_light1-wd_0.1-do_0.0-sd_0.0--oxford_iiit_pet-steps_0k-lr_0.003-res_384'\n",
        "\n",
        "tfds_name = filename.split('--')[1].split('-')[0]\n",
        "model_config = models_config.AUGREG_CONFIGS[filename.split('-')[0]]\n",
        "resolution = int(filename.split('_')[-1])\n",
        "path = f'gs://vit_models/augreg/{filename}.npz'\n",
        "\n",
        "print(f'{tf.io.gfile.stat(path).length / 1024 / 1024:.1f} MiB - {path}')"
      ],
      "execution_count": 24,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "137.8 MiB - gs://vit_models/augreg/R26_S_32-i21k-300ep-lr_0.001-aug_light1-wd_0.1-do_0.0-sd_0.0--oxford_iiit_pet-steps_0k-lr_0.003-res_384.npz\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 1000,
          "referenced_widgets": [
            "225ebb8a26264dca90b312edb1e73fb6",
            "7d3d39e266ab4a81983fbf829795f3a7",
            "c529bc24ac1943dd8628873d1648c275",
            "d4e82b87ab4a485c9e15f3fd6e77a969",
            "7bdc1b98ea574b8e99fc1bd3bbd2b456",
            "63c6ec28cbad46359e38d1f921b15ffa",
            "d9d5245d91fe42999d8b8693674880f7",
            "113db5d71db84749805a9394e89b9e0b",
            "a5a294a2b18f4c078b04b8126be1d84b",
            "c2ce9e71f07a4be59374466f7a8f33ff",
            "608ff4fab08743d8ade2cf99ac8dcd69",
            "107c6e554e9a4a2a913f1684761d3736",
            "0e51cd83a88140728b937542605c7e62",
            "a6f978d87f114e4ea5fd99ca45c67ba1",
            "c9598f8547494601afc86dbdcfc88516",
            "c625f7758d544bb6937f265f209cd2c8",
            "132048250a42471785d742ecfa679a7b",
            "a9d7f91c071345ce9012a39adc3037b5",
            "79577fafaa0c46cc85e057d50e1f640f",
            "a0e3deff575141949907cb8647a769be",
            "770dd8577d454c2e82774043376809c1",
            "a73ff8ef581e44d8a8503f4cf25df02f",
            "a54f014c7dff445f846b97281652c14f",
            "b3f5112e065d4cb884c514a8e4097b38",
            "43a06b56ab8742cabec36cd19b51945c",
            "b0e5b9a1cd5c4b918300b2006daacbdb",
            "acc7c1e9e5ec470bb7d22a49fbab75dc",
            "5f70500c82aa4c189b85c4765047305f",
            "df9bdd9cfdb74967a5cd83fbcbc3d57e",
            "8529b470a86640a1819098d4b809d7c8",
            "8f0adf499e434dc2863bdf6537c08454",
            "e5cafaad07cd45dfb2329b4fa3b18ad4",
            "f0da817110a947d2b0b54a43c7879ec6",
            "13b28f5de68b428bba06b648c201c5e2",
            "5a54036091c84c6582ccf8ea545f7f4c",
            "ee0a499dc13b4a6eb36627a242d6e92a",
            "febe24498bc0445ab747c120c6abe261",
            "56f06ed9e69d40f9b5f39e1550fd2eb1",
            "7d7deac790b4445cb170356ae88b36e9",
            "747b920f1510475dad43d49ba782be4d",
            "3f995dd7906341eb806cb4cfc52c78eb",
            "80048ea0d80c4bf5a9172380e948799d",
            "716975e14c8f4b59ade02dc968d58b5e",
            "2acfe7938b3543bf8b9542556f5a23b8",
            "817ad1768fa14aa0b548e259d2162944",
            "413ba3075e7d4f54989b8de944aa0759",
            "2a3f564c7c884ff1936d1c6db79ac40b",
            "3b863d829bbc4c1ea618ed5dfb3f0538",
            "ee27f121919d4345acd966ad98518b3a",
            "fa485b969a6c4893bbb42904fed24eeb",
            "7d4bff3e31c4431e91054562d51f80e8",
            "902db9dfbf8543a4956409a22c2efa80",
            "74a0420bbcc64cd0a662f8e27a23cf5f",
            "9ff4e5d0a6cd4953a9b04cdd115de440",
            "661c19ff44f048eaa8b21ee203319c68",
            "e9bc91cf5f67406689c7fe90800f729f"
          ]
        },
        "id": "MKW2flB_Sk1u",
        "outputId": "f7be2692-d725-4799-f165-9724cd941799"
      },
      "source": [
        "# Fetch dataset that the checkpoint was finetuned on.\n",
        "# (Note that automatic download does not work with imagenet2012)\n",
        "ds, ds_info = tfds.load(tfds_name, with_info=True)\n",
        "ds_info"
      ],
      "execution_count": 25,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "INFO:absl:Load pre-computed DatasetInfo (eg: splits, num examples,...) from GCS: oxford_iiit_pet/3.2.0\n",
            "INFO:absl:Load dataset info from /tmp/tmpjdu62jcgtfds\n",
            "INFO:absl:Generating dataset oxford_iiit_pet (/root/tensorflow_datasets/oxford_iiit_pet/3.2.0)\n"
          ],
          "name": "stderr"
        },
        {
          "output_type": "stream",
          "text": [
            "\u001b[1mDownloading and preparing dataset oxford_iiit_pet/3.2.0 (download: 773.52 MiB, generated: 774.69 MiB, total: 1.51 GiB) to /root/tensorflow_datasets/oxford_iiit_pet/3.2.0...\u001b[0m\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "display_data",
          "data": {
            "application/vnd.jupyter.widget-view+json": {
              "model_id": "225ebb8a26264dca90b312edb1e73fb6",
              "version_minor": 0,
              "version_major": 2
            },
            "text/plain": [
              "HBox(children=(FloatProgress(value=1.0, bar_style='info', description='Dl Completed...', max=1.0, style=Progre…"
            ]
          },
          "metadata": {
            "tags": []
          }
        },
        {
          "output_type": "display_data",
          "data": {
            "application/vnd.jupyter.widget-view+json": {
              "model_id": "a5a294a2b18f4c078b04b8126be1d84b",
              "version_minor": 0,
              "version_major": 2
            },
            "text/plain": [
              "HBox(children=(FloatProgress(value=1.0, bar_style='info', description='Dl Size...', max=1.0, style=ProgressSty…"
            ]
          },
          "metadata": {
            "tags": []
          }
        },
        {
          "output_type": "display_data",
          "data": {
            "application/vnd.jupyter.widget-view+json": {
              "model_id": "132048250a42471785d742ecfa679a7b",
              "version_minor": 0,
              "version_major": 2
            },
            "text/plain": [
              "HBox(children=(FloatProgress(value=1.0, bar_style='info', description='Extraction completed...', max=1.0, styl…"
            ]
          },
          "metadata": {
            "tags": []
          }
        },
        {
          "output_type": "stream",
          "text": [
            "INFO:absl:Downloading http://www.robots.ox.ac.uk/~vgg/data/pets/data/annotations.tar.gz into /root/tensorflow_datasets/downloads/robots.ox.ac.uk_vgg_pets_annotationsL8B_bMDM0_iUhyDddl6Yu-1togBm2g5ZH8E6r80wOQo.tar.gz.tmp.4604a2666aa54d328f47b684e6b8c5c7...\n",
            "INFO:absl:Downloading http://www.robots.ox.ac.uk/~vgg/data/pets/data/images.tar.gz into /root/tensorflow_datasets/downloads/robots.ox.ac.uk_vgg_pets_imageswMR1o1DWRq_DHWToagdXedb7P88RHpceK3WqG77VVwU.tar.gz.tmp.87f3d2b8bfb14abcadb60abb95bc060d...\n",
            "INFO:absl:Generating split train\n"
          ],
          "name": "stderr"
        },
        {
          "output_type": "stream",
          "text": [
            "\n",
            "\n",
            "\n",
            "\n",
            "\n",
            "\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "display_data",
          "data": {
            "application/vnd.jupyter.widget-view+json": {
              "model_id": "43a06b56ab8742cabec36cd19b51945c",
              "version_minor": 0,
              "version_major": 2
            },
            "text/plain": [
              "HBox(children=(FloatProgress(value=1.0, bar_style='info', max=1.0), HTML(value='')))"
            ]
          },
          "metadata": {
            "tags": []
          }
        },
        {
          "output_type": "stream",
          "text": [
            "\rShuffling and writing examples to /root/tensorflow_datasets/oxford_iiit_pet/3.2.0.incompleteADR8WE/oxford_iiit_pet-train.tfrecord\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "display_data",
          "data": {
            "application/vnd.jupyter.widget-view+json": {
              "model_id": "f0da817110a947d2b0b54a43c7879ec6",
              "version_minor": 0,
              "version_major": 2
            },
            "text/plain": [
              "HBox(children=(FloatProgress(value=0.0, max=3680.0), HTML(value='')))"
            ]
          },
          "metadata": {
            "tags": []
          }
        },
        {
          "output_type": "stream",
          "text": [
            "INFO:absl:Done writing /root/tensorflow_datasets/oxford_iiit_pet/3.2.0.incompleteADR8WE/oxford_iiit_pet-train.tfrecord. Shard lengths: [920, 920, 920, 920]\n",
            "INFO:absl:Generating split test\n"
          ],
          "name": "stderr"
        },
        {
          "output_type": "stream",
          "text": [
            "\r"
          ],
          "name": "stdout"
        },
        {
          "output_type": "display_data",
          "data": {
            "application/vnd.jupyter.widget-view+json": {
              "model_id": "3f995dd7906341eb806cb4cfc52c78eb",
              "version_minor": 0,
              "version_major": 2
            },
            "text/plain": [
              "HBox(children=(FloatProgress(value=1.0, bar_style='info', max=1.0), HTML(value='')))"
            ]
          },
          "metadata": {
            "tags": []
          }
        },
        {
          "output_type": "stream",
          "text": [
            "\rShuffling and writing examples to /root/tensorflow_datasets/oxford_iiit_pet/3.2.0.incompleteADR8WE/oxford_iiit_pet-test.tfrecord\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "display_data",
          "data": {
            "application/vnd.jupyter.widget-view+json": {
              "model_id": "ee27f121919d4345acd966ad98518b3a",
              "version_minor": 0,
              "version_major": 2
            },
            "text/plain": [
              "HBox(children=(FloatProgress(value=0.0, max=3669.0), HTML(value='')))"
            ]
          },
          "metadata": {
            "tags": []
          }
        },
        {
          "output_type": "stream",
          "text": [
            "INFO:absl:Done writing /root/tensorflow_datasets/oxford_iiit_pet/3.2.0.incompleteADR8WE/oxford_iiit_pet-test.tfrecord. Shard lengths: [917, 917, 918, 917]\n",
            "INFO:absl:Skipping computing stats for mode ComputeStatsMode.SKIP.\n",
            "INFO:absl:Constructing tf.data.Dataset for split None, from /root/tensorflow_datasets/oxford_iiit_pet/3.2.0\n"
          ],
          "name": "stderr"
        },
        {
          "output_type": "stream",
          "text": [
            "\u001b[1mDataset oxford_iiit_pet downloaded and prepared to /root/tensorflow_datasets/oxford_iiit_pet/3.2.0. Subsequent calls will reuse this data.\u001b[0m\n",
            "\r"
          ],
          "name": "stdout"
        },
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "tfds.core.DatasetInfo(\n",
              "    name='oxford_iiit_pet',\n",
              "    version=3.2.0,\n",
              "    description='The Oxford-IIIT pet dataset is a 37 category pet image dataset with roughly 200\n",
              "images for each class. The images have large variations in scale, pose and\n",
              "lighting. All images have an associated ground truth annotation of breed.',\n",
              "    homepage='http://www.robots.ox.ac.uk/~vgg/data/pets/',\n",
              "    features=FeaturesDict({\n",
              "        'file_name': Text(shape=(), dtype=tf.string),\n",
              "        'image': Image(shape=(None, None, 3), dtype=tf.uint8),\n",
              "        'label': ClassLabel(shape=(), dtype=tf.int64, num_classes=37),\n",
              "        'segmentation_mask': Image(shape=(None, None, 1), dtype=tf.uint8),\n",
              "        'species': ClassLabel(shape=(), dtype=tf.int64, num_classes=2),\n",
              "    }),\n",
              "    total_num_examples=7349,\n",
              "    splits={\n",
              "        'test': 3669,\n",
              "        'train': 3680,\n",
              "    },\n",
              "    supervised_keys=('image', 'label'),\n",
              "    citation=\"\"\"@InProceedings{parkhi12a,\n",
              "      author       = \"Parkhi, O. M. and Vedaldi, A. and Zisserman, A. and Jawahar, C.~V.\",\n",
              "      title        = \"Cats and Dogs\",\n",
              "      booktitle    = \"IEEE Conference on Computer Vision and Pattern Recognition\",\n",
              "      year         = \"2012\",\n",
              "    }\"\"\",\n",
              "    redistribution_info=,\n",
              ")"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 25
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "KtmV4tRCSX8n"
      },
      "source": [
        "# Get model instance - no weights are initialized yet.\n",
        "model = models.VisionTransformer(\n",
        "    num_classes=ds_info.features['label'].num_classes, **model_config)"
      ],
      "execution_count": 26,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "GU6Tw1KvScBA"
      },
      "source": [
        "# Load a checkpoint from cloud - for large checkpoints this can take a while...\n",
        "params = checkpoint.load(path)"
      ],
      "execution_count": 27,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Js-QtAqeUJCo"
      },
      "source": [
        "# Get a single example from dataset for inference.\n",
        "d = next(iter(ds['test']))"
      ],
      "execution_count": 28,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 269
        },
        "id": "xpovSMYFUHAj",
        "outputId": "b0a70f79-2ddb-40f7-9e48-16783404aeec"
      },
      "source": [
        "def pp(img, sz):\n",
        "  \"\"\"Simple image preprocessing.\"\"\"\n",
        "  img = tf.cast(img, float) / 255.0\n",
        "  img = tf.image.resize(img, [sz, sz])\n",
        "  return img\n",
        "\n",
        "plt.imshow(pp(d['image'], resolution));"
      ],
      "execution_count": 29,
      "outputs": [
        {
          "output_type": "display_data",
          "data": {
            "image/png": "[encoded data, 145852 characters]\n",
            "text/plain": [
              "\u003cFigure size 432x288 with 1 Axes\u003e"
            ]
          },
          "metadata": {
            "tags": [],
            "needs_background": "light"
          }
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "dBgTdceOT9vz",
        "outputId": "6bc9ec9b-8f30-4f8e-a820-08928f30ecaa",
        "colab": {
          "base_uri": "https://localhost:8080/"
        }
      },
      "source": [
        "# Inference on batch with single example.\n",
        "logits, = model.apply({'params': params}, pp(d['image'], resolution).numpy()[None], train=False)"
      ],
      "execution_count": 30,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "INFO:absl:Starting the local TPU driver.\n",
            "INFO:absl:Unable to initialize backend 'tpu_driver': Not found: Unable to find driver in registry given worker: local://\n",
            "INFO:absl:Unable to initialize backend 'tpu': Invalid argument: TpuPlatform is not available.\n"
          ],
          "name": "stderr"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 382
        },
        "id": "QIQM--g4VwY9",
        "outputId": "57cb35e7-c9e3-45f2-8ead-38a21f6646d4"
      },
      "source": [
        "# Plot logits (you can use tf.nn.softmax() to show probabilities instead).\n",
        "plt.figure(figsize=(10, 4))\n",
        "plt.bar(list(map(ds_info.features['label'].int2str, range(len(logits)))), logits)\n",
        "plt.xticks(rotation=90);"
      ],
      "execution_count": 31,
      "outputs": [
        {
          "output_type": "display_data",
          "data": {
            "image/png": "[encoded data, 35824 characters]\n",
            "text/plain": [
              "\u003cFigure size 720x288 with 1 Axes\u003e"
            ]
          },
          "metadata": {
            "tags": [],
            "needs_background": "light"
          }
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "1nMyWmDycpAo"
      },
      "source": [
        "#### Using `timm`\n",
        "\n",
        "If you know PyTorch, you're probably already familiar with `timm`.\n",
        "\n",
        "If not yet - it's your lucky day! Please check out their docs here:\n",
        "\n",
        "https://rwightman.github.io/pytorch-image-models/"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "yV6Sr1TdXUMp",
        "outputId": "5afc4a13-6c7e-4a01-cef0-27edf08a2a59"
      },
      "source": [
        "# Checkpoints can also be loaded directly into timm...\n",
        "!pip install timm\n",
        "import timm\n",
        "import torch"
      ],
      "execution_count": 32,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Collecting timm\n",
            "  Downloading timm-0.4.12-py3-none-any.whl (376 kB)\n",
            "\u001b[?25l\r\u001b[K     |▉                               | 10 kB 28.4 MB/s eta 0:00:01\r\u001b[K     |█▊                              | 20 kB 32.1 MB/s eta 0:00:01\r\u001b[K     |██▋                             | 30 kB 14.0 MB/s eta 0:00:01\r\u001b[K     |███▌                            | 40 kB 10.7 MB/s eta 0:00:01\r\u001b[K     |████▍                           | 51 kB 7.5 MB/s eta 0:00:01\r\u001b[K     |█████▏                          | 61 kB 7.9 MB/s eta 0:00:01\r\u001b[K     |██████                          | 71 kB 7.4 MB/s eta 0:00:01\r\u001b[K     |███████                         | 81 kB 8.3 MB/s eta 0:00:01\r\u001b[K     |███████▉                        | 92 kB 6.4 MB/s eta 0:00:01\r\u001b[K     |████████▊                       | 102 kB 6.9 MB/s eta 0:00:01\r\u001b[K     |█████████▋                      | 112 kB 6.9 MB/s eta 0:00:01\r\u001b[K     |██████████▍                     | 122 kB 6.9 MB/s eta 0:00:01\r\u001b[K     |███████████▎                    | 133 kB 6.9 MB/s eta 0:00:01\r\u001b[K     |████████████▏                   | 143 kB 6.9 MB/s eta 0:00:01\r\u001b[K     |█████████████                   | 153 kB 6.9 MB/s eta 0:00:01\r\u001b[K     |██████████████                  | 163 kB 6.9 MB/s eta 0:00:01\r\u001b[K     |██████████████▊                 | 174 kB 6.9 MB/s eta 0:00:01\r\u001b[K     |███████████████▋                | 184 kB 6.9 MB/s eta 0:00:01\r\u001b[K     |████████████████▌               | 194 kB 6.9 MB/s eta 0:00:01\r\u001b[K     |█████████████████▍              | 204 kB 6.9 MB/s eta 0:00:01\r\u001b[K     |██████████████████▎             | 215 kB 6.9 MB/s eta 0:00:01\r\u001b[K     |███████████████████▏            | 225 kB 6.9 MB/s eta 0:00:01\r\u001b[K     |████████████████████            | 235 kB 6.9 MB/s eta 0:00:01\r\u001b[K     |████████████████████▉           | 245 kB 6.9 MB/s eta 0:00:01\r\u001b[K     |█████████████████████▊          | 256 kB 6.9 MB/s eta 0:00:01\r\u001b[K     |██████████████████████▋         | 266 kB 6.9 MB/s eta 0:00:01\r\u001b[K     |███████████████████████▌        | 276 kB 6.9 MB/s eta 0:00:01\r\u001b[K     |████████████████████████▍       | 286 kB 6.9 MB/s eta 0:00:01\r\u001b[K     |█████████████████████████▏      | 296 kB 6.9 MB/s eta 0:00:01\r\u001b[K     |██████████████████████████      | 307 kB 6.9 MB/s eta 0:00:01\r\u001b[K     |███████████████████████████     | 317 kB 6.9 MB/s eta 0:00:01\r\u001b[K     |███████████████████████████▉    | 327 kB 6.9 MB/s eta 0:00:01\r\u001b[K     |████████████████████████████▊   | 337 kB 6.9 MB/s eta 0:00:01\r\u001b[K     |█████████████████████████████▌  | 348 kB 6.9 MB/s eta 0:00:01\r\u001b[K     |██████████████████████████████▍ | 358 kB 6.9 MB/s eta 0:00:01\r\u001b[K     |███████████████████████████████▎| 368 kB 6.9 MB/s eta 0:00:01\r\u001b[K     |████████████████████████████████| 376 kB 6.9 MB/s \n",
            "\u001b[?25hRequirement already satisfied: torch\u003e=1.4 in /usr/local/lib/python3.7/dist-packages (from timm) (1.9.0+cu102)\n",
            "Requirement already satisfied: torchvision in /usr/local/lib/python3.7/dist-packages (from timm) (0.10.0+cu102)\n",
            "Requirement already satisfied: typing-extensions in /usr/local/lib/python3.7/dist-packages (from torch\u003e=1.4-\u003etimm) (3.7.4.3)\n",
            "Requirement already satisfied: numpy in /usr/local/lib/python3.7/dist-packages (from torchvision-\u003etimm) (1.19.5)\n",
            "Requirement already satisfied: pillow\u003e=5.3.0 in /usr/local/lib/python3.7/dist-packages (from torchvision-\u003etimm) (7.1.2)\n",
            "Installing collected packages: timm\n",
            "Successfully installed timm-0.4.12\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "gwV4IvfYXioI",
        "outputId": "ca3651e1-b892-4218-d59e-bca10faaa7b5"
      },
      "source": [
        "# For available model names, see here:\n",
        "# https://github.com/rwightman/pytorch-image-models/blob/master/timm/models/vision_transformer.py\n",
        "# https://github.com/rwightman/pytorch-image-models/blob/master/timm/models/vision_transformer_hybrid.py\n",
        "timm_model = timm.create_model(\n",
        "    'vit_small_r26_s32_384', num_classes=ds_info.features['label'].num_classes)\n",
        "\n",
        "# Non-default checkpoints need to be loaded from local files.\n",
        "if not tf.io.gfile.exists(f'{filename}.npz'):\n",
        "  tf.io.gfile.copy(f'gs://vit_models/augreg/{filename}.npz', f'{filename}.npz')\n",
        "timm.models.load_checkpoint(timm_model, f'{filename}.npz')"
      ],
      "execution_count": 33,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "/usr/local/lib/python3.7/dist-packages/torch/nn/functional.py:718: UserWarning: Named tensors and all their associated APIs are an experimental feature and subject to change. Please do not use them for anything important until they are released as stable. (Triggered internally at  /pytorch/c10/core/TensorImpl.h:1156.)\n",
            "  return torch.max_pool2d(input, kernel_size, stride, padding, dilation, ceil_mode)\n"
          ],
          "name": "stderr"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "gd6NVWNLYOSI"
      },
      "source": [
        "def pp_torch(img, sz):\n",
        "  \"\"\"Simple image preprocessing for PyTorch.\"\"\"\n",
        "  img = pp(img, sz)\n",
        "  img = img.numpy().transpose([2, 0, 1])  # PyTorch expects NCHW format.\n",
        "  return torch.tensor(img[None])\n",
        "\n",
        "with torch.no_grad():\n",
        "  logits, = timm_model(pp_torch(d['image'], resolution)).detach().numpy()"
      ],
      "execution_count": 34,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 382
        },
        "id": "vA-mEJofeHxd",
        "outputId": "883ea202-af9c-4be9-b80d-ff207b7208e8"
      },
      "source": [
        "# Same results as above (since we loaded the same checkpoint).\n",
        "plt.figure(figsize=(10, 4))\n",
        "plt.bar(list(map(ds_info.features['label'].int2str, range(len(logits)))), logits)\n",
        "plt.xticks(rotation=90);"
      ],
      "execution_count": 35,
      "outputs": [
        {
          "output_type": "display_data",
          "data": {
            "image/png": "[encoded data, 35640 characters]\n",
            "text/plain": [
              "\u003cFigure size 720x288 with 1 Axes\u003e"
            ]
          },
          "metadata": {
            "tags": [],
            "needs_background": "light"
          }
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "iAruT3YOxqB6"
      },
      "source": [
        "### Fine-tune\n",
        "\n",
        "You want to be connected to a TPU or GPU runtime for fine-tuning.\n",
        "\n",
        "Note that here we're just calling into the code. For more details see the\n",
        "annotated Colab\n",
        "\n",
        "https://colab.research.google.com/github/google-research/vision_transformer/blob/linen/vit_jax.ipynb\n",
        "\n",
        "Also note that Colab GPUs and TPUs are not very powerful. To run this code on\n",
        "more powerful machines, see:\n",
        "\n",
        "https://github.com/google-research/vision_transformer/#running-on-cloud\n",
        "\n",
        "In particular, note that due to the Colab \"TPU Node\" setup, transfering data to\n",
        "the TPUs is realtively slow (for example the smallest `R+Ti/16` model trains\n",
        "faster on a single GPU than on 8 TPUs...)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "TiHEFXlBbZiV"
      },
      "source": [
        "#### TensorBoard"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Fu6uaCGDXLmr"
      },
      "source": [
        "# Launch tensorboard before training - maybe click \"reload\" during training.\n",
        "%load_ext tensorboard\n",
        "%tensorboard --logdir=./workdirs"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "KYU9EdV2bWV7"
      },
      "source": [
        "#### From tfds"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 35
        },
        "id": "Z1Ee2-GzZ9sl",
        "outputId": "bb6c3056-1e2d-4f76-f490-29daa38c30c8"
      },
      "source": [
        "# Create a new temporary workdir.\n",
        "workdir = f'./workdirs/{int(time.time())}'\n",
        "workdir"
      ],
      "execution_count": 37,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "application/vnd.google.colaboratory.intrinsic+json": {
              "type": "string"
            },
            "text/plain": [
              "'./workdirs/1627539136'"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 37
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "SB73ZFGCxsXo"
      },
      "source": [
        "# Get config for specified model.\n",
        "\n",
        "# Note that we can specify simply the model name (in which case the recommended\n",
        "# checkpoint for that model is taken), or it can be specified by its full\n",
        "# name.\n",
        "config = augreg_config.get_config('R_Ti_16')\n",
        "\n",
        "# A very small tfds dataset that only has a \"train\" split. We use this single\n",
        "# split both for training \u0026 evaluation by splitting it further into 90%/10%.\n",
        "config.dataset = 'tf_flowers'\n",
        "config.pp.train = 'train[:90%]'\n",
        "config.pp.test = 'train[90%:]'\n",
        "# tf_flowers only has 3670 images - so the 10% evaluation split will contain\n",
        "# 360 images. We specify batch_eval=120 so we evaluate on all but 7 of those\n",
        "# images (remainder is dropped).\n",
        "config.batch_eval = 120\n",
        "\n",
        "# Some more parameters that you will often want to set manually.\n",
        "# For example for VTAB we used steps={500, 2500} and lr={.001, .003, .01, .03}\n",
        "config.base_lr = 0.01\n",
        "config.shuffle_buffer = 1000\n",
        "config.total_steps = 100\n",
        "config.warmup_steps = 10\n",
        "config.accum_steps = 0  # Not needed with R+Ti/16 model.\n",
        "config.pp['crop'] = 224"
      ],
      "execution_count": 38,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "ldKZFMnmYjY_",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 1000,
          "referenced_widgets": [
            "5ed291cf3157447684e4c732e15713b8",
            "06ea24f8073b4185965fb4b1494e0f19",
            "6b6f717b6613435ca113a68a1e1f4126",
            "a9deea48888c423cbb5b3e310b106147",
            "e2388f624f314ae68d3c6a84127ed400",
            "d51511f4aa504dc2a56c089a167b1c6a",
            "05ec97fde7084f5c9cf54b707bf0edf1",
            "61ee19bf89fb4955bda1e759dfd99cb7"
          ]
        },
        "outputId": "d6fe8e95-83c1-4c55-e96f-1a0b57b7d84a"
      },
      "source": [
        "# Call main training loop. See repository and above Colab for details.\n",
        "state = train.train_and_evaluate(config, workdir)"
      ],
      "execution_count": 39,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "INFO:absl:Load pre-computed DatasetInfo (eg: splits, num examples,...) from GCS: tf_flowers/3.0.1\n",
            "INFO:absl:Load dataset info from /tmp/tmprx4ogu7ltfds\n",
            "INFO:absl:Reading dataset from tfds \"tf_flowers\"\n",
            "INFO:absl:Load pre-computed DatasetInfo (eg: splits, num examples,...) from GCS: tf_flowers/3.0.1\n",
            "INFO:absl:Load dataset info from /tmp/tmpd_sq954qtfds\n",
            "INFO:absl:Generating dataset tf_flowers (/root/tensorflow_datasets/tf_flowers/3.0.1)\n",
            "WARNING:absl:Dataset tf_flowers is hosted on GCS. It will automatically be downloaded to your\n",
            "local data directory. If you'd instead prefer to read directly from our public\n",
            "GCS bucket (recommended if you're running on GCP), you can instead pass\n",
            "`try_gcs=True` to `tfds.load` or set `data_dir=gs://tfds-data/datasets`.\n",
            "\n"
          ],
          "name": "stderr"
        },
        {
          "output_type": "stream",
          "text": [
            "\u001b[1mDownloading and preparing dataset tf_flowers/3.0.1 (download: 218.21 MiB, generated: 221.83 MiB, total: 440.05 MiB) to /root/tensorflow_datasets/tf_flowers/3.0.1...\u001b[0m\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "display_data",
          "data": {
            "application/vnd.jupyter.widget-view+json": {
              "model_id": "5ed291cf3157447684e4c732e15713b8",
              "version_minor": 0,
              "version_major": 2
            },
            "text/plain": [
              "HBox(children=(FloatProgress(value=0.0, description='Dl Completed...', max=5.0, style=ProgressStyle(descriptio…"
            ]
          },
          "metadata": {
            "tags": []
          }
        },
        {
          "output_type": "stream",
          "text": [
            "INFO:absl:Load dataset info from /root/tensorflow_datasets/tf_flowers/3.0.1.incompleteXSDU19\n",
            "INFO:absl:Constructing tf.data.Dataset for split train[:90%], from /root/tensorflow_datasets/tf_flowers/3.0.1\n",
            "INFO:absl:Load dataset info from /root/tensorflow_datasets/tf_flowers/3.0.1\n"
          ],
          "name": "stderr"
        },
        {
          "output_type": "stream",
          "text": [
            "\n",
            "\n",
            "\u001b[1mDataset tf_flowers downloaded and prepared to /root/tensorflow_datasets/tf_flowers/3.0.1. Subsequent calls will reuse this data.\u001b[0m\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "stream",
          "text": [
            "INFO:absl:Load dataset info from /root/tensorflow_datasets/tf_flowers/3.0.1\n",
            "INFO:absl:Reusing dataset tf_flowers (/root/tensorflow_datasets/tf_flowers/3.0.1)\n",
            "INFO:absl:Constructing tf.data.Dataset for split train[90%:], from /root/tensorflow_datasets/tf_flowers/3.0.1\n",
            "INFO:absl:Load dataset info from /root/tensorflow_datasets/tf_flowers/3.0.1\n",
            "INFO:absl:\u003cPrefetchDataset shapes: {image: (1, 512, 224, 224, 3), label: (1, 512, 5)}, types: {image: tf.float32, label: tf.float32}\u003e\n",
            "INFO:absl:\u003cPrefetchDataset shapes: {image: (1, 120, 224, 224, 3), label: (1, 120, 5)}, types: {image: tf.float32, label: tf.float32}\u003e\n",
            "INFO:absl:Selected fillename=\"R_Ti_16-i21k-300ep-lr_0.001-aug_none-wd_0.03-do_0.0-sd_0.0\" for \"R_Ti_16\" with final_val=0.390\n",
            "INFO:absl:Found no checkpoint file at ./workdirs/1627539136\n",
            "INFO:absl:Will start/continue training at initial_step=1\n",
            "/usr/local/lib/python3.7/dist-packages/jax/lib/xla_bridge.py:387: UserWarning: jax.host_count has been renamed to jax.process_count. This alias will eventually be removed; please update your code.\n",
            "  \"jax.host_count has been renamed to jax.process_count. This alias \"\n",
            "/usr/local/lib/python3.7/dist-packages/jax/lib/xla_bridge.py:374: UserWarning: jax.host_id has been renamed to jax.process_index. This alias will eventually be removed; please update your code.\n",
            "  \"jax.host_id has been renamed to jax.process_index. This alias \"\n",
            "INFO:absl:Hyperparameters: {'accum_steps': 0, 'base_lr': 0.01, 'batch': 512, 'batch_eval': 120, 'checkpoint_every': 1000, 'dataset': 'tf_flowers', 'decay_type': 'cosine', 'eval_every': 100, 'grad_norm_clip': 1.0, 'model': {'classifier': 'token', 'hidden_size': 192, 'name': 'R+ViT-Ti_16', 'patches': {'size': (8, 8)}, 'representation_size': None, 'resnet': {'num_layers': (), 'width_factor': 1}, 'transformer': {'attention_dropout_rate': 0.0, 'dropout_rate': 0.0, 'mlp_dim': 768, 'num_heads': 3, 'num_layers': 12}}, 'model_or_filename': 'R_Ti_16', 'optim_dtype': 'bfloat16', 'pp': {'crop': 224, 'resize': 448, 'test': 'train[90%:]', 'train': 'train[:90%]'}, 'prefetch': 2, 'pretrained_dir': 'gs://vit_models/augreg', 'progress_every': 10, 'shuffle_buffer': 1000, 'tfds_data_dir': None, 'tfds_manual_dir': None, 'total_steps': 100, 'trainer': 'train', 'warmup_steps': 10}\n",
            "INFO:absl:Starting training loop; initial compile can take a while...\n",
            "/usr/local/lib/python3.7/dist-packages/jax/_src/profiler.py:167: UserWarning: StepTraceContext has been renamed to StepTraceAnnotation. This alias will eventually be removed; please update your code.\n",
            "  \"StepTraceContext has been renamed to StepTraceAnnotation. This alias \"\n",
            "INFO:absl:First step took 40.1 seconds.\n",
            "INFO:absl:[10] img_sec_core_train=497.200154, train_loss=0.522652\n",
            "INFO:absl:Step: 10/100 10.0%, img/sec/core: 497.2, ETA: 0.02h\n",
            "INFO:absl:Created artifact [10] Profile of type ArtifactType.URL and value None.\n",
            "INFO:absl:[20] img_sec_core_train=348.195610, train_loss=0.200892\n",
            "INFO:absl:Step: 20/100 20.0%, img/sec/core: 348.2, ETA: 0.03h\n",
            "INFO:absl:[30] img_sec_core_train=408.326154, train_loss=0.223066\n",
            "INFO:absl:Step: 30/100 30.0%, img/sec/core: 408.3, ETA: 0.02h\n",
            "INFO:absl:[40] img_sec_core_train=410.988480, train_loss=0.210076\n",
            "INFO:absl:Step: 40/100 40.0%, img/sec/core: 411.0, ETA: 0.02h\n",
            "INFO:absl:Setting work unit notes: 0.8 steps/s, 49.0% @49, ETA: 1m\n",
            "INFO:absl:[49] steps_per_sec=0.797867\n",
            "INFO:absl:[50] img_sec_core_train=410.102459, train_loss=0.255746\n",
            "INFO:absl:Step: 50/100 50.0%, img/sec/core: 410.1, ETA: 0.02h\n",
            "INFO:absl:[60] img_sec_core_train=389.702776, train_loss=0.172555\n",
            "INFO:absl:Step: 60/100 60.0%, img/sec/core: 389.7, ETA: 0.01h\n",
            "INFO:absl:[70] img_sec_core_train=383.566635, train_loss=0.175968\n",
            "INFO:absl:Step: 70/100 70.0%, img/sec/core: 383.6, ETA: 0.01h\n",
            "INFO:absl:[80] img_sec_core_train=407.022965, train_loss=0.147852\n",
            "INFO:absl:Step: 80/100 80.0%, img/sec/core: 407.0, ETA: 0.01h\n",
            "INFO:absl:[90] img_sec_core_train=406.781729, train_loss=0.137439\n",
            "INFO:absl:Step: 90/100 90.0%, img/sec/core: 406.8, ETA: 0.00h\n",
            "INFO:absl:Setting work unit notes: 0.8 steps/s, 96.0% @96, ETA: 0m\n",
            "INFO:absl:[96] steps_per_sec=0.776280\n",
            "INFO:absl:Setting work unit notes: 0.8 steps/s, 100.0% @100, ETA: 0m\n",
            "INFO:absl:[100] steps_per_sec=0.805336\n",
            "INFO:absl:[100] img_sec_core_train=407.049767, train_loss=0.143796\n",
            "INFO:absl:Step: 100/100 100.0%, img/sec/core: 407.0, ETA: 0.00h\n",
            "INFO:absl:Step: 100 Learning rate: 0.0000000, Test accuracy: 0.95000, img/sec/core: 41.9\n",
            "INFO:absl:[100] accuracy_test=0.9500000476837158, img_sec_core_test=41.864635, lr=0.000000\n",
            "INFO:absl:Saving checkpoint at step: 100\n",
            "INFO:absl:Saved checkpoint at workdirs/1627539136/checkpoint_100\n",
            "INFO:absl:Stored checkpoint at step 100 to \"workdirs/1627539136/checkpoint_100\"\n"
          ],
          "name": "stderr"
        },
        {
          "output_type": "stream",
          "text": [
            "workdirs/1627539136/checkpoint_100\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "JOtCfWFOfVFq"
      },
      "source": [
        "#### From JPG files\n",
        "\n",
        "The codebase supports training directly form JPG files on the local filesystem\n",
        "instead of `tfds` datasets. Note that the throughput is somewhat reduced, but\n",
        "that only is noticeable for very small models.\n",
        "\n",
        "The main advantage of `tfds` datasets is that they are versioned and available\n",
        "globally."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "5VxlDlnxKAa3"
      },
      "source": [
        "base = '.'  # Store data on VM (ephemeral)."
      ],
      "execution_count": 40,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "yr6JY8OBmKkj"
      },
      "source": [
        "# Uncomment below lines if you want to download \u0026 persist files in your Google\n",
        "# Drive instead. Note that Colab VMs are reset (i.e. files are deleted) after\n",
        "# some time of inactivity. Storing data to Google Drive guarantees that it is\n",
        "# still available next time you connect from a new VM.\n",
        "\n",
        "# Note that this is significantly slower than reading from the VMs locally\n",
        "# attached file system!\n",
        "\n",
        "# from google.colab import drive\n",
        "# drive.mount('/gdrive')\n",
        "# base = '/gdrive/My Drive/vision_transformer_images'"
      ],
      "execution_count": 41,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "VZJdKYv3mLqd",
        "outputId": "552207d7-cfbe-4a1f-e140-bbbbfd0da31c"
      },
      "source": [
        "# Download some dataset \u0026 unzip.\n",
        "! rm -rf $base/flower_photos; mkdir -p $base\n",
        "! (cd $base \u0026\u0026 curl https://storage.googleapis.com/download.tensorflow.org/example_images/flower_photos.tgz | tar xz)"
      ],
      "execution_count": 42,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "  % Total    % Received % Xferd  Average Speed   Time    Time     Time  Current\n",
            "                                 Dload  Upload   Total   Spent    Left  Speed\n",
            "100  218M  100  218M    0     0  81.8M      0  0:00:02  0:00:02 --:--:-- 81.8M\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "0QRnYHLzmPxL",
        "outputId": "5d464050-9e3f-4ba0-eed6-ce49cfe01276"
      },
      "source": [
        "# Since the default file format of above \"tf_flowers\" dataset is\n",
        "# flower_photos/{class_name}/{filename}.jpg\n",
        "# we first need to split it into a \"train\" (90%) and a \"test\" (10%) set:\n",
        "# flower_photos/train/{class_name}/{filename}.jpg\n",
        "# flower_photos/test/{class_name}/{filename}.jpg\n",
        "\n",
        "def split(base_dir, test_ratio=0.1):\n",
        "  paths = glob.glob(f'{base_dir}/*/*.jpg')\n",
        "  random.shuffle(paths)\n",
        "  counts = dict(test=0, train=0)\n",
        "  for i, path in enumerate(paths):\n",
        "    split = 'test' if i \u003c test_ratio * len(paths) else 'train'\n",
        "    *_, class_name, basename = path.split('/')\n",
        "    dst = f'{base_dir}/{split}/{class_name}/{basename}'\n",
        "    if not os.path.isdir(os.path.dirname(dst)):\n",
        "      os.makedirs(os.path.dirname(dst))\n",
        "    shutil.move(path, dst)\n",
        "    counts[split] += 1\n",
        "  print(f'Moved {counts[\"train\"]:,} train and {counts[\"test\"]:,} test images.')\n",
        "\n",
        "split(f'{base}/flower_photos')"
      ],
      "execution_count": 43,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Moved 3,303 train and 367 test images.\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 35
        },
        "id": "1DKwlCXTKJjJ",
        "outputId": "83f39dfb-e764-4bf8-e486-95d26db6e371"
      },
      "source": [
        "# Create a new temporary workdir.\n",
        "workdir = f'./workdirs/{int(time.time())}'\n",
        "workdir"
      ],
      "execution_count": 44,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "application/vnd.google.colaboratory.intrinsic+json": {
              "type": "string"
            },
            "text/plain": [
              "'./workdirs/1627539429'"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 44
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "N2ojstZeKJjK"
      },
      "source": [
        "# Read data from directory containing files.\n",
        "# (See cell above for more config settings)\n",
        "config.dataset = f'{base}/flower_photos'"
      ],
      "execution_count": 45,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "L8hljjZlKJjK",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "b9880743-3984-475e-a3f1-7c60b6f904e4"
      },
      "source": [
        "# And fine-tune on images provided\n",
        "opt = train.train_and_evaluate(config, workdir)"
      ],
      "execution_count": 46,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "INFO:absl:Reading dataset from directories \"./flower_photos/train\" and \"./flower_photos/test\"\n",
            "INFO:absl:\u003cPrefetchDataset shapes: {image: (1, 512, 224, 224, 3), label: (1, 512, 5)}, types: {image: tf.float32, label: tf.float32}\u003e\n",
            "INFO:absl:\u003cPrefetchDataset shapes: {image: (1, 120, 224, 224, 3), label: (1, 120, 5)}, types: {image: tf.float32, label: tf.float32}\u003e\n",
            "INFO:absl:Selected fillename=\"R_Ti_16-i21k-300ep-lr_0.001-aug_none-wd_0.03-do_0.0-sd_0.0\" for \"R_Ti_16\" with final_val=0.390\n",
            "INFO:absl:Found no checkpoint file at ./workdirs/1627539429\n",
            "INFO:absl:Will start/continue training at initial_step=1\n",
            "/usr/local/lib/python3.7/dist-packages/jax/lib/xla_bridge.py:387: UserWarning: jax.host_count has been renamed to jax.process_count. This alias will eventually be removed; please update your code.\n",
            "  \"jax.host_count has been renamed to jax.process_count. This alias \"\n",
            "/usr/local/lib/python3.7/dist-packages/jax/lib/xla_bridge.py:374: UserWarning: jax.host_id has been renamed to jax.process_index. This alias will eventually be removed; please update your code.\n",
            "  \"jax.host_id has been renamed to jax.process_index. This alias \"\n",
            "INFO:absl:Hyperparameters: {'accum_steps': 0, 'base_lr': 0.01, 'batch': 512, 'batch_eval': 120, 'checkpoint_every': 1000, 'dataset': './flower_photos', 'decay_type': 'cosine', 'eval_every': 100, 'grad_norm_clip': 1.0, 'model': {'classifier': 'token', 'hidden_size': 192, 'name': 'R+ViT-Ti_16', 'patches': {'size': (8, 8)}, 'representation_size': None, 'resnet': {'num_layers': (), 'width_factor': 1}, 'transformer': {'attention_dropout_rate': 0.0, 'dropout_rate': 0.0, 'mlp_dim': 768, 'num_heads': 3, 'num_layers': 12}}, 'model_or_filename': 'R_Ti_16', 'optim_dtype': 'bfloat16', 'pp': {'crop': 224, 'resize': 448, 'test': 'train[90%:]', 'train': 'train[:90%]'}, 'prefetch': 2, 'pretrained_dir': 'gs://vit_models/augreg', 'progress_every': 10, 'shuffle_buffer': 1000, 'tfds_data_dir': None, 'tfds_manual_dir': None, 'total_steps': 100, 'trainer': 'train', 'warmup_steps': 10}\n",
            "INFO:absl:Starting training loop; initial compile can take a while...\n",
            "/usr/local/lib/python3.7/dist-packages/jax/_src/profiler.py:167: UserWarning: StepTraceContext has been renamed to StepTraceAnnotation. This alias will eventually be removed; please update your code.\n",
            "  \"StepTraceContext has been renamed to StepTraceAnnotation. This alias \"\n",
            "INFO:absl:First step took 26.7 seconds.\n",
            "INFO:absl:[10] img_sec_core_train=504.439041, train_loss=0.517574\n",
            "INFO:absl:Step: 10/100 10.0%, img/sec/core: 504.4, ETA: 0.02h\n",
            "INFO:absl:Created artifact [10] Profile of type ArtifactType.URL and value None.\n",
            "INFO:absl:[20] img_sec_core_train=364.853150, train_loss=0.232779\n",
            "INFO:absl:Step: 20/100 20.0%, img/sec/core: 364.9, ETA: 0.03h\n",
            "INFO:absl:[30] img_sec_core_train=358.281401, train_loss=0.251437\n",
            "INFO:absl:Step: 30/100 30.0%, img/sec/core: 358.3, ETA: 0.02h\n",
            "INFO:absl:[40] img_sec_core_train=359.572849, train_loss=0.196024\n",
            "INFO:absl:Step: 40/100 40.0%, img/sec/core: 359.6, ETA: 0.02h\n",
            "INFO:absl:Setting work unit notes: 0.8 steps/s, 47.0% @47, ETA: 1m\n",
            "INFO:absl:[47] steps_per_sec=0.753239\n",
            "INFO:absl:[50] img_sec_core_train=381.908006, train_loss=0.194640\n",
            "INFO:absl:Step: 50/100 50.0%, img/sec/core: 381.9, ETA: 0.02h\n",
            "INFO:absl:[60] img_sec_core_train=380.494434, train_loss=0.135086\n",
            "INFO:absl:Step: 60/100 60.0%, img/sec/core: 380.5, ETA: 0.01h\n",
            "INFO:absl:[70] img_sec_core_train=376.330891, train_loss=0.181564\n",
            "INFO:absl:Step: 70/100 70.0%, img/sec/core: 376.3, ETA: 0.01h\n",
            "INFO:absl:[80] img_sec_core_train=354.874495, train_loss=0.100307\n",
            "INFO:absl:Step: 80/100 80.0%, img/sec/core: 354.9, ETA: 0.01h\n",
            "INFO:absl:[90] img_sec_core_train=367.600830, train_loss=0.128104\n",
            "INFO:absl:Step: 90/100 90.0%, img/sec/core: 367.6, ETA: 0.00h\n",
            "INFO:absl:Setting work unit notes: 0.7 steps/s, 91.0% @91, ETA: 0m\n",
            "INFO:absl:[91] steps_per_sec=0.721651\n",
            "INFO:absl:Setting work unit notes: 0.7 steps/s, 100.0% @100, ETA: 0m\n",
            "INFO:absl:[100] steps_per_sec=0.745245\n",
            "INFO:absl:[100] img_sec_core_train=376.285043, train_loss=0.141850\n",
            "INFO:absl:Step: 100/100 100.0%, img/sec/core: 376.3, ETA: 0.00h\n",
            "INFO:absl:Step: 100 Learning rate: 0.0000000, Test accuracy: 0.95278, img/sec/core: 49.7\n",
            "INFO:absl:[100] accuracy_test=0.9527778625488281, img_sec_core_test=49.667636, lr=0.000000\n",
            "INFO:absl:Saving checkpoint at step: 100\n",
            "INFO:absl:Saved checkpoint at workdirs/1627539429/checkpoint_100\n",
            "INFO:absl:Stored checkpoint at step 100 to \"workdirs/1627539429/checkpoint_100\"\n"
          ],
          "name": "stderr"
        },
        {
          "output_type": "stream",
          "text": [
            "workdirs/1627539429/checkpoint_100\n"
          ],
          "name": "stdout"
        }
      ]
    }
  ]
}
