Notebooks/Ultimate_Training_Loop.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": [],
      "toc_visible": true,
      "authorship_tag": "ABX9TyMnVJZ99JAmwo1oxhRI3x/t",
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    },
    "accelerator": "GPU",
    "gpuClass": "standard",
    "widgets": {
      "application/vnd.jupyter.widget-state+json": {
        "63dfa8a74fd045ffaac822f9f43e72e9": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HBoxModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HBoxModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HBoxView",
            "box_style": "",
            "children": [
              "IPY_MODEL_8d60b02aabbe47f8acc7cedfc41104f7",
              "IPY_MODEL_3e410451d4b14ba99f13e0f011a98e95",
              "IPY_MODEL_2c4227810a454ccd885fe022920fa60b"
            ],
            "layout": "IPY_MODEL_44fb6bf5859d4991b74636f180e80dcf"
          }
        },
        "8d60b02aabbe47f8acc7cedfc41104f7": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HTMLModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HTMLModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HTMLView",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_9ab1f173c1e849d98370bbb7fd9ca00c",
            "placeholder": "​",
            "style": "IPY_MODEL_093b4063ab614edfa9b3382f0eafa8a2",
            "value": "100%"
          }
        },
        "3e410451d4b14ba99f13e0f011a98e95": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "FloatProgressModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "FloatProgressModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "ProgressView",
            "bar_style": "success",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_d108b428d12f459eb46fec76fb4bed0d",
            "max": 26421880,
            "min": 0,
            "orientation": "horizontal",
            "style": "IPY_MODEL_eb2d57debdb845718418b316f5f9af9c",
            "value": 26421880
          }
        },
        "2c4227810a454ccd885fe022920fa60b": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HTMLModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HTMLModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HTMLView",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_1a48f249f0724d94930d9f799d1393b8",
            "placeholder": "​",
            "style": "IPY_MODEL_d63b69038301453aaf86b94f1825ef9c",
            "value": " 26421880/26421880 [00:03&lt;00:00, 13844165.76it/s]"
          }
        },
        "44fb6bf5859d4991b74636f180e80dcf": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "9ab1f173c1e849d98370bbb7fd9ca00c": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "093b4063ab614edfa9b3382f0eafa8a2": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "DescriptionStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "DescriptionStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "description_width": ""
          }
        },
        "d108b428d12f459eb46fec76fb4bed0d": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "eb2d57debdb845718418b316f5f9af9c": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "ProgressStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "ProgressStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "bar_color": null,
            "description_width": ""
          }
        },
        "1a48f249f0724d94930d9f799d1393b8": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "d63b69038301453aaf86b94f1825ef9c": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "DescriptionStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "DescriptionStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "description_width": ""
          }
        },
        "b8f819130ff249f6b400af84844b265d": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HBoxModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HBoxModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HBoxView",
            "box_style": "",
            "children": [
              "IPY_MODEL_431b8ee6ed1f4f83a650baca18ffb55f",
              "IPY_MODEL_ae6f057646884135a7e79c6b615e14ea",
              "IPY_MODEL_d25ca09757e840419340a3652d8f58ed"
            ],
            "layout": "IPY_MODEL_d0b19436553f41b29eca0b6a6c9b1019"
          }
        },
        "431b8ee6ed1f4f83a650baca18ffb55f": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HTMLModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HTMLModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HTMLView",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_7eab30d1ddd147e3adbdcf374cab6f04",
            "placeholder": "​",
            "style": "IPY_MODEL_7fb9c12460f241c39b58e8d402879033",
            "value": "100%"
          }
        },
        "ae6f057646884135a7e79c6b615e14ea": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "FloatProgressModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "FloatProgressModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "ProgressView",
            "bar_style": "success",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_19a2183d08ae4deea471b6d18734ee19",
            "max": 29515,
            "min": 0,
            "orientation": "horizontal",
            "style": "IPY_MODEL_901710e665684edc96e9e5600c3a026e",
            "value": 29515
          }
        },
        "d25ca09757e840419340a3652d8f58ed": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HTMLModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HTMLModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HTMLView",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_879ce688924641d1bb03cef1cea05245",
            "placeholder": "​",
            "style": "IPY_MODEL_6b79bb7745ec4aac8dc4b26a248a87ba",
            "value": " 29515/29515 [00:00&lt;00:00, 144492.07it/s]"
          }
        },
        "d0b19436553f41b29eca0b6a6c9b1019": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "7eab30d1ddd147e3adbdcf374cab6f04": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "7fb9c12460f241c39b58e8d402879033": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "DescriptionStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "DescriptionStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "description_width": ""
          }
        },
        "19a2183d08ae4deea471b6d18734ee19": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "901710e665684edc96e9e5600c3a026e": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "ProgressStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "ProgressStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "bar_color": null,
            "description_width": ""
          }
        },
        "879ce688924641d1bb03cef1cea05245": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "6b79bb7745ec4aac8dc4b26a248a87ba": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "DescriptionStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "DescriptionStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "description_width": ""
          }
        },
        "1d240e6e82a147019d1bf3e741bd55ba": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HBoxModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HBoxModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HBoxView",
            "box_style": "",
            "children": [
              "IPY_MODEL_67966b3ed0d9411585737d1d35e60ee5",
              "IPY_MODEL_9cdd7f4fab59464e90b606c6dbdfd916",
              "IPY_MODEL_f99791e702b84b1ebbd15e0bd66b1333"
            ],
            "layout": "IPY_MODEL_68670028637a420d90c999a05dd4206c"
          }
        },
        "67966b3ed0d9411585737d1d35e60ee5": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HTMLModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HTMLModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HTMLView",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_5ebf724b7af443f49dd3a2ba954e030a",
            "placeholder": "​",
            "style": "IPY_MODEL_80f2e1a5bd854f88a13db57817296c8c",
            "value": "100%"
          }
        },
        "9cdd7f4fab59464e90b606c6dbdfd916": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "FloatProgressModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "FloatProgressModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "ProgressView",
            "bar_style": "success",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_26765f58c1d44bf192e080cb1c5e7492",
            "max": 4422102,
            "min": 0,
            "orientation": "horizontal",
            "style": "IPY_MODEL_d3046fe79ebf4627a69aebecc412a2ea",
            "value": 4422102
          }
        },
        "f99791e702b84b1ebbd15e0bd66b1333": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HTMLModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HTMLModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HTMLView",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_ef508d90634b44d6add55edabd37399c",
            "placeholder": "​",
            "style": "IPY_MODEL_520cecdef46949f5beb24bb7a33e187f",
            "value": " 4422102/4422102 [00:01&lt;00:00, 4618943.63it/s]"
          }
        },
        "68670028637a420d90c999a05dd4206c": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "5ebf724b7af443f49dd3a2ba954e030a": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "80f2e1a5bd854f88a13db57817296c8c": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "DescriptionStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "DescriptionStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "description_width": ""
          }
        },
        "26765f58c1d44bf192e080cb1c5e7492": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "d3046fe79ebf4627a69aebecc412a2ea": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "ProgressStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "ProgressStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "bar_color": null,
            "description_width": ""
          }
        },
        "ef508d90634b44d6add55edabd37399c": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "520cecdef46949f5beb24bb7a33e187f": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "DescriptionStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "DescriptionStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "description_width": ""
          }
        },
        "80393295b52d4845972d3a454ae432b4": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HBoxModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HBoxModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HBoxView",
            "box_style": "",
            "children": [
              "IPY_MODEL_f982afd1605542fdb0b042802ff7d7ef",
              "IPY_MODEL_6189961fdb89450c852672574ed30d2c",
              "IPY_MODEL_9e203c1650854dc0bcf54ea88023fe18"
            ],
            "layout": "IPY_MODEL_50db8a1f9edf46d28072e739c0ca33cd"
          }
        },
        "f982afd1605542fdb0b042802ff7d7ef": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HTMLModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HTMLModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HTMLView",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_83faa07de0e1441c8cc29d23316319fc",
            "placeholder": "​",
            "style": "IPY_MODEL_77d73a672ede4f4e8b399629021b041a",
            "value": "100%"
          }
        },
        "6189961fdb89450c852672574ed30d2c": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "FloatProgressModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "FloatProgressModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "ProgressView",
            "bar_style": "success",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_d83d7739f4e34486b3259bfcfea67a4b",
            "max": 5148,
            "min": 0,
            "orientation": "horizontal",
            "style": "IPY_MODEL_9907eb9b2da849dd8c9b1a14e427ea5f",
            "value": 5148
          }
        },
        "9e203c1650854dc0bcf54ea88023fe18": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HTMLModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HTMLModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HTMLView",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_720d011a63b14bf7abab20d3f17b0785",
            "placeholder": "​",
            "style": "IPY_MODEL_ff66ca34ec49449181fbee8546f44e58",
            "value": " 5148/5148 [00:00&lt;00:00, 215911.97it/s]"
          }
        },
        "50db8a1f9edf46d28072e739c0ca33cd": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "83faa07de0e1441c8cc29d23316319fc": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "77d73a672ede4f4e8b399629021b041a": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "DescriptionStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "DescriptionStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "description_width": ""
          }
        },
        "d83d7739f4e34486b3259bfcfea67a4b": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "9907eb9b2da849dd8c9b1a14e427ea5f": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "ProgressStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "ProgressStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "bar_color": null,
            "description_width": ""
          }
        },
        "720d011a63b14bf7abab20d3f17b0785": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "ff66ca34ec49449181fbee8546f44e58": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "DescriptionStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "DescriptionStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "description_width": ""
          }
        }
      }
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/Ankur-singh/personal_projects/blob/master/Notebooks/Ultimate_Training_Loop.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "# The Ultimate Training Loop\n",
        "\n",
        "Writing a train loops is one of the most tedious tasks in Deep Learning. \n",
        "- It seems simple, but requires a lot of small changes in every other use-case.\n",
        "- Very error prone\n",
        "- It quickly becomes messy.\n",
        "\n",
        "In this blog, I will discuss the training loop proposed by Jeremy Howard in Part 2 of the 2022 course. I think this is the last training loop that you will ever need. Its infinitely flexible and super easy to use once you get the hang of it. Lets take a look at it. \n",
        "\n",
        "**Disclaimer:** Most of the code snippets and ideas used in this blog are taken from fastai. So, credit for everything good should go to Jeremy Howard, and everything bad or mediocre is my work. The main reason for writing the blog was to consolidate my learning."
      ],
      "metadata": {
        "id": "Ryq90JJL27Or"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Setting Up\n",
        "\n",
        "I will start by writing all the necessary code. Its okay if you don't understand it completely. In fact, even I am exploring. But don't worry, I will try my best to highlight and explain the code snippets that are important to understanding this blog. The main focus of this blog is to get you excited about this training loop and show you \"how to use it\" through examples."
      ],
      "metadata": {
        "id": "JJ8Pis77fLD0"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "import torch\n",
        "import fastcore.all as fc\n",
        "from torch import nn, optim\n",
        "from functools import partial\n",
        "from operator import attrgetter\n",
        "from torch.nn import functional as F"
      ],
      "metadata": {
        "id": "t6fSx1j94FtE"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# Callback Utilities\n",
        "\n",
        "class CancelFitException(Exception): pass\n",
        "class CancelBatchException(Exception): pass\n",
        "class CancelEpochException(Exception): pass\n",
        "\n",
        "class Callback(): order = 0\n",
        "\n",
        "def run_cbs(cbs, method_nm, learn=None):\n",
        "    for cb in sorted(cbs, key=attrgetter('order')):\n",
        "        method = getattr(cb, method_nm, None)\n",
        "        if method is not None: method(learn)"
      ],
      "metadata": {
        "id": "i1cJJX0h4ham"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "In the above cell, we defined three custom exceptions, namely `CancelFitExeption`, `CancelBatchException`, and `CancelEpochException`. We will be raising them to cancel different parts of our training loop. \n",
        "\n",
        "Next, we defined a `Callback` class and `run_cbs` function. The `run_cbs` function is quite interesting. It takes a *list of callbacks*, a *method name*, and a *learner object* as input. \n",
        "\n",
        "We start by sorting our callbacks based on `order` (hence, the `Callback` class has a static variable named `order`). Then, for each callback in the sorted list if the particular method is present, then it will be called with *learner* as an argument. "
      ],
      "metadata": {
        "id": "EvTGY_U5g5_B"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "class with_cbs:\n",
        "    def __init__(self, nm): self.nm = nm\n",
        "    def __call__(self, f):\n",
        "        def _f(o, *args, **kwargs):\n",
        "            try:\n",
        "                o.callback(f'before_{self.nm}')\n",
        "                f(o, *args, **kwargs)\n",
        "                o.callback(f'after_{self.nm}')\n",
        "            except globals()[f'Cancel{self.nm.title()}Exception']: pass\n",
        "            finally: o.callback(f'cleanup_{self.nm}')\n",
        "        return _f"
      ],
      "metadata": {
        "id": "9wlr7B-rBdX5"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "The `with_cbs` class is a super compact way of defining a decorator that returns a new function wrapping `learn.callback(\"before_<context_name>\")` and `learn.callback(\"after_<context_name>)` calls around the main function. Furthermore, its put in a *try-except-finally* block that handles `Cancel<Context_name>Exception` and runs `cleanup_<context_name>` methods in the finally block."
      ],
      "metadata": {
        "id": "CQP3NTKipuWz"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "sRcErNTS2hof"
      },
      "outputs": [],
      "source": [
        "# Main Training Loop\n",
        "\n",
        "class Learner():\n",
        "    def __init__(self, model, dls=(0,), loss_func=F.mse_loss, lr=0.1, cbs=None, opt_func=optim.SGD):\n",
        "        cbs = fc.L(cbs)\n",
        "        fc.store_attr()\n",
        "\n",
        "    @with_cbs('batch')\n",
        "    def _one_batch(self):\n",
        "        self.predict()\n",
        "        self.callback('after_predict')\n",
        "        self.get_loss()\n",
        "        self.callback('after_loss')\n",
        "        if self.training:\n",
        "            self.backward()\n",
        "            self.callback('after_backward')\n",
        "            self.step()\n",
        "            self.callback('after_step')\n",
        "            self.zero_grad()\n",
        "\n",
        "    @with_cbs('epoch')\n",
        "    def _one_epoch(self):\n",
        "        for self.iter,self.batch in enumerate(self.dl): self._one_batch()\n",
        "\n",
        "    def one_epoch(self, training):\n",
        "        self.model.train(training)\n",
        "        self.dl = self.dls.train if training else self.dls.valid\n",
        "        self._one_epoch()\n",
        "\n",
        "    @with_cbs('fit')\n",
        "    def _fit(self, train, valid):\n",
        "        for self.epoch in self.epochs:\n",
        "            if train: self.one_epoch(True)\n",
        "            if valid: torch.no_grad()(self.one_epoch)(False)\n",
        "\n",
        "    def fit(self, n_epochs=1, train=True, valid=True, cbs=None, lr=None):\n",
        "        cbs = fc.L(cbs)\n",
        "        # `add_cb` and `rm_cb` were added in lesson 18\n",
        "        for cb in cbs: self.cbs.append(cb)\n",
        "        try:\n",
        "            self.n_epochs = n_epochs\n",
        "            self.epochs = range(n_epochs)\n",
        "            if lr is None: lr = self.lr\n",
        "            if self.opt_func: self.opt = self.opt_func(self.model.parameters(), lr)\n",
        "            self._fit(train, valid)\n",
        "        finally:\n",
        "            for cb in cbs: self.cbs.remove(cb)\n",
        "\n",
        "    def __getattr__(self, name):\n",
        "        if name in ('predict','get_loss','backward','step','zero_grad'): return partial(self.callback, name)\n",
        "        raise AttributeError(name)\n",
        "\n",
        "    def callback(self, method_nm): run_cbs(self.cbs, method_nm, self)\n",
        "    \n",
        "    @property\n",
        "    def training(self): return self.model.training"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        " The above code might look daunting so here is a simplified version of it.\n",
        "\n",
        "```python\n",
        "def one_epoch(self, train):\n",
        "    self.model.train(train)\n",
        "    self.dl = self.dls.train if train else self.dls.valid\n",
        "    self.callback('before_epoch')             # Epoch start\n",
        "    for self.iter,self.batch in enumerate(self.dl):\n",
        "        self.callback('before_batch')         # Batch Start\n",
        "        self.predict()\n",
        "        self.callback('after_predict')        # After populating self.preds\n",
        "        self.get_loss()\n",
        "        self.callback('after_loss')           # After calculating self.loss\n",
        "        if self.training:\n",
        "            self.backward()\n",
        "            self.callback('after_backward')  # After calling backward\n",
        "            self.step()\n",
        "            self.callback('after_step')      # After performing weight updates\n",
        "            self.zero_grad()\n",
        "        self.callback('after_batch')         # Batch End\n",
        "    self.callback('after_epoch')             # Epoch End\n",
        "\n",
        "\n",
        "def fit(self, n_epochs=1, train=True, valid=True, cbs=None, lr=None)\n",
        "    self.n_epochs = n_epochs\n",
        "    self.epochs = range(n_epochs)\n",
        "    if lr is None: lr = self.lr\n",
        "    if self.opt_func: \n",
        "        self.opt = self.opt_func(self.model.parameters(), lr)\n",
        "    self.callback('before_fit')              # Fit Start\n",
        "    for self.epoch in self.epochs:\n",
        "        if train: self.one_epoch(True)\n",
        "        if valid: torch.no_grad()(self.one_epoch)(False)\n",
        "    self.callback('after_fit')               # Fit End\n",
        "\n",
        "\n",
        "```"
      ],
      "metadata": {
        "id": "Pe3yXJ9f5yq_"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "There are two interesting things in this training loop:\n",
        "1. It doesn't have implementation for `predict`, `get_loss`, `backward`, `step`, and `zero_grad` methods. So, one can easily inherit this learner class and implement these methods based on one's use-case. We will discuss this in more detail in a minute. \n",
        "2. It supports callbacks at various stages of the training process like `before_fit`, `after_fit`, `before_epoch`, so on and so forth. This will allow us to run arbitrary code at any of these various points in the training loop. I have highlighted all these points in the above code snippet.\n",
        "\n",
        "Both these features might look very simple at first, but trust me its very powerful. I will be discussing both the features and their applications in detail so hold on tight. Its going to be fun!\n",
        "\n",
        "But before we can test this loop, we will need some data and a model, right? I plan to keep it simple so that we can focus on the topic at hand. When it comes to keeping things simple, whats better than *fashion_mnist* and basic CNN?"
      ],
      "metadata": {
        "id": "7oqu7rgxoSLd"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "from torchvision import transforms as T\n",
        "from torchvision.datasets import FashionMNIST\n",
        "from torch.utils.data import DataLoader"
      ],
      "metadata": {
        "id": "kBA2EMCp40mn"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "BS = 64\n",
        "\n",
        "tfms = T.Compose([T.ToTensor(), T.Normalize((0.5,), (0.5,))])\n",
        "\n",
        "train_ds = FashionMNIST('./data', train=True, transform=tfms, download=True)\n",
        "test_ds = FashionMNIST('./data', train=False, transform=tfms, download=True)\n",
        "\n",
        "train_dl = DataLoader(train_ds, batch_size=BS, shuffle=True)\n",
        "test_dl = DataLoader(test_ds, batch_size=BS, shuffle=False)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 437,
          "referenced_widgets": [
            "63dfa8a74fd045ffaac822f9f43e72e9",
            "8d60b02aabbe47f8acc7cedfc41104f7",
            "3e410451d4b14ba99f13e0f011a98e95",
            "2c4227810a454ccd885fe022920fa60b",
            "44fb6bf5859d4991b74636f180e80dcf",
            "9ab1f173c1e849d98370bbb7fd9ca00c",
            "093b4063ab614edfa9b3382f0eafa8a2",
            "d108b428d12f459eb46fec76fb4bed0d",
            "eb2d57debdb845718418b316f5f9af9c",
            "1a48f249f0724d94930d9f799d1393b8",
            "d63b69038301453aaf86b94f1825ef9c",
            "b8f819130ff249f6b400af84844b265d",
            "431b8ee6ed1f4f83a650baca18ffb55f",
            "ae6f057646884135a7e79c6b615e14ea",
            "d25ca09757e840419340a3652d8f58ed",
            "d0b19436553f41b29eca0b6a6c9b1019",
            "7eab30d1ddd147e3adbdcf374cab6f04",
            "7fb9c12460f241c39b58e8d402879033",
            "19a2183d08ae4deea471b6d18734ee19",
            "901710e665684edc96e9e5600c3a026e",
            "879ce688924641d1bb03cef1cea05245",
            "6b79bb7745ec4aac8dc4b26a248a87ba",
            "1d240e6e82a147019d1bf3e741bd55ba",
            "67966b3ed0d9411585737d1d35e60ee5",
            "9cdd7f4fab59464e90b606c6dbdfd916",
            "f99791e702b84b1ebbd15e0bd66b1333",
            "68670028637a420d90c999a05dd4206c",
            "5ebf724b7af443f49dd3a2ba954e030a",
            "80f2e1a5bd854f88a13db57817296c8c",
            "26765f58c1d44bf192e080cb1c5e7492",
            "d3046fe79ebf4627a69aebecc412a2ea",
            "ef508d90634b44d6add55edabd37399c",
            "520cecdef46949f5beb24bb7a33e187f",
            "80393295b52d4845972d3a454ae432b4",
            "f982afd1605542fdb0b042802ff7d7ef",
            "6189961fdb89450c852672574ed30d2c",
            "9e203c1650854dc0bcf54ea88023fe18",
            "50db8a1f9edf46d28072e739c0ca33cd",
            "83faa07de0e1441c8cc29d23316319fc",
            "77d73a672ede4f4e8b399629021b041a",
            "d83d7739f4e34486b3259bfcfea67a4b",
            "9907eb9b2da849dd8c9b1a14e427ea5f",
            "720d011a63b14bf7abab20d3f17b0785",
            "ff66ca34ec49449181fbee8546f44e58"
          ]
        },
        "id": "3oyUPIDV8mQ6",
        "outputId": "e2949f6a-2f87-4c44-f8f6-797208f8b96f"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Downloading http://fashion-mnist.s3-website.eu-central-1.amazonaws.com/train-images-idx3-ubyte.gz\n",
            "Downloading http://fashion-mnist.s3-website.eu-central-1.amazonaws.com/train-images-idx3-ubyte.gz to ./data/FashionMNIST/raw/train-images-idx3-ubyte.gz\n"
          ]
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "  0%|          | 0/26421880 [00:00<?, ?it/s]"
            ],
            "application/vnd.jupyter.widget-view+json": {
              "version_major": 2,
              "version_minor": 0,
              "model_id": "63dfa8a74fd045ffaac822f9f43e72e9"
            }
          },
          "metadata": {}
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Extracting ./data/FashionMNIST/raw/train-images-idx3-ubyte.gz to ./data/FashionMNIST/raw\n",
            "\n",
            "Downloading http://fashion-mnist.s3-website.eu-central-1.amazonaws.com/train-labels-idx1-ubyte.gz\n",
            "Downloading http://fashion-mnist.s3-website.eu-central-1.amazonaws.com/train-labels-idx1-ubyte.gz to ./data/FashionMNIST/raw/train-labels-idx1-ubyte.gz\n"
          ]
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "  0%|          | 0/29515 [00:00<?, ?it/s]"
            ],
            "application/vnd.jupyter.widget-view+json": {
              "version_major": 2,
              "version_minor": 0,
              "model_id": "b8f819130ff249f6b400af84844b265d"
            }
          },
          "metadata": {}
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Extracting ./data/FashionMNIST/raw/train-labels-idx1-ubyte.gz to ./data/FashionMNIST/raw\n",
            "\n",
            "Downloading http://fashion-mnist.s3-website.eu-central-1.amazonaws.com/t10k-images-idx3-ubyte.gz\n",
            "Downloading http://fashion-mnist.s3-website.eu-central-1.amazonaws.com/t10k-images-idx3-ubyte.gz to ./data/FashionMNIST/raw/t10k-images-idx3-ubyte.gz\n"
          ]
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "  0%|          | 0/4422102 [00:00<?, ?it/s]"
            ],
            "application/vnd.jupyter.widget-view+json": {
              "version_major": 2,
              "version_minor": 0,
              "model_id": "1d240e6e82a147019d1bf3e741bd55ba"
            }
          },
          "metadata": {}
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Extracting ./data/FashionMNIST/raw/t10k-images-idx3-ubyte.gz to ./data/FashionMNIST/raw\n",
            "\n",
            "Downloading http://fashion-mnist.s3-website.eu-central-1.amazonaws.com/t10k-labels-idx1-ubyte.gz\n",
            "Downloading http://fashion-mnist.s3-website.eu-central-1.amazonaws.com/t10k-labels-idx1-ubyte.gz to ./data/FashionMNIST/raw/t10k-labels-idx1-ubyte.gz\n"
          ]
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "  0%|          | 0/5148 [00:00<?, ?it/s]"
            ],
            "application/vnd.jupyter.widget-view+json": {
              "version_major": 2,
              "version_minor": 0,
              "model_id": "80393295b52d4845972d3a454ae432b4"
            }
          },
          "metadata": {}
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Extracting ./data/FashionMNIST/raw/t10k-labels-idx1-ubyte.gz to ./data/FashionMNIST/raw\n",
            "\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "Its much easier if we put both the dataloaders together in a single object. So, lets make a very simple class for this."
      ],
      "metadata": {
        "id": "Dn3vqNEF99Mz"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "class DataLoaders:\n",
        "    def __init__(self, *dls): self.train,self.valid = dls[:2]"
      ],
      "metadata": {
        "id": "hX2gI0As91Tr"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "dls = DataLoaders(train_dl, test_dl)\n",
        "dls.train, dls.valid"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "iv_UAOXf-KeV",
        "outputId": "e4a3922e-0b85-4081-8616-ff0fe1fb8bcd"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "(<torch.utils.data.dataloader.DataLoader at 0x7f0106396ca0>,\n",
              " <torch.utils.data.dataloader.DataLoader at 0x7f01040c9d90>)"
            ]
          },
          "metadata": {},
          "execution_count": 7
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "Perfect, we have our data in place. Next, we need to build our model. We will build a very simple CNN model"
      ],
      "metadata": {
        "id": "3gJUuAw5-s7i"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "def get_model():\n",
        "    return nn.Sequential(nn.Conv2d(1, 8, 3), nn.ReLU(),\n",
        "                      nn.Conv2d(8, 16, 3), nn.ReLU(),\n",
        "                      nn.Conv2d(16, 32, 3), nn.ReLU(),\n",
        "                      nn.AdaptiveAvgPool2d(1), nn.Flatten(),\n",
        "                      nn.Linear(32, 10))\n",
        "\n",
        "model = get_model()\n",
        "model"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "bbGTGJYT-S7G",
        "outputId": "6078ded6-a9b3-4bce-c685-1cd217ba1702"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "Sequential(\n",
              "  (0): Conv2d(1, 8, kernel_size=(3, 3), stride=(1, 1))\n",
              "  (1): ReLU()\n",
              "  (2): Conv2d(8, 16, kernel_size=(3, 3), stride=(1, 1))\n",
              "  (3): ReLU()\n",
              "  (4): Conv2d(16, 32, kernel_size=(3, 3), stride=(1, 1))\n",
              "  (5): ReLU()\n",
              "  (6): AdaptiveAvgPool2d(output_size=1)\n",
              "  (7): Flatten(start_dim=1, end_dim=-1)\n",
              "  (8): Linear(in_features=32, out_features=10, bias=True)\n",
              ")"
            ]
          },
          "metadata": {},
          "execution_count": 8
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "One last thing, our loss function"
      ],
      "metadata": {
        "id": "VDI58xjZFx5m"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "loss_func = nn.CrossEntropyLoss()"
      ],
      "metadata": {
        "id": "aV624bZUD1dT"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "I am cheating a bit by using a callback even though I have not discussed it. But bear with me, it will all make sense when we will talk about callbacks. For now, just think of it as a class responsible for tracking metrics and loss, and also printing them after each epoch."
      ],
      "metadata": {
        "id": "kL8TlLaoIQmj"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "!pip install -q torcheval"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "dhVai3FcKSrR",
        "outputId": "9f0d3236-05cc-49d8-f6a3-c93651bbb431"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m121.1/121.1 KB\u001b[0m \u001b[31m12.6 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m78.1/78.1 KB\u001b[0m \u001b[31m2.5 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25h"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "from copy import copy\n",
        "from collections.abc import Mapping\n",
        "from torcheval.metrics import MulticlassAccuracy, Mean\n",
        "\n",
        "def to_cpu(x):\n",
        "    if isinstance(x, Mapping): return {k:to_cpu(v) for k,v in x.items()}\n",
        "    if isinstance(x, list): return [to_cpu(o) for o in x]\n",
        "    if isinstance(x, tuple): return tuple(to_cpu(list(x)))\n",
        "    return x.detach().cpu()\n",
        "\n",
        "class MetricsCB(Callback):\n",
        "    def __init__(self, *ms, **metrics):\n",
        "        for o in ms: metrics[type(o).__name__] = o\n",
        "        self.metrics = metrics\n",
        "        self.all_metrics = copy(metrics)\n",
        "        self.all_metrics['loss'] = self.loss = Mean()\n",
        "\n",
        "    def _log(self, d): print(d)\n",
        "    def before_fit(self, learn): learn.metrics = self\n",
        "    def before_epoch(self, learn): [o.reset() for o in self.all_metrics.values()]\n",
        "\n",
        "    def after_epoch(self, learn):\n",
        "        log = {k:f'{v.compute():.3f}' for k,v in self.all_metrics.items()}\n",
        "        log['epoch'] = learn.epoch\n",
        "        log['train'] = 'train' if learn.model.training else 'eval'\n",
        "        self._log(log)\n",
        "\n",
        "    def after_batch(self, learn):\n",
        "        x,y,*_ = to_cpu(learn.batch)\n",
        "        for m in self.metrics.values(): m.update(to_cpu(learn.preds), y)\n",
        "        self.loss.update(to_cpu(learn.loss), weight=len(x))"
      ],
      "metadata": {
        "id": "MHRZk0gkHRVW"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "metrics = MetricsCB(accuracy=MulticlassAccuracy())"
      ],
      "metadata": {
        "id": "7wrIaxkjOPqp"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Now we are all set to test our training loop. I will start by discussing the first feature in detail, followed by the second feature."
      ],
      "metadata": {
        "id": "ju6UJFc2F71P"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Inheriting Learner\n",
        "\n",
        "The very first thing that we should do is define `predict`, `get_loss`, `backward`, `step`, and `zero_grad` methods so that we can train our CNN model on our dataset.\n",
        "\n",
        "To accomplish this, we can simply inherit the `Learner` class and implement these method as follows:"
      ],
      "metadata": {
        "id": "VsvB4qsGGXAt"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "class TrainLearner(Learner):\n",
        "    def predict(self): self.preds = self.model(self.batch[0])   # batch -> (x, y)\n",
        "    def get_loss(self): self.loss = self.loss_func(self.preds, self.batch[1])\n",
        "    def backward(self): self.loss.backward()\n",
        "    def step(self): self.opt.step()\n",
        "    def zero_grad(self): self.opt.zero_grad()"
      ],
      "metadata": {
        "id": "OXHx8tNKFOB9"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "As you can see, all our methods are so concise and to the point. Lets train our model."
      ],
      "metadata": {
        "id": "Vu9Wx6b3s81q"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "learner = TrainLearner(get_model(), dls, cbs=metrics, loss_func=loss_func)\n",
        "learner.fit(2)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "5MH_KvXdJIgl",
        "outputId": "778e9ab0-b46d-428c-a4db-3dbf85e4c8e5"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "{'accuracy': '0.400', 'loss': '1.572', 'epoch': 0, 'train': 'train'}\n",
            "{'accuracy': '0.544', 'loss': '1.086', 'epoch': 0, 'train': 'eval'}\n",
            "{'accuracy': '0.651', 'loss': '0.945', 'epoch': 1, 'train': 'train'}\n",
            "{'accuracy': '0.636', 'loss': '0.960', 'epoch': 1, 'train': 'eval'}\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "This `TrainLearner` class can now be used to train any model on any dataset if you have a simple training loop. But what if you have a more complex training loop? Easy! just update the corresponding method implementation.\n",
        "\n",
        "For example, you want to introduce *momentum*"
      ],
      "metadata": {
        "id": "6EYVfHlotZ5i"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "class MomentumLearner(TrainLearner):\n",
        "    def __init__(self, model, dls=(0,), loss_func=F.mse_loss, lr=0.1, cbs=None, opt_func=optim.SGD, mom=0.85):\n",
        "        self.mom = mom\n",
        "        super().__init__(model, dls, loss_func, lr, cbs, opt_func)\n",
        "    \n",
        "    def zero_grad(self):\n",
        "        for p in self.model.parameters():\n",
        "            p.grad *= self.mom"
      ],
      "metadata": {
        "id": "Xg8xYPyqK5EH"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "We will start by inheriting the `TrainLearner` class. We will update the `__init__` method to take an extra parameter (`mom`). Next, we just need to update the `zero_grad` method to multiply all the gradients with `self.mom` instead of zeroing them. "
      ],
      "metadata": {
        "id": "u4k1s_H6ugDI"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "learner = MomentumLearner(get_model(), dls, cbs=metrics, loss_func=loss_func)\n",
        "learner.fit(2)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "otX9Wu7tMbFa",
        "outputId": "142d1894-6e8f-4e9c-c21d-39e56a30ddde"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "{'accuracy': '0.561', 'loss': '1.162', 'epoch': 0, 'train': 'train'}\n",
            "{'accuracy': '0.707', 'loss': '0.932', 'epoch': 0, 'train': 'eval'}\n",
            "{'accuracy': '0.735', 'loss': '0.738', 'epoch': 1, 'train': 'train'}\n",
            "{'accuracy': '0.754', 'loss': '0.681', 'epoch': 1, 'train': 'eval'}\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "Hurry! we jumped from ~65% accuracy to ~75% just by adding momentum. And also the loss is much lower.\n",
        "\n",
        "One can already appreciate how easy it was to add momentum to our learner. Things get even more exciting once you start using callbacks to run code at other stages of the training loop. "
      ],
      "metadata": {
        "id": "yGa_PQuFwvPJ"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Callbacks\n",
        "\n",
        "Callbacks are pretty widely used in programming. They basically allow you to inject and run arbitrary code in your program. One can pass functions as variables to other functions in Python. Hence, using callbacks in python feels very natural.\n",
        "\n",
        "We will start by building some simple callbacks and gradually build our way to more advanced callbacks."
      ],
      "metadata": {
        "id": "cilRFGLpyhlx"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Single Batch & Epoch Callback"
      ],
      "metadata": {
        "id": "7gTUZkGu0Me3"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "Lets write a basic callback that will stop the training after one batch. To stop the training after one batch, you can just raise `CancelFitException` after batch. Heres the code"
      ],
      "metadata": {
        "id": "_kiuEem30zSd"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "class SingleBatchCB(Callback):\n",
        "    def after_batch(self, learner): \n",
        "        print('[SingleBatchCB] Terminating . . . ')\n",
        "        raise CancelFitException()"
      ],
      "metadata": {
        "id": "q3ZfHQuSMenO"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Its pretty straight forward. Lets see it in action."
      ],
      "metadata": {
        "id": "x0RgZkvo0WsM"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "cbs = [metrics, SingleBatchCB()]\n",
        "learner = MomentumLearner(get_model(), dls, cbs=cbs, loss_func=loss_func)\n",
        "learner.fit(2)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "axwm2-Qf0LSU",
        "outputId": "5a798db8-b0a7-4f69-e3b3-8dc4be4a7f84"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "[SingleBatchCB] Terminating . . . \n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "The training terminates immediately even after calling `fit()` with 2 epochs.\n",
        "\n",
        "Similarly, heres the code for `SingleEpochCB`"
      ],
      "metadata": {
        "id": "Ot82j8XG3G9M"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "class SingleEpochCB(Callback):\n",
        "    def after_epoch(self, learner): \n",
        "        if not learner.model.training:\n",
        "            print('[SingleEpochCB] Terminating . . . ') \n",
        "            raise CancelFitException()"
      ],
      "metadata": {
        "id": "zRmCeBx93PjF"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Here, we have one extra thing. Both training and validation are epochs. To differentiate them, we can use the `learner.model.training` attribute. In the above cell, we will only raise exception if its a validation epoch."
      ],
      "metadata": {
        "id": "hQcikjxJ3h4_"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "cbs = [metrics, SingleEpochCB()]\n",
        "learner = MomentumLearner(get_model(), dls, cbs=cbs, loss_func=loss_func)\n",
        "learner.fit(2)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "-dHNia7i4_Bd",
        "outputId": "5c3726a2-ed69-4e23-8388-86c422d6013c"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "{'accuracy': '0.597', 'loss': '1.076', 'epoch': 0, 'train': 'train'}\n",
            "{'accuracy': '0.678', 'loss': '0.873', 'epoch': 0, 'train': 'eval'}\n",
            "[SingleEpochCB] Terminating . . . \n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "So, our training terminates after one epoch. You can see the loss printed after both training and validation.  \n",
        "\n",
        "Both these callbacks can be used to quickly test if the code is working properly or not. Minor errors like dimension mismatch, incorrect data types, etc can be caught very easily. Furthermore, these callbacks can be used in scenarios where you have to do a forward pass like looking at the model summary. We will be implementing a summary callback toward the end of this notebook."
      ],
      "metadata": {
        "id": "Ie73-49Y5Es0"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Device and Channel Last Callback"
      ],
      "metadata": {
        "id": "0JxvgCWK6jJ4"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "The training is pretty slow on CPU. So lets write a callback that would move our model and batch to GPU."
      ],
      "metadata": {
        "id": "SuL4WrimyHs4"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "def_device = 'mps' if torch.backends.mps.is_available() else 'cuda' if torch.cuda.is_available() else 'cpu'\n",
        "\n",
        "def to_device(x, device=def_device):\n",
        "    if isinstance(x, torch.Tensor): return x.to(device)\n",
        "    if isinstance(x, Mapping): return {k:v.to(device) for k,v in x.items()}\n",
        "    return type(x)(to_device(o, device) for o in x)\n",
        "\n",
        "class DeviceCB(Callback):\n",
        "    def __init__(self, device=def_device): self.device=device\n",
        "    def before_fit(self, learn): learn.model.to(self.device)\n",
        "    def before_batch(self, learn): learn.batch = to_device(learn.batch, self.device)"
      ],
      "metadata": {
        "id": "JKortHMXPD3o"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "`DeviceCB` simply moves the *model* to device *before_fit*, and moves the *batch* to device *before_batch*."
      ],
      "metadata": {
        "id": "3EMHDJzozp1Q"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "cbs = [metrics, DeviceCB()]\n",
        "learner = MomentumLearner(get_model(), dls, cbs=cbs, loss_func=loss_func)\n",
        "learner.fit(2)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "GkHOcQFyV7qV",
        "outputId": "b8bef1a5-a242-4dbe-eb3a-ccc1373d36f9"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "{'accuracy': '0.587', 'loss': '1.106', 'epoch': 0, 'train': 'train'}\n",
            "{'accuracy': '0.765', 'loss': '0.715', 'epoch': 0, 'train': 'eval'}\n",
            "{'accuracy': '0.770', 'loss': '0.656', 'epoch': 1, 'train': 'train'}\n",
            "{'accuracy': '0.791', 'loss': '0.615', 'epoch': 1, 'train': 'eval'}\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "This time around, the training was much quicker. We can make the training even faster by using *channel last* memory format. You can read more about it [here](https://pytorch.org/blog/accelerating-pytorch-vision-models-with-channels-last-on-cpu/). Lets write a callback for that\n",
        "\n"
      ],
      "metadata": {
        "id": "uTuCMgGUy6LB"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "class ChannelLastCB(Callback):\n",
        "    def before_fit(self, learn): learn.model = learn.model.to(memory_format=torch.channels_last)\n",
        "    def before_batch(self, learn): learn.batch[0] = learn.batch[0].contiguous(memory_format=torch.channels_last)"
      ],
      "metadata": {
        "id": "E-OGb3Aocc36"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "cbs = [metrics, ChannelLastCB(), DeviceCB()]\n",
        "learner = MomentumLearner(get_model(), dls, cbs=cbs, loss_func=loss_func)\n",
        "learner.fit(2)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "RoI1VBXT1akw",
        "outputId": "e9d778f7-cf58-46d0-cdb7-9e8b11c63baa"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "{'accuracy': '0.600', 'loss': '1.065', 'epoch': 0, 'train': 'train'}\n",
            "{'accuracy': '0.738', 'loss': '0.719', 'epoch': 0, 'train': 'eval'}\n",
            "{'accuracy': '0.763', 'loss': '0.659', 'epoch': 1, 'train': 'train'}\n",
            "{'accuracy': '0.785', 'loss': '0.606', 'epoch': 1, 'train': 'eval'}\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "As you can see its so easy to introduce new ideas. We can write a callback for each new idea and pass them all to our learner. Once you start working with multiple callbacks it becomes quite powerful."
      ],
      "metadata": {
        "id": "oGsitdPg2exI"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "### TrainCB"
      ],
      "metadata": {
        "id": "Ld7rkZiU_L5N"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "Earlier we inherited the `Learner` class to define `TrainLearner`. We can also implement the same using callbacks. \n",
        "\n",
        "Here is a `TrainCB` that will define `predict`, `get_loss`, `backward`, `step`, and `zero_grad` methods for training"
      ],
      "metadata": {
        "id": "X4uaGAYq_OZQ"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "class TrainCB(Callback):\n",
        "    def predict(self, learn): learn.preds = learn.model(learn.batch[0])\n",
        "    def get_loss(self, learn): learn.loss = learn.loss_func(learn.preds, learn.batch[1])\n",
        "    def backward(self, learn): learn.loss.backward()\n",
        "    def step(self, learn): learn.opt.step()\n",
        "    def zero_grad(self, learn): learn.opt.zero_grad()"
      ],
      "metadata": {
        "id": "oMwIY0t1_LRB"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Now we can pass this callback to our `Learner` class, not `TrainLearner`, to start training."
      ],
      "metadata": {
        "id": "gmNr22vx__ei"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "cbs = [TrainCB(), metrics, DeviceCB()]\n",
        "learner = Learner(get_model(), dls, loss_func=loss_func, cbs=cbs)\n",
        "learner.fit(2)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "4N96RykrAG0w",
        "outputId": "ba26bbd9-c810-4752-ed53-afb99e0ab2d2"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "{'accuracy': '0.391', 'loss': '1.637', 'epoch': 0, 'train': 'train'}\n",
            "{'accuracy': '0.519', 'loss': '1.208', 'epoch': 0, 'train': 'eval'}\n",
            "{'accuracy': '0.633', 'loss': '0.983', 'epoch': 1, 'train': 'train'}\n",
            "{'accuracy': '0.628', 'loss': '0.927', 'epoch': 1, 'train': 'eval'}\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "This callback highlights one very important insight. Even though, `predict` and other methods are not defined in the `Learner` class its able to use the `predict` method from the callback. \n",
        "\n",
        "The `Learner` class first checks for `predict` method inside `self`, if no `predict` method is found then it will search for it inside callbacks. Error is only raised when there is no `predict` method in any callback.\n",
        "\n",
        "Lets introduce more ideas using callbacks"
      ],
      "metadata": {
        "id": "VTABcssOAUXL"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Gradient Accumulation & Clipping Callback"
      ],
      "metadata": {
        "id": "Vo5_vwTG66l5"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "Next, we will write a callback for [gradient clipping](https://stackoverflow.com/a/56069467). Just like all above callbacks, it would not be more than a single line of code"
      ],
      "metadata": {
        "id": "aHqVZBRg2drj"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "class GradientClippingCB(Callback):\n",
        "    def __init__(self, max_norm=1.0, norm_type=2): \n",
        "        self.max_norm, self.norm_type = max_norm, norm_type\n",
        "        \n",
        "    def after_backward(self, learner):\n",
        "        torch.nn.utils.clip_grad_norm_(learner.model.parameters(), self.max_norm, self.norm_type)"
      ],
      "metadata": {
        "id": "xoQDAXyqZotl"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Ever heard about Gradient Accumulation? This technique allows you to train models with large batch size even when your GPU can hold only a few samples before it runs out of memory.\n",
        "\n",
        "You can read more about it in [this blog](https://medium.com/huggingface/training-larger-batches-practical-tips-on-1-gpu-multi-gpu-distributed-setups-ec88c3e51255). Here is a short excert from the blog to give you a gist of how its done.\n",
        "> Accumulating gradients just means that, before calling `optimizer.step()` to perform a step of gradient descent, we will sum the gradients of several backward operations in the `parameter.grad` tensors. This is straightforward to do in PyTorch as the gradient tensors are not reset unless we call `model.zero_grad()` or `optimizer.zero_grad()`. We’ll also need to divide by the number of accumulation steps if our loss is averaged over the training samples."
      ],
      "metadata": {
        "id": "cdi6QnNz4-7K"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "class GradientAccumulationCB(Callback):\n",
        "    def __init__(self, n): \n",
        "        self.counter = 0\n",
        "        self.n = n\n",
        "\n",
        "    def after_loss(self, learner):\n",
        "        learner.loss = learner.loss / self.n\n",
        "        \n",
        "    def after_backward(self, learner):\n",
        "        self.counter += 1\n",
        "        if self.counter % self.n != 0:\n",
        "            raise CancelBatchException()"
      ],
      "metadata": {
        "id": "YF1Ri_iqWN6K"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "In the `after_backward()` method we are only zeroing the gradient when `self.counter` is divisible by `self.n`. At other times, we are just raising the `CancelBatchException` exception so neither `self.step()` nor `self.zero_grad()` methods are executed.\n",
        "\n",
        "But there is one shuttle problem with this callback. What if you have 19 batches in total and we are accumulating gradients every 5 batches? Clearly, the gradients will be accumulated in the last 4 batches but weights will not be updated.\n",
        "\n",
        "How do you fix this? Simple, call `self.step()` and `self.zero_grad()` after each training epoch. Here is the updated code"
      ],
      "metadata": {
        "id": "OUtZW3uIBJmu"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "class GradientAccumulationCB(Callback):\n",
        "    def __init__(self, n): \n",
        "        self.counter = 0\n",
        "        self.n = n\n",
        "\n",
        "    def after_loss(self, learner):\n",
        "        learner.loss = learner.loss / self.n\n",
        "        \n",
        "    def after_backward(self, learner):\n",
        "        self.counter += 1\n",
        "        if self.counter % self.n != 0:\n",
        "            raise CancelBatchException()\n",
        "\n",
        "    def after_epoch(self, learner):\n",
        "        if learner.model.training: \n",
        "            learner.step()\n",
        "            learner.zero_grad()"
      ],
      "metadata": {
        "id": "CAxueS4oB4Ro"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Lets train our model with both gradient clipping and accumulation."
      ],
      "metadata": {
        "id": "HRrjrQfLGITj"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "cbs = [metrics, ChannelLastCB(), DeviceCB(), GradientClippingCB(), GradientAccumulationCB(2)]\n",
        "learner = MomentumLearner(get_model(), dls, cbs=cbs, loss_func=loss_func)\n",
        "learner.fit(2)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "lxLeKbb0JBi1",
        "outputId": "0a559de7-b387-4220-9f20-2c2f9463a14a"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "{'accuracy': '0.469', 'loss': '0.710', 'epoch': 0, 'train': 'train'}\n",
            "{'accuracy': '0.542', 'loss': '0.657', 'epoch': 0, 'train': 'eval'}\n",
            "{'accuracy': '0.660', 'loss': '0.465', 'epoch': 1, 'train': 'train'}\n",
            "{'accuracy': '0.598', 'loss': '0.555', 'epoch': 1, 'train': 'eval'}\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "it would be exciting to see if we can train the model with a large learning rate"
      ],
      "metadata": {
        "id": "Ur1pB5KIJlsC"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Mixed Precision Training"
      ],
      "metadata": {
        "id": "uYfnlCKe7L1f"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "Lets implement mixed precision training. For details, refer [this pytorch tutorial](https://pytorch.org/tutorials/recipes/recipes/amp_recipe.html)"
      ],
      "metadata": {
        "id": "xzL3iHO4TJSt"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "class MixedTrainingCB(Callback):\n",
        "    def before_fit(self, learn): self.scaler = torch.cuda.amp.GradScaler()\n",
        "    def before_batch(self, learn):\n",
        "        self.autocast = torch.autocast('cuda', dtype=torch.float16)\n",
        "        self.autocast.__enter__()\n",
        "    def after_loss(self, learn):\n",
        "        self.autocast.__exit__(None, None, None)\n",
        "    def backward(self, learn): self.scaler.scale(learn.loss).backward()\n",
        "    def step(self, learn): \n",
        "        self.scaler.step(learn.opt)\n",
        "        self.scaler.update()"
      ],
      "metadata": {
        "id": "Vh0AehRbcioL"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "cbs = [metrics, DeviceCB(), MixedTrainingCB()]\n",
        "learner = MomentumLearner(get_model(), dls, cbs=cbs, loss_func=loss_func)\n",
        "learner.fit(2)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "ncPO0FFLS3fZ",
        "outputId": "aeda60fd-d7a5-4acd-ed75-ab01f71239c5"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "{'accuracy': '0.599', 'loss': '1.071', 'epoch': 0, 'train': 'train'}\n",
            "{'accuracy': '0.748', 'loss': '0.701', 'epoch': 0, 'train': 'eval'}\n",
            "{'accuracy': '0.779', 'loss': '0.617', 'epoch': 1, 'train': 'train'}\n",
            "{'accuracy': '0.813', 'loss': '0.534', 'epoch': 1, 'train': 'eval'}\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "There's an even easier way to implement mixed precision training using the `accelerate` python package. Its pretty simple to integrate "
      ],
      "metadata": {
        "id": "0GJWkQuET7hh"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "!pip install -q accelerate"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "Po78wdGSgG9A",
        "outputId": "3baac68a-cb33-4129-885e-fdaf8b455d58"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "\u001b[?25l     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m0.0/191.5 KB\u001b[0m \u001b[31m?\u001b[0m eta \u001b[36m-:--:--\u001b[0m\r\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m191.5/191.5 KB\u001b[0m \u001b[31m16.3 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25h"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "from accelerate import Accelerator"
      ],
      "metadata": {
        "id": "sEyFfgRwgFjq"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "class AccelerateCB(TrainCB):\n",
        "    def __init__(self, mixed_precision=\"fp16\"):\n",
        "        super().__init__()\n",
        "        self.acc = Accelerator(mixed_precision=mixed_precision)\n",
        "        \n",
        "    def before_fit(self, learn):\n",
        "        learn.model,learn.opt,learn.dls.train,learn.dls.valid = self.acc.prepare(\n",
        "            learn.model, learn.opt, learn.dls.train, learn.dls.valid)\n",
        "\n",
        "    def backward(self, learn): self.acc.backward(learn.loss)"
      ],
      "metadata": {
        "id": "q1DDqBNvfkJD"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Note that we are inheriting from `TrainCB` instead of `Callback` because we are updating the `backward` method. Hence, we will be using `Learner` and not `TrainLearner` or `MomentumLearner` class for training."
      ],
      "metadata": {
        "id": "I53SeF4CaQtJ"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "cbs = [metrics, DeviceCB(), AccelerateCB()]\n",
        "learner = Learner(get_model(), dls, cbs=cbs, loss_func=loss_func)\n",
        "learner.fit(2)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "_eWqktGuUUDL",
        "outputId": "30d99465-e56d-466c-cc68-f649ddf4ce51"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "{'accuracy': '0.341', 'loss': '1.761', 'epoch': 0, 'train': 'train'}\n",
            "{'accuracy': '0.557', 'loss': '1.285', 'epoch': 0, 'train': 'eval'}\n",
            "{'accuracy': '0.647', 'loss': '0.964', 'epoch': 1, 'train': 'train'}\n",
            "{'accuracy': '0.692', 'loss': '0.887', 'epoch': 1, 'train': 'eval'}\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "You might have already realized the power of this callback feature. The most important part is isolation. When you are implementing some feature or research paper, you can just focus on the feature at hand. Everything else will just work. This highly reduces the chances of error, and also increases the speed of experimentation. "
      ],
      "metadata": {
        "id": "efuyfYUIXBnG"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "### LR Finder"
      ],
      "metadata": {
        "id": "QUj-7_IC7V8v"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "Lets build the iconic `lr_find()`. The idea is to train the model with growing learning rates until the current loss becomes greater than some multiple of minimum loss. Here, we will use 3 as the multiple. Then plots the losses vs the learning rates with a log scale."
      ],
      "metadata": {
        "id": "k60NvaCUbKVl"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "import math\n",
        "import matplotlib.pyplot as plt"
      ],
      "metadata": {
        "id": "-C8hXtIwb41C"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "class LRFinderCB(Callback):\n",
        "    def __init__(self, lr_mult=1.3): \n",
        "        self.lr_mult = lr_mult\n",
        "    def before_fit(self, learn): \n",
        "        self.lrs, self.losses = [], []\n",
        "        self.min = math.inf\n",
        "    def after_batch(self, learn):\n",
        "        if not learn.training: raise CancelEpochException()\n",
        "        self.lrs.append(learn.opt.param_groups[0]['lr'])\n",
        "        loss = to_cpu(learn.loss)\n",
        "        self.losses.append(loss)\n",
        "        if loss < self.min: self.min = loss\n",
        "        if loss > self.min*3: raise CancelFitException()\n",
        "        for g in learn.opt.param_groups: g['lr'] *= self.lr_mult"
      ],
      "metadata": {
        "id": "Xg14vZWMbROE"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "lrfind = LRFinderCB()\n",
        "cbs = [DeviceCB(), lrfind]\n",
        "learn = MomentumLearner(get_model(), dls, lr=1e-4, cbs=cbs, loss_func=loss_func)\n",
        "learn.fit(1)\n",
        "plt.plot(lrfind.lrs, lrfind.losses)\n",
        "plt.xscale('log')"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 269
        },
        "id": "3hOIRJlGcpNy",
        "outputId": "dc3596b1-f151-442c-b65e-7952bffba863"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "<Figure size 432x288 with 1 Axes>"
            ],
            "image/png": "iVBORw0KGgoAAAANSUhEUgAAAW4AAAD8CAYAAABXe05zAAAABHNCSVQICAgIfAhkiAAAAAlwSFlzAAALEgAACxIB0t1+/AAAADh0RVh0U29mdHdhcmUAbWF0cGxvdGxpYiB2ZXJzaW9uMy4yLjIsIGh0dHA6Ly9tYXRwbG90bGliLm9yZy+WH4yJAAAbyklEQVR4nO3deZRcd3nm8e9bW2/qbm2tXbJkQDYGL5jGNhAcMDETwIYZhglgtkCCGLKwhBkGTnLCJGdmApmECSQMIHAgBMxmYMAEMzhgDPbES9sWso1sRbIttWRZanVr6a5bXes7f1R1W25rqequqntv+fmcU6eWvl31/nSrn7p676/uNXdHRETiIxF2ASIi0hgFt4hIzCi4RURiRsEtIhIzCm4RkZhRcIuIxEyqFU+6fPly37hxYyueWkSkI919992H3X2onmVbEtwbN25kZGSkFU8tItKRzGxPvcuqVSIiEjMKbhGRmFFwi4jEjIJbRCRmFNwiIjGj4BYRiRkFt4hIEzzw2DFu2TnWltdScIuINMFXbt/Lf/rWL9vyWgpuEZEmCAolejPJtryWgltEpAmCQpneTEu+jP4UCm4RkSbIFcrR2uI2s/eZ2f1m9oCZvb/VRYmIxE02Sq0SM3su8C7gEuBC4Coze2arCxMRiZOobXE/G7jD3QN3LwG3AK9rbVkiIvEStR73/cBLzGyZmfUCrwLWt7YsEZF4CQoletq0xX3Gjwd332FmHwd+DGSBbUB57nJmtgXYArBhw4YmlykiEm1BoUxfhFoluPu17v58d78cOALsPMkyW9192N2Hh4bqOomDiEhHqFScXLFMT5taJXW9ipmtcPdDZraBan/7staWJSISH9OlMu60bedkvR8P3zazZUAR+H13P9rCmkREYiUoVLvH7WqV1BXc7v6SVhciIhJXuVpwt6tVom9OiogsULZQAtrXKlFwi4gs0EyrRMEtIhITudngVqtERCQWsnm1SkREYiVXVKtERCRWArVKRETiZaZV0q5jlSi4RUQWKKdZJSIi8RIUy2SSCdLJ9kSqgltEZIGCfPsO6QoKbhGRBWvnIV1BwS0ismBBsawtbhGROAnypbZNBQQFt4jIggVtPFEwKLhFRBYsV1Rwi4jESlatEhGReMkVtHNSRCRWgqKmA4qIxEqQb98Z3qHO4DazD5jZA2Z2v5l9zcy6W12YiEgclMoVCuVKtHZOmtla4L3AsLs/F0gCb2x1YSIicRC0+VjcUH+rJAX0mFkK6AUea11JIiLxEeTbeyxuqCO43X0/8FfAXuAAcMzdfzx3OTPbYmYjZjYyNjbW/EpFRCIoaPMZ3qG+VskS4LXAJmAN0Gdmb5m7nLtvdfdhdx8eGhpqfqUiIhHU7jO8Q32tkt8AHnH3MXcvAt8BXtTaskRE4qHdpy2D+oJ7L3CZmfWamQEvB3a0tiwRkXiYaZVE6gs47n4HcD1wD3Bf7Xe2trguEZFYmDltWV9X+4K7rm17d/8o8NEW1yIiEjvZmVZJOlqtEhEROYVcFFslIiJyakEIrRIFt4jIAsy0SrpTCm4RkVjIFUr0pJMkEta211Rwi4gsQFAot7VNAgpuEZEFCdp8EgVQcIuILEhQKLV1KiAouEVEFiQolOlVq0REJD6CQnvP8A4KbhGRBQkKZXrUKhERiY9coaRZJSIicZJVq0REJF5yapWIiMSHu5NVq0REJD7ypQru7T0yICi4RUTmbfa0ZWkFt4hILGTztTO8d6nHLSISC7li+8/wDnUEt5mdY2bbTrgcN7P3t6M4EZEoe+IM7+0N7jNu37v7Q8BFAGaWBPYD321xXSIikRfMtEoy0W6VvBzY7e57WlGMiEichLXF3WhwvxH4WisKERGJmyCqPe4ZZpYBXgN86xQ/32JmI2Y2MjY21qz6REQiKw6tklcC97j7wZP90N23uvuwuw8PDQ01pzoRkQiLQ6vkTahNIiIya2Y6YCS/OWlmfcCVwHdaW46ISHxk8yVSCSOTbO9XYupqzLh7FljW4lpERGJl5kTBZtbW19U3J0VE5ikXwrG4QcEtIjJv2UKJvjbPKAEFt4jIvOVqrZJ2U3CLiMxTGGd4BwW3iMi8BYVS2798AwpuEZF50xa3iEjMBOpxi4jES6BZJSIi8aJWiYhIjJQrTr5UUatERCQugkL1kK5qlYiIxESuEM6RAUHBLSIyL2EdixsU3CIi85IthHP2G1Bwi4jMS05b3CIi8aJWiYhIzARqlYiIxIu2uEVEYibywW1mi83sejN70Mx2mNkLW12YiEiUzbZKutrfKqn3FT8J/MjdX29mGaC3hTWJiETezBZ3T7r9W9xnDG4zGwQuB34bwN0LQKG1ZYmIRFuuUKYrlSCZaO8Z3qG+VskmYAz4opnda2ZfMLO+uQuZ2RYzGzGzkbGxsaYXKiISJdlCib4Q2iRQX3CngIuBz7j784As8OG5C7n7VncfdvfhoaGhJpcpIhItQaEcSpsE6gvufcA+d7+jdv96qkEuIvK0lQvpWNxQR3C7++PAqJmdU3vo5cCvWlqViEjEZQvlUGaUQP2zSv4Q+GptRsnDwDtaV5KISPTlCiV6Q2qV1BXc7r4NGG5xLSIisREUyqwaSIfy2vrmpIjIPAQhtkoU3CIi8xCE2CpRcIuIzENQKIdy2jJQcIuINMzdCQpl+roU3CIisVAoVyhXPJRjcYOCW0SkYbkQDzAFCm4RkYZla8GtVomISEzkasfi7lGrREQkHmbPfqNWiYhIPGTzteBWq0REJB5yxfDO8A4KbhGRhoV5omBQcIuINCzIK7hFRGJl9gzvapWIiMRDUNQWt4hIrAT5MgmDrlQ4EargFhFpUFAo05tJYWahvL6CW0SkQbliKbRDuoKCW0SkYdl8mb4Qg7uuXaJm9igwCZSBkrvr/JMi8rRVPYlCODNKoP6zvAO8zN0Pt6wSEZGYyBVLoc0oAbVKREQals2XYxHcDvzYzO42sy0nW8DMtpjZiJmNjI2NNa9CEZGIyRXiEdy/5u4XA68Eft/MLp+7gLtvdfdhdx8eGhpqapEiIlESFEuhfWsS6gxud99fuz4EfBe4pJVFiYhEWRD1VomZ9ZlZ/8xt4BXA/a0uTEQkqoKQWyX1bOuvBL5b+4ZQCrjO3X/U0qpERCKqUnFyxYhPB3T3h4EL21CLiEjk5WoHmArzCziaDigi0oCwT6IACm4RkYbkasEdZqtEwS0i0oDs7EkUtMUtIhILapWIiMRMbja41SoREYkFtUpERGImp1aJiEg0TU4X+eXoUdz9SY8HapWIiETTZ2/ZzWs/fRtX/e2t/PC+A5Qr1QAPaq2SME9dFt5HhohIhO06NMWyvgxBoczvffUezh7q4z2//gyO54pAuK0SBbeIyEnsnchx0frFbH3bMDfef4BP37yb/3z9dswgk0yQTobXsFBwi4jM4e6MTgRcumkpyYRx1QVrePX5q/nZzjE+c/NuqsfcC4+CW0RkjiNBkal8ifVLe2cfMzNeds4KXnbOihArq9LOSRGROUYnAgA2nBDcUaLgFhGZY6+CW0QkXmaCe92SnpArOTkFt4jIHKMTAcsXZejriuZuQAW3iMgco0eCJ+2YjJq6g9vMkmZ2r5n9oJUFiYiEbe9EENn+NjS2xf0+YEerChERiYJSucJjR6dZvyTmwW1m64BXA19obTkiIuE6cGyacsU7Yov7b4APAZUW1iIiErqZGSWx7nGb2VXAIXe/+wzLbTGzETMbGRsba1qBIiLt9ERwR3MqINS3xf1i4DVm9ijwdeAKM/vK3IXcfau7D7v78NDQUJPLFBFpj70TAamEsXowxsHt7h9x93XuvhF4I/BTd39LyysTEQnB6ETAuiU9JBMhH0nqNDSPW0TkBKMT0Z7DDQ0Gt7v/zN2valUxIiJh29tpwS0i0skmp4scCYqRngoICm4RkVmjEzkgukcFnKHgFhGpifrhXGcouEVEamZOoBDlr7uDgltEZNbokYCB7hSDvemwSzktBbeISM3eiYANy6K9tQ0KbhGRWXsngsi3SUDBLSICQKXi7JvIRX7HJCi4RUQAODSZp1CuRP7LN6DgFhEB4jMVEBTcIiJAPI7DPUPBLSJCNbjNYO3i6B7OdYaCW0QE2DcRsGawh0wq+rEY/QpFRNqgelTA6G9tg4JbRASIzxxuUHCLiDBdLHNoMh+LGSWg4BYRYd+R2lTAGHzdHRTcIiKzUwHXqVUiIhIPe8fj8+UbqCO4zazbzO40s1+a2QNm9mftKExEpF1Gj+ToSSdZvigTdil1SdWxTB64wt2nzCwN3GpmN7r77S2uTUSkLfZOBGxY2ouZhV1KXc4Y3O7uwFTtbrp28VYWJSLSTqMxmsMNdfa4zSxpZtuAQ8BN7n7HSZbZYmYjZjYyNjbW7DpFRFrC3WvBHY/+NtQZ3O5edveLgHXAJWb23JMss9Xdh919eGhoqNl1ioi0xES2QLZQjs2OSWhwVom7HwVuBn6zNeWIiLRXnA7nOqOeWSVDZra4drsHuBJ4sNWFiYi0Q5wO5zqjnlklq4F/MLMk1aD/prv/oLVliYi0x8ijRzCDdUvis3Oynlkl24HntaEWEZG2+vnOMf7x9j286ZIN9Gbq2Y6NBn1zUkSelg5NTvNH39zG5pWL+NOrzgu7nIbE5yNGRKRJKhXnA9/YxlS+xHXvuoyeTDLskhqi4BaRp53P3LKb23aN87HXnc/mlf1hl9MwtUpEOlyxXOETN+1kx4HjYZcSCSOPTvCJm3Zy9YVreMML1oddzrwouEU63I8fOMinfvKvvOFz/8K20aNhlxOqo0GB937tXtYu7uF//LvnxubYJHMpuEU63Nfu3MuqgW4W92Z4yxfu4K5HJ8IuqaWmi2X+Zfc49+49wq5Dkxw8Pk02X6JScT50/XbGpvL83TXPo787HXap86Yet0gH2zOe5dZdh/mjKzfzW8Pruebzt/O2a+/k2rcP86JnLg+7vKZ75HCW93zlbh58fPIpPzMDd/iTVz+bC9YtDqG65lFwi3Swb9w1SsLgt4bXs2qwm6+/+zLe8oU7eMeX7uJzb30+Lz1nRdglntHDY1PsPDjJ5ZuHTjvX+of3HeBD128nnTT+5g0XMdibZnK6xNR0icnpIlP5Eiv6u3jLZWe1sfrWUHCLdKhiucI3R/ZxxbkrWTXYDcCK/m6+vuWFvPXaO3jXl0f49DUX84rnrAq50qfadyTgn7Yf4Ibtj3H//upO1cGeNG++dANvf9FGVg50zy5bKFX4ixt38MXbHuV5Gxbz6WsuZs3i+HwLcj4U3CId6ic7DnJ4Ks+bLnnyzImlfRmu+93LePsX7+Q9X72H3/m1TbzrJWcz1N8VUqVVuUKZb9y1lxu2H+DuPUcAuHDdIH/y6mezeWU/192xl8/cspvP/+Jhrr5wDb/7a2cz2JvmD667h3v3HuWdL97Eh195LplU5++6s+p5EppreHjYR0ZGmv68IlK/t/39nfzrwUl+8aGXkUo+Ncwmp4v86fce4Hvb9pNJJXjzpWfx7svPZsUJW7P1GJ0IuHf0KGct7eXc1f10pRr/MssDjx3jfV/fxq5DU5y7qp+rL1zDVRes5qxlfU9abs94li/e9ijfHBklKJTpTidIJRL85esv4FXnr274daPEzO529+G6lo1DcBdKFdJJi+3UnacTd8cdEgmtqzCNTgRc/j9v5r1XPIsPXLn5tMs+PDbFp2/ezf/Ztp9kwrjmkg28+9fPZvXgydsNpXKFe0eP8pMdh/jpgwfZeXBq9mfppHHOqn7OXzvI+WsXc+H6Qc5bPXDKv91Kxbn21kf4y//7IEt6M/zVf7iQyzef+Xj+x4IiX71zDw89Psn7Xv4szh5adMbfibpYBre78/EfPcT4VJ6JbIHxbKF6PZUnWyizrC/D+esGuWDtIOevW8wF6wZZOdCNuzM2lWfPeMCjh7PsGQ8YPRLQlUqwbFEXy/oyLF/UxdK+DEv7MnSlEiQSRsKMpBlm1ZCZLpbJ5ktk82WCQolsoUyuUMLMyCQTpJJGOpkgXbtOJaq3U8kEqYSRSVWvU4nqsqmEkazdtwRk8yUmaztJjk+XZneaTBfL5EuVOddlposVcoUy06Vy7bpCvlgmYUYiAQmrfpAlDDLJBCsGulk92M3KE66X9mUo1J4zV7tM154zX6yQL1XIn3DbDFb0d7FioJtVA9XnWL4oQyqZoFJxsoVSrfYix3MlJrJ59h3J1S7B7O2pfIn+7hSDPWkW96YZ7Kle+rvSdKUTdKUSdKWS1et0goQZxbJTKlcolisUard7M0mesWIRzxhaxNlDfWc8CJC7c+DYNA8dnOShx5+4BIUSAz1pBrrTDPSkGOiu1tOTSZJOJsjU1msmlSSdNCruTBer/27TxQq52jrpSiVZVnsfLVuUYVlfF0trJ5fN5ktM5Uu191D1fbSkL82m5YtYt6SH9Em2eFvpr3/8EH938y5u/S9XsLbOfu+e8Sz/++bdfPuefTiwpDdDf3eK/u4Ui7qq1wB3PDLB0aBIKmFcsmkpV5y7gks3LWPfkYDt+49x375jbN93lOPTJaB61L2rL1zDay9aw7mrBmZf7/Fj03zwW9u4bdc4/+Y5K/nY6y5gSV88TtbbCrEMboDh/3YTqURi9g9jJmwX92TYdyTgvv3H2Hlwkkqt5OWLuggKJYJCefY5kglj9WA3hVKFiWyBUiU+p8fMpJ4ItZ5Mgp50ku4TLpnaH7+7U3Gn7NXb+WKFg5PTPH5smnypMq/XntlAnvvPlTDoy6SYKpQ41VtlUVeKdUt6WLekl3VLehjoSXM8V+RY7XI0KHAsV92rny9VyBcrTJfKp3y+mSDNFctPqmft4h6euWIRS/sysx9G1evqh9pjR3OzYQGwaqCbc1b1M9CTrn3YVD80Z+oq1PlvlUlW18t0qUyx3Pj7KZUw1i/tZdPyPjYu62NR98k/gIb6u9i4rJeNy/pYs7iH5Dz/11IqV3jxx3/KeasH+OI7Lmn490cnAr41MsrYVIGpfImp6WJ1QyNf3dC4+KwlvPzclbxk83IGTjEXuno6sBy3PzLOD7Yf4LZdhylXnM0rF3H1BWtYNdjNf//hDvLFCh+9+jze8IL1T/v/Ucc2uN39jCsvKJTYceA42/cd44HHjtPfnWLjsj7Oqr3h156wdePuHM+VOJytbcVPFSiWK1RqwVepQNkdd6c7naQvk6K3K8mirhS9mRQ9mSSVilOqOMXa1mCx/MTtUtkpVZ54rHrfKVcqtWunVK6+Vl9ti6W/O01/d4qB7hSLutL0pJN0patbfQttL7g7R4Mijx+vhviRoDD7ITAT/k98GJyw1ZtKkEomKFec8Wyeg8fyPH58moO1y+R0qVZzrfae6vWS3gzrlvQw2JNu+I/OvfpvNV0sU6lUP7TSyer/UmaeK18qs2c8YNehKXYdmmL3WPX6+HSR7lSSnkyS7lSS7kyS7lSCof4uzl3Vz+aV/Zy7aoDB3tN/waJScYq19Vco1bb2SxUSCaM79cS/2UyAujuT+RLjUwUmsvnadQEz6OtKVS+ZFH1d1ffSeLbAI4ezPHJ4ikcOZ3l4LMuj41mmi2f+wEgnq2G/cVkfqwe7Gervql4WVf9HtKK/i9WD3Sf9d7/pVwd515dH2PrW50dmxsjhqTw33neA7217jJHajsfz1w7yyTde1BFtjmaIbXCLPF1VKs6hyTyPHM6yZzzLI+NZ9hwOeHQ8y6HJ6obHXJduWspHr34O560ZeNLj7/zSXdy//xj/78NXnHSnZNj2HQl48EB1XvbTYQZIvRoJbk0HFImARMJYNdjNqsFuXviMZU/5ebFcYXyqwNhknrGpaXYenOJzt+zmqr/9BW+6ZAMffMU5LO3LsP9ojp89dIjfe+kzIxnaQK2lFp/ThEWRglskBtLJxGywwyBXnLuSN71gA//rn3fyj7fv4YZfPsYHrtzM+FQBh9ge9U7qc8bgNrP1wJeBlYADW939k60uTEROb7A3zX99zXN486Ub+PMf/Io/u+FXALzkWctjdeJbaVw9W9wl4IPufo+Z9QN3m9lN7v6rFtcmInV41sp+vvzOS/jnHYfY+vPd/MHLnhl2SdJi9Zws+ABwoHZ70sx2AGsBBbdIRJgZV563kivPWxl2KdIGDe29MLONVM/4fkcrihERkTOrO7jNbBHwbeD97v6UcyCZ2RYzGzGzkbGxsWbWKCIiJ6gruM0sTTW0v+ru3znZMu6+1d2H3X14aOjMxxoQEZH5OWNwW/WrWdcCO9z9E60vSURETqeeLe4XA28FrjCzbbXLq1pcl4iInEI9s0puBZ7eR38REYmQaH4nVkRETknBLSISMy05OqCZjQF7ancHgWMn/PjE+zO3T3xsOXB4AS8/9/UaWeZkj5/psTPdDnM8p/pZPevkdLcXMqZ6xnO65Rodz9z7UVtH9T7eye+5ufebnQtxec+d5e71Tcnz2vGoW3WhemyTk96fuT3nsZFmvl4jy5zs8TM9dqbbYY6n3jE1Mp6Fjqme8Sx0HTU6vqi95xa6juL2njvNOPSeO8WlHa2SG05z/4ZTLNPM12tkmZM9fqbH6rm9EAsZz6l+Vs86Od3thaj3eRayjuYzvoVo9nvuZI9HbR018z03936zc6Hj3nMtaZUshJmNeJ0HE4+DThsPdN6YNJ7o67QxLXQ8Udw5uTXsApqs08YDnTcmjSf6Om1MCxpP5La4RUTk9KK4xS0iIqeh4BYRiRkFt4hIzMQquM2sr3bM76vCrqUZzOzZZvZZM7vezN4Tdj0LZWb/1sw+b2bfMLNXhF1PM5jZ2WZ2rZldH3Yt81X7u/mH2rp5c9j1LFQnrJO5Gv7bWcgk8AYmwP89cAi4f87jvwk8BOwCPlzH8/w58CHgqnbU3Y4x1X4nAXylg8azBLi2w9bR9WGPZ75jo3p0z6trt78Rdu3NWldRWydNGlNdfzvtGsDlwMUnDgBIAruBs4EM8EvgPOB84AdzLiuAK4E3Ar8dkeBe8Jhqv/Ma4Ebgmk4YT+33/hq4uFPWUe33IhUSDY7tI8BFtWWuC7v2hY4nquukSWOq62+nnrO8L5i7/7x2vsoTXQLscveHAczs68Br3f0vgKe0QszspUAf1Tdizsx+6O6VVtZ9Os0YU+15vg9838z+CbiudRWfXpPWkQEfA25093taW/GZNWsdRVEjYwP2AeuAbUS0PdrgeGJxovJGxlQ7CXvdfzthrsS1wOgJ9/fVHjspd/9jd38/1XD7fJihfRoNjcnMXmpmnzKzzwE/bHVx89DQeIA/BH4DeL2Z/cdWFrYAja6jZWb2WeB5ZvaRVhe3QKca23eAf29mn6G5h5dotZOOJ2brZK5TraOG/nbassXdTO7+pbBraBZ3/xnws5DLaBp3/xTwqbDraCZ3Hwei+iFUF3fPAu8Iu45m6YR1MlejfzthbnHvB9afcH9d7bE467Qxddp4oDPHNKPTxtZp44EmjSnM4L4LeJaZbTKzDNUdj98PsZ5m6LQxddp4oDPHNKPTxtZp44FmjalNe1e/BhwAilR7Or9Te/xVwE6qe1n/OOy9wE/nMXXaeDp1TJ06tk4bT6vHpINMiYjETCSnBomIyKkpuEVEYkbBLSISMwpuEZGYUXCLiMSMgltEJGYU3CIiMaPgFhGJGQW3iEjM/H9zi88VD63o+AAAAABJRU5ErkJggg==\n"
          },
          "metadata": {
            "needs_background": "light"
          }
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "Doing this every time, is quite a tedious job. So, lets create a learner method. \n",
        "\n",
        "The amazing thing is that you can pass callbacks to our `fit()` method as well. The fit method will automatically remove the callback after training is complete."
      ],
      "metadata": {
        "id": "o4089shufkE4"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "@fc.patch\n",
        "def lr_find(self:Learner, lr_mult=1.3, start_lr=1e-5, max_epochs=10):\n",
        "    lrfind = LRFinderCB(lr_mult=lr_mult)\n",
        "    self.fit(max_epochs, lr=start_lr, cbs=lrfind)\n",
        "    plt.plot(lrfind.lrs, lrfind.losses)\n",
        "    plt.xscale('log')"
      ],
      "metadata": {
        "id": "w-Xcq6xCdVMK"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "learner = MomentumLearner(get_model(), dls, F.cross_entropy, cbs=cbs)\n",
        "learner.lr_find()"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 269
        },
        "id": "6FNuMIU6dY4w",
        "outputId": "af333b21-7e2e-4290-e6df-f7251cb00a67"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "<Figure size 432x288 with 1 Axes>"
            ],
            "image/png": "iVBORw0KGgoAAAANSUhEUgAAAXQAAAD8CAYAAABn919SAAAABHNCSVQICAgIfAhkiAAAAAlwSFlzAAALEgAACxIB0t1+/AAAADh0RVh0U29mdHdhcmUAbWF0cGxvdGxpYiB2ZXJzaW9uMy4yLjIsIGh0dHA6Ly9tYXRwbG90bGliLm9yZy+WH4yJAAAc5klEQVR4nO3de3Bb53nn8e9DgOCdFEVSF0uyaEuOLN9iR7Ib27ltbnVTj7Nt3NRNk9bZpI6zTZvMdDZdNxPPJLOz23amTTeb2aZu0o3XuTnVeryON07HXtubZJzaphxbsizKlhzLulGkCPAGkAQIPPsHQJqmSQkkcQAc5PeZwfDg4BDn0RH4w8GL932PuTsiIhJ+dZUuQERESkOBLiJSIxToIiI1QoEuIlIjFOgiIjVCgS4iUiOildpxd3e39/b2Vmr3IiKhtHfv3jPu3rPYYxUL9N7eXvr6+iq1exGRUDKzo0s9piYXEZEaoUAXEakRCnQRkRqhQBcRqREKdBGRGqFAFxGpEQp0EZEyeviF0xweHA/kuRXoIiJl4u78++/sZc/eE4E8vwJdRKRMJqZnyGSdrpZYIM+vQBcRKZNEMgNApwJdRCTc4qk0AGtb6gN5fgW6iEiZJJL5QO9s1hm6iEioDSdnz9AV6CIioTZ3hq5AFxEJt3gqTbTOaGsIZuZyBbqISJkkkmk6W2KYWSDPX3Sgm1nEzH5hZg8u8titZjZkZs8Wbp8sbZkiIuEXT6YD64MOy7ti0WeBg0D7Eo/f6+6fWX1JIiK1KZFKB9bDBYo8QzezzcBvAt8IrBIRkRoXT6YD6+ECxTe5/B3weSB3lm0+ZGb7zGyPmW1ZbAMzu83M+sysb2hoaLm1ioiEWiKVoTOgQUVQRKCb2Y3AoLvvPctmPwR63f0K4GHg7sU2cve73H23u+/u6Vn0otUiIjUpm3MSqTRrK9zkcj1wk5m9AnwfeLeZfXv+Bu4+7O7ThbvfAHaVtEoRkZAbnczgHlwfdCgi0N39Dnff7O69wC3Ao+7+0fnbmNnGeXdvIv/lqYiIFMQDHiUKy+vl8jpm9mWgz90fAP7UzG4CZoA4cGtpyhMRqQ2JVJUFurs/DjxeWL5z3vo7gDtKWZiISC2JBzwxF2ikqIhIWSTK0OSiQBcRKYPZudB1hi4iEnKJZJqm+ghNsUhg+1Cgi4iUwXDAo0RBgS4iUhb5mRaDGyUKCnQRkbKIpzKBtp+DAl1EpCwSAU+dCwp0EZGymL24RZAU6CIiAUvP5Bifngl0Yi5QoIuIBG4kFezFoWcp0EVEAjZchlGioEAXEQlcogzzuIACXUQkcPEyzLQICnQRkcDNnaFrYJGISLjFkxlATS4iIqGXSKVpb4xSHwk2chXoIiIBi5dhYi5QoIuIBC6RCn6UKCjQRUQCNzyRDnyUKCjQRUQCpzN0EZEa4O5qQxcRqQWTmSzTMzkFuohI2MVn53FRG7qISLglZgcV6QxdRCTcXpvHJdhh/6BAFxEJVLlmWgQFuohIoMo1Fzoo0EVEApVIpqkzaG9Uk4uISKjFU2k6m2PU1Vng+1Kgi4gEKFGmQUWgQBcRCVQ8WZ5h/6BAFxEJVCJVnom5QIEuIhKoeDJTfWfoZhYxs1+Y2YOLPNZgZvea2WEze9LMektZpIhIGLl7/gy9DIOKYHln6J8FDi7x2CeAhLtvB74C/NVqCxMRCbuxyRmyOS/LoCIoMtDNbDPwm8A3ltjkg8DdheU9wHvMLPg+OiIiVey1Yf9VFOjA3wGfB3JLPL4JOAbg7jPAKNC1cCMzu83M+sysb2hoaAXlioiEx+xMi1XThm5mNwKD7r53tTtz97vcfbe77+7p6Vnt04mIVLXZeVy6qiXQgeuBm8zsFeD7wLvN7NsLtjkBbAEwsyjQAQyXsE4RkdCZbXKpmjZ0d7/D3Te7ey9wC/Cou390wWYPAH9YWL65sI2XtFIRkZBJlHFiLoDoSn/RzL4M9Ln7A8A3gXvM7DAQJx/8IiK/0uKpNLFoHc2xSFn2t6xAd/fHgccLy3fOWz8F/E4pCxMRCbtEMj9KtFyd/jRSVEQkIOWcxwUU6CIigYknyzdKFBToIiKBSaQyZevhAgp0EZHAxJPpsvVBBwW6iEggZrI5RifLN9MiKNBFRAIxMpkBytcHHRToIiKBmB1UpDZ0EZGQGy7zKFFQoIuIBEJn6CIiNaLcc6GDAl1EJBBzZ+gaWCQiEm7xZIbWhigN0fJMzAUKdBGRQCRS6bKenYMCXUQkEPHCTIvlpEAXEQlA/gxdgS4iEnrDEzpDFxGpCTpDFxGpAVOZLKl0tqx90EGBLiJScokKDCoCBbqISMnFKzDsHxToIiIll0iWf+pcUKCLiJTca/O4aGCRiEioVWKmRVCgi4iU3HAyjRl0NOkMXUQk1BLJNB1N9UQj5Y1YBbqISInFU+UfJQoKdBGRkksk02Xv4QIKdBGRkosnyz/sHxToIiIll1CTi4hI+Lk7iWRGZ+giImGXTGdJZ3NlH1QECnQRkZKKT1RmUBEo0EVESipeoZkWoYhAN7NGM3vKzJ4zswNm9qVFtrnVzIbM7NnC7ZPBlCsiUt1mh/1XItCjRWwzDbzb3SfMrB74mZk95O7/umC7e939M6UvUUQkPOLVHOju7sBE4W594eZBFiUiElazF7eo2l4uZhYxs2eBQeBhd39ykc0+ZGb7zGyPmW1Z4nluM7M+M+sbGhpaRdkiItUpnkwTrTPaGoppACmtogLd3bPufiWwGbjGzC5bsMkPgV53vwJ4GLh7iee5y913u/vunp6e1dQtIlKVZi8ObWZl3/eyerm4+wjwGHDDgvXD7j5duPsNYFdpyhMRCZd4sjKjRKG4Xi49ZramsNwEvA/oX7DNxnl3bwIOlrJIEZGwyM/jUv5BRVBcL5eNwN1mFiH/BvADd3/QzL4M9Ln7A8CfmtlNwAwQB24NqmARkWoWT6bZsaGtIvsuppfLPuCqRdbfOW/5DuCO0pYmIhI+iVSmIl0WQSNFRURKJptzRio00yIo0EVESmZsMkPOK9MHHRToIiIlU8l5XECBLiJSMrPzuFRipkVQoIuIlMxwBedxAQW6iEjJzJ2hK9BFRMJtrg1dTS4iIuGWSKZpqo/QFItUZP8KdBGREoknKzeoCBToIiIlk59psTLzuIACXUSkZOLJdMW6LIICXUSkZBKptJpcRERqQXxCZ+giIqGXnskxPj2jM3QRkbAbqeDFoWcp0EVESmB2UFGXAl1EJNziFZ6YCxToIiIlkUhmgMpNzAUKdBGRkojPtaFrYJGISKhVei50UKCLiJREPJmmrTFKfaRysapAFxEpgXiysqNEQYEuIlISlR72Dwp0EZGSiCfTFbuwxSwFuohICSSS6YqOEgUFuohIScTV5CIiEn6T6SxTmVxFuyyCAl1EZNUGx6eAys7jAgp0EZFVe/H0BADb17dWtA4FuojIKvWfGgPgTevbKlqHAl1EZJX6T4+zZW0TrQ3RitahQBcRWaX+U2NcvKG90mUo0EVEVmMqk+WV4RQXb6hscwsUEehm1mhmT5nZc2Z2wMy+tMg2DWZ2r5kdNrMnzaw3iGJFRKrN4cEJsjkPzRn6NPBud38zcCVwg5m9dcE2nwAS7r4d+ArwV6UtU0SkOvUPjAOwIwxn6J43UbhbX7j5gs0+CNxdWN4DvMfMrGRViohUqUMDY8SidfR2NVe6lOLa0M0sYmbPAoPAw+7+5IJNNgHHANx9BhgFuhZ5ntvMrM/M+oaGhlZXuYhIFegfGOdN61uJVnAe9FlFVeDuWXe/EtgMXGNml61kZ+5+l7vvdvfdPT09K3kKEZGq0j8wzo71lW8/h2X2cnH3EeAx4IYFD50AtgCYWRToAIZLUaCISLUanphmaHyanRsr334OxfVy6TGzNYXlJuB9QP+CzR4A/rCwfDPwqLsvbGcXEakph6roC1GAYoY1bQTuNrMI+TeAH7j7g2b2ZaDP3R8AvgncY2aHgThwS2AVi4hUiWrq4QJFBLq77wOuWmT9nfOWp4DfKW1pIiLVrX9gjK6WGD2tDZUuBdBIURGRFTs0MM6ODW1USy9tBbqIyApkc86h0+NVMUJ0lgJdRGQFXo2nmMrkqmIOl1kKdBGRFTg0kJ8D/eIq6bIICnQRkRU5eGocM7honQJdRCTUDg2M09vVQlMsUulS5ijQRURWoH9grKraz0GBLiKybKn0DEfjqaoZUDRLgS4iskwvnZ7AnarqsggKdBGRZeuf7eGiM3QRkXDrHxinqT7C+Wsrf1GL+RToIiLL1H9qnDdtaKOurjqG/M9SoIuILIN7Ycj/+upqbgEFuojIsgxNTBNPpqtqhOgsBbqIyDL0n6quOdDnU6CLiCzD7FWKqq3LIijQRUSW5eDAGOvaGljbEqt0KW+gQBcRWYbZi1pUIwW6iEiRZrI5XhqcYOfG6mtuAQW6iEjRXhlOkp7JsaMKuyyCAl1EpGj9A9XbwwUU6CIiRes/NU6kzti+rrXSpSxKgS4iUqT+gXEu6G6hsb56LmoxnwJdRKRIh05X30Ut5lOgi4gUYWJ6hmPxSQW6iEjYVfMI0VkKdBGRIsxe1KJae7iAAl1EpCiHBsZpbYiyubOp0qUsSYEuIlKE/sKQf7PquqjFfAp0EZFzcHf6T41VdXMLKNBFRM5pYGyKsakZdirQRUTC7bWLWlRvDxdQoIuInNPcHC5VOinXrHMGupltMbPHzOwFMztgZp9dZJt3mdmomT1buN0ZTLkiIuXXPzDGeR2NdDTXV7qUs4oWsc0M8Gfu/oyZtQF7zexhd39hwXY/dfcbS1+iiEhlVfNFLeY75xm6u59y92cKy+PAQWBT0IWJiFSD9EyOI0MTXFylF7WYb1lt6GbWC1wFPLnIw9ea2XNm9pCZXbrE799mZn1m1jc0NLTsYkVEyu3lMxNksl7Vc7jMKjrQzawV+F/A59x9bMHDzwBb3f3NwH8D7l/sOdz9Lnff7e67e3p6VlqziEjZHKryi1rMV1Sgm1k9+TD/jrvft/Bxdx9z94nC8o+AejPrLmmlIiIVcPDUOPUR48Lu6ryoxXzF9HIx4JvAQXf/2yW22VDYDjO7pvC8w6UsVESkEg4NjLGtp5VYtPp7eRfTy+V64GPAfjN7trDuL4DzAdz968DNwKfNbAaYBG5xdw+gXhGRkhibytAai1JXd/a5WQ4NjHPNBWvLVNXqnDPQ3f1nwFn/xe7+NeBrpSpKRCQoqfQMf3Hffu5/9iTROqOnrYF17Y1saG9gfXvjvFsDbY31nBydqvoRorOKOUMXEakJR4eTfOqevRw6Pc7Hr++lORZhYHSawfEpXh5K8vMjw4xNzbzh93ZurP4vREGBLiIVcmZimjvu28/5a5v5wgd2nrPpY7Ue6x/ks9//BWbGtz5+De980+I97SbTWU6PTXF6bIqBsSmmMznetj0cfTwU6CJSdvuOj/Cpe/YyOD5NNueMpDL89c1XEAkg1HM552uPHeYrj7zIzg3tfP2juzi/q3nJ7ZtiEXq7W+jtbil5LUFToItIWf1z3zG+cP/z9LQ28L//+Hr+78FBvvLIi6SzOf72w2+mPlK63iSjkxn+7AfP8sjBQX7rqk3859+6nKZYpGTPX20U6CJSFplsjv/04Avc/fOjXLeti6995C2sbYlx2aYOGurr+MuH+knPZPnq711FQ3T1oXtoYJzbv72XY/EUX7rpUv7g2q1VfbWhUlCgi0jghsan+ePvPMNTr8T5o7dfwJ/fcDHReWfit79zGw3ROr70wxe4/Z69/P1Hd9FYv/JQf3DfST6/Zx8tDVG+d9tbubo3HN0OV0uBLiKB+sWrCT797WcYmUzzX2+5kg9eufjcfh+//gIaohG+cP9+Pnl3H3f9wS6aY8uLqJlsjr/+l0Pc9ZOX2bW1k//++29hfXtjKf4ZoaBAF5nH3TkxMglASyxKUyxCQ7Su5j+qB+Xep1/li/cfYH1HA/d9+nouOe/s/bk/8mvnE4vW8fk9z3Hr/3iaf7r1alobzh1Tp8em+Oe+Y3z/6WMcT0zysbdu5Ys3XhKK0Z2lVFOBHk+m2X9ilOdPjLL/+Cgvnh5nJufUGdSZQeHn3H3yPyN1xprm+vwAg7ZG1rU1FJbzAw7WtTXQssSLyt1JZ3NMZXJMZbJMprNMzWRZ0xRjfXtDoEGwcN9TmSypdJaRVIaRVJpEKkMilZ5bnr8uNZ1l19ZObrhsA//m4nVF/dGUw/hUhqd+GeeJI8P88kyS2QHHs8OO3ecv55eaYxGu7l3Lddu6uXhD27K7v2WyOZ76ZZxHDp7mkYOnORaffN3jkTqjORahJRaluSFCcyxCcyxKa0OUi9a1srt3Lbu2drK2Jbaaf3pNSc/k+NIPD/CdJ1/l7Rd189VbrqKzyONz867NNETr+Ny9z/Kxbz7Jtz5+DR1Nb7ywRDbn/OTFIb771Ks82j9INudct62LO2+8hPdfuqHU/6RQsEqN0N+9e7f39fWt+PcXhvf+E6NzZ1YAW7ua2bmhnaZYhJw7OYdcIQ3y9/Pr3J1szomnMpwZzw8wyGTfeExaYhHWtTdSZ7wW3oUQzS1xCLtbG7hicweXbergik0dXLG5g3VFfPybTGc5MjTBS4PjvHR6gsODExxPTM6Fdn6/OaZmshTz39dYX0dnc4w1zTE6m+vpbI5RHzF+dvgMZybSxKJ1vOOiHn7jsg28d+f6oq/Kkss5rwwnOXByjBdOjVFfZ1zQ08IF3a1c0N2y6B/hYv/WvUcTPHHkDE8cGWb/iVGyOScWrWN7TyuROmP2PXEups3mls3yr4WjwykA1rbEuPbCLq7d1sX127vp7Wpe9E11JJXm8UNDPHLwNP/v0BDj0zPEonW8bXs379rRQ2N9hNT0DMl0llR6huR0/s06mZ4hlc6SnJ5hbGqGw4Pjc6+XbT0tXF0I96t717J1iX3Pmp7JcmpkihMjk5xITHJ8ZBJ357w1TWxa08SmzibO62haUa+MXM4Zncy/eY9NzTA2mWF8aoaxqcyiy9MzObava+Xywuv0wsKxL5a7c3Q4xb4To+w7NsJPXzrDodPj3P7ObfyHX9+xou6IP35+gD/53jPs2NDGPf/u1+beEE6NTvKDp49z79OvcnJ0iu7WGDfv2sItV28JZVfD5TKzve6+e9HHwhboj/af5ov3H3hDeM+G5uWbOrh0U0dRYbIY93yf2MFCuA+NT+eXx/L33aGxPkJjfR1N9REa6yNzH8ubYpG5dYNjU+w/Mcb+EyMcHpyYC/11bfmQv3zTGi7f3M6a5hiHByfmbi8NjnM8MTkX1NE6o7e7ha1rm2mKFfZX2H9jYV8N0deWm+ojdDbX58O7JR/eS325lM05e48meOj5U/zL8wOcHJ0iWmdct72b37hsA++/ZD1drQ1A/iz2pdMTPH9ylBdOjnGg8DOZzgJQHzGyOX/dm1t3a4wLulsKt3zIX9jTwthkhieODPPEkTM8c3SEdDZHtM5485Y1XLctH8ZvOb9zWV+KnRyZnHvOJw4PMzA2BcDGjkau3dbFddu62bG+jX99eZhHDp6m72iCbM7pbo3xnovX856d63jbRd3LbrOdymTZd3yUvqNx+l5JsPdogtHJTOHf38DurZ3s7u2kIVrH8UJwzwb40MT0696QCx8i33CC0NUSmwv3TZ35sG9tjDKayhBPpUkk08STaRKp2Z/5T2NLnWgA1Bm0N9XT1hilvbGeaJ3x4ukJJjP5/8/mWIRLz2vn8k1r5k5KLuxuoa7OcHdOjk6x//gI+46PFm4jcyMsY9E6LtnYzm3vuJAPXL5xWcdzocf6B/nUt/dyYXcLn3n3du7/xQke7R8k5/D2i7r5vWvO57071/9KNa3UVKDvPz7K139ypCThXS6p9AwvnBxj3/H8J4p9J0Y5MjTxuj/mWLSObT2tbF/XykWzt/WtbO1qKWm/3KW4O88dH+Wh50/x4+cHODqcos5g19ZOJjNZXhyYIJ3NAfk/9ks2tnPpee1cel4Hl25q56J1bTjOsXiKl4eS/PJM/vZy4efQ+PTr9mcGl57XznXburl2WxdX964tWbOPu/PLM0meODLMzwshn0hl5h6/eEMb79m5jvfuXM+bN68p6QjFXM45PDTB06/E2ftKgqePxueacGKROjauacyffa95LZw3dTaxeU0zGzrynwAHxqY4kZjk5Oi8N4CRKU4kUpwYmWQqk5vbX33E6GyOsbYl9trPlnrWNsfobImxprmejqZ62hvraWusp70pSltjPS2xyBs+PWRzzpGhiddep8dHeOHU2Nz+WhuibF/XyvFEijMTaSB/wrFjQxtXbM4H/xWbO3jT+raSvmZ/9tIZPvk/n2Yqk6OnrYEP797M7+4+/6yDg2pZTQV6rZiYnuHAiVHGp2bYvq6VLWubAxkltxLuzsFT4/z4wACP9Q/S0VTPpZsK4X1eO71dLcuudXwqwytnUrx8ZoKGaIS3XriWNc3laXPO5Zz+gXFePD3Orq2dbFlb3iAYHJvCgZ7WhlW/ebg78WSaVDpLZ0ts0WAupZlsjsNDE3PNmi+eHmdzZ3MhvNdw8Ya2VXUvLNaBk6OcGpninTt6ynKCU80U6CIiNeJsgf6r/VYnIlJDFOgiIjVCgS4iUiMU6CIiNUKBLiJSIxToIiI1QoEuIlIjFOgiIjWiYgOLzGwUeKlwtwMYXbA8+7MbOLOCXcx/zmIfX7jubPdLXfNK6l1OzdV6jBeuC/MxXqxOHePiHtcxLt5Wd1/8CtfuXpEbcNfZluf97Fvt8xf7+MJ1Z7tf6ppXUu9yaq7WY1zMcQ3LMV6iTh1jHeOS/+0tdatkk8sPz7E8f91qn7/YxxeuO9v9Ute8knoXW7/U/Wo9xgvXhfkYz1/WMV7e4zrGJVCxJpdimVmfLzFvQbUKW81hqxfCV3PY6oXw1Ry2eqH0NYfhS9G7Kl3ACoSt5rDVC+GrOWz1QvhqDlu9UOKaq/4MXUREihOGM3QRESmCAl1EpEYo0EVEakSoA93M3mVmPzWzr5vZuypdT7HMrMXM+szsxkrXci5mtrNwfPeY2acrXU8xzOzfmtk/mtm9Zvb+StdzLmZ2oZl908z2VLqWsym8bu8uHNvfr3Q95xKW4zrfal+7FQt0M/snMxs0s+cXrL/BzA6Z2WEz+4/neBoHJoBG4HhQtc6rrRQ1A/w58INgqnxdXauu190PuvvtwIeB64Ost1BbKWq+393/CLgd+N0Q1Puyu38iyDqXssz6fxvYUzi2N5W9WJZXbyWP64LallPz6l67pRyltMzRWe8A3gI8P29dBDgCXAjEgOeAS4DLgQcX3NYBdYXfWw98JyQ1vw+4BbgVuLHa6y38zk3AQ8BHwnCM5/3e3wBvCVG9e4I+vqus/w7gysI23y13rcutt5LHtQQ1r+i1G6VC3P0nZta7YPU1wGF3fxnAzL4PfNDd/wtwtuaJBNAQRJ3zlaLmQtNQC/k/kEkz+5G756q13sLzPAA8YGb/B/huELXO21cpjrEBfwk85O7PVHu9lbSc+sl/Ct4MPEuFPt0vs94Xylvd4pZTs5kdZBWv3WprQ98EHJt3/3hh3aLM7LfN7B+Ae4CvBVzbUpZVs7t/wd0/Rz4Y/zGoMD+L5R7jd5nZVwvH+UdBF7eEZdUM/AnwXuBmM7s9yMKWsNxj3GVmXweuMrM7gi6uCEvVfx/wITP7ewIaur5Ci9Zbhcd1vqWO8apeuxU7Qy8Fd7+P/IssdNz9W5WuoRju/jjweIXLWBZ3/yrw1UrXUSx3HybfZlrV3D0JfLzSdRQrLMd1vtW+dqvtDP0EsGXe/c2FddUsbDWHrV4IX81hq3ehsNUftnohoJqrLdCfBi4yswvMLEb+y8MHKlzTuYSt5rDVC+GrOWz1LhS2+sNWLwRVcwW/+f0ecArIkG8/+kRh/QeAF8l/A/yFSn47Hfaaw1ZvGGsOW71hrz9s9Za7Zk3OJSJSI6qtyUVERFZIgS4iUiMU6CIiNUKBLiJSIxToIiI1QoEuIlIjFOgiIjVCgS4iUiMU6CIiNeL/A/uusJBJ9nv0AAAAAElFTkSuQmCC\n"
          },
          "metadata": {
            "needs_background": "light"
          }
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "Now we are talking. This is such a handy API to quickly find the right learning rate.\n",
        "\n",
        "By this time, you should have realized how powerful these callbacks can get. I know this is very new for a lot of people and will take some getting used to. To build your muscle try reading the `MetricsCB` that we have defined at the very beginning of this notebook.\n",
        "\n",
        "Another thing that one can try is to go through the Fastai Callbacks docs and try reading them all. This will give you a pretty good idea about all the different things that can be accomplished using callbacks. You can even try reading the source code. The syntax is a bit different but still it would be good practice. "
      ],
      "metadata": {
        "id": "LpGH5TjjlK81"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Order of Callbacks"
      ],
      "metadata": {
        "id": "DiB66_W2ljft"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "One caveat when using multiple callbacks is the order of execution. For all particular events like `after_loss` all the callbacks that define these methods will be executed. So, if there are 4 callbacks that define the `after_loss` method then all of them will be executed. But what if one of them raises `CancelFitException`? Its will cancel the complete training and other methods will not be executed. Hence, the order in which they are executed is really important.  \n",
        "\n",
        "Here are examples to substantiate the idea. For example, you want to use `SingleEpochCB` and `MetricsCB` together. Both of them define the `after_epoch` method. The `after_epoch` method in `SingleEpochCB` simply raises `CancelFitException`. Its really important that `MetricsCB` is executed before `SingleEpochCB`. Else, loss and other metrics will not be aggregated or logged. \n",
        "\n",
        "Here is an example, first we will execute `SingleEpochCB`"
      ],
      "metadata": {
        "id": "XLgqI4sR0vwq"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "cbs = [SingleEpochCB(), metrics, DeviceCB()] # important\n",
        "learner = MomentumLearner(get_model(), dls, loss_func=loss_func, cbs=cbs)\n",
        "learner.fit(3)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "6rirZxS23reU",
        "outputId": "e7c19b40-f2bb-4d56-fd72-c0c2021f1ca1"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "{'accuracy': '0.584', 'loss': '1.111', 'epoch': 0, 'train': 'train'}\n",
            "[SingleEpochCB] Terminating . . . \n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "Here, only training metrics are logged because in validation loop `SingleEpochCB` callback is executed first that terminates the training. Next, we will change the order of callbacks"
      ],
      "metadata": {
        "id": "-NvUnlUK71qC"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "cbs = [metrics, SingleEpochCB(), DeviceCB()] # update\n",
        "learner = MomentumLearner(get_model(), dls, loss_func=loss_func, cbs=cbs)\n",
        "learner.fit(3)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "CU8g3OyK75iF",
        "outputId": "ed0ca9ba-c7cc-4f7d-aee9-45b432b68059"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "{'accuracy': '0.584', 'loss': '1.111', 'epoch': 0, 'train': 'train'}\n",
            "{'accuracy': '0.744', 'loss': '0.726', 'epoch': 0, 'train': 'eval'}\n",
            "[SingleEpochCB] Terminating . . . \n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "As you can see both training and validation metrics are logged. Validation metrics are only logged when `MetricsCB` is executed before `SingleEpochCB()`. But this approach of changing the callbacks order in the list is not very useful if you don't know the internal workings of the callback. Also, it introduces an element of uncertainty in the code. So, we need a better mechanism to better handle the order of execution.\n",
        "\n",
        "If you might have noticed, the `Callback` class has an attribute called `order`. `run_cbs` function sorts all the callbacks based on this `order` attribute. So, if you want some callback to be executed before others then make sure its `order` is less than others. And if you want to execute last, then make its order really large.\n",
        "\n",
        "Lets update our `SingleEpochCB` so that its executed last everytime."
      ],
      "metadata": {
        "id": "DZ_OUl797-Y1"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "class SingleEpochCB(Callback):\n",
        "    order = Callback.order + 1  # update\n",
        "    def after_epoch(self, learner): \n",
        "        if not learner.model.training:\n",
        "            print('[SingleEpochCB] Terminating . . . ') \n",
        "            raise CancelFitException()"
      ],
      "metadata": {
        "id": "kXZYVc2y9PhZ"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "After this single line of change, no matter the order of `SingleEpochCB` in the list, it will always be executed after `MetricsCB`."
      ],
      "metadata": {
        "id": "-1hvE4Cw9Zx1"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "cbs = [SingleEpochCB(), metrics, DeviceCB()] # important\n",
        "learner = MomentumLearner(get_model(), dls, loss_func=loss_func, cbs=cbs)\n",
        "learner.fit(3)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "1lZ402JV9qyT",
        "outputId": "8ad7b492-0203-4fee-f772-4818d117fc58"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "{'accuracy': '0.519', 'loss': '1.297', 'epoch': 0, 'train': 'train'}\n",
            "{'accuracy': '0.719', 'loss': '0.767', 'epoch': 0, 'train': 'eval'}\n",
            "[SingleEpochCB] Terminating . . . \n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "Even though `SingleEpochCB` comes before `MetricsCB` in the `cbs` list, it is executed after `MetricsCB`. \n",
        "\n",
        "**Note:** I must admit, this is the most error prone aspect of this approach. Even if there is an error, its difficult to detect and debug because of multiple levels of abstraction. Often I add `print` statements to check if things are working as expected.\n",
        "\n",
        "Generally the idea is to execute the callbacks that raise exceptions towards the end. But there can be some scenarios where this is not true. So, it  depends on the use-case."
      ],
      "metadata": {
        "id": "DyE6Lc4n30jk"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Summary\n",
        "\n",
        "Here are some of the most important takeaways:\n",
        "\n",
        "- Base `Learner` class doesn't have `predict`, `get_loss`, `backward`, `step`, and `zero_grad` methods defined. You can implement these methods either by inheriting the `Learner` class or by using callbacks.\n",
        "- Using callbacks, you can inject arbitrary code at various points inside the training loop.\n",
        "- The order of callbacks is extremely important when passing multiple callbacks.\n",
        "- You can also pass callbacks to the `fit()` method. These callbacks are automatically removed after that training loop.\n",
        "\n",
        "\n",
        "You can accomplish much more when you combine Callbacks with Pytorch Hooks. In the next blog, we will be discussing Hooks in detail and will write some amazing callbacks. Until next time 👋"
      ],
      "metadata": {
        "id": "Db_aU9c8ldyV"
      }
    }
  ]
}