examples/early_stopping.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 🛑 Early Stopping\n",
    "\n",
    "In composer, Callbacks modify trainer behavior and are called at the relevant Events in the training loop. This tutorial focuses on two callbacks, the [EarlyStopper](https://docs.mosaicml.com/en/stable/api_reference/composer.callbacks.early_stopper.html#composer.callbacks.early_stopper.EarlyStopper) and [ThresholdStopper](https://docs.mosaicml.com/en/stable/api_reference/composer.callbacks.threshold_stopper.html#composer.callbacks.threshold_stopper.ThresholdStopper), both of which halt training early depending on different criteria."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Setup\n",
    "\n",
    "In this tutorial, we'll train a ComposerModel and halt training for criteria that we'll set. We'll use the same model as in the [Getting Started](https://docs.mosaicml.com/en/stable/examples/getting_started.html) tutorial."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Install Composer\n",
    "\n",
    "First, install Composer if you haven't already:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%pip install mosaicml"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Seed\n",
    "\n",
    "Next, we'll set the seed for reproducibility:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from composer.utils.reproducibility import seed_all\n",
    "\n",
    "seed_all(42)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Dataloader Setup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch.utils.data\n",
    "from torchvision import datasets, transforms\n",
    "\n",
    "data_directory = \"./data\"\n",
    "\n",
    "# Normalization constants\n",
    "mean = (0.507, 0.487, 0.441)\n",
    "std = (0.267, 0.256, 0.276)\n",
    "\n",
    "batch_size = 1024\n",
    "\n",
    "cifar10_transforms = transforms.Compose([transforms.ToTensor(), transforms.Normalize(mean, std)])\n",
    "\n",
    "train_dataset = datasets.CIFAR10(data_directory, train=True, download=True, transform=cifar10_transforms)\n",
    "eval_dataset = datasets.CIFAR10(data_directory, train=False, download=True, transform=cifar10_transforms)\n",
    "\n",
    "# Setting shuffle=False to allow for easy overfitting in this example\n",
    "train_dataloader = torch.utils.data.DataLoader(train_dataset, batch_size=batch_size, shuffle=False)\n",
    "eval_dataloader = torch.utils.data.DataLoader(eval_dataset, batch_size=batch_size, shuffle=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Model, Optimizer, Scheduler, and Evaluator Setup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from composer import models\n",
    "from composer.optim import DecoupledSGDW, LinearWithWarmupScheduler\n",
    "from composer.core import Evaluator\n",
    "from torchmetrics.classification.accuracy import Accuracy\n",
    "\n",
    "\n",
    "model = models.composer_resnet_cifar(model_name='resnet_56', num_classes=10)\n",
    "\n",
    "optimizer = DecoupledSGDW(\n",
    "    model.parameters(), # Model parameters to update\n",
    "    lr=0.05, # Peak learning rate\n",
    "    momentum=0.9,\n",
    "    weight_decay=2.0e-3 # If this looks large, it's because its not scaled by the LR as in non-decoupled weight decay\n",
    ")\n",
    "\n",
    "lr_scheduler = LinearWithWarmupScheduler(\n",
    "    t_warmup=\"1ep\", # Warm up over 1 epoch\n",
    "    alpha_i=1.0, # Flat LR schedule achieved by having alpha_i == alpha_f\n",
    "    alpha_f=1.0\n",
    ")\n",
    "\n",
    "evaluator = Evaluator(\n",
    "    dataloader = eval_dataloader,\n",
    "    label = \"eval\",\n",
    "    metrics = Accuracy()\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## EarlyStopper"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The `EarlyStopper` callback tracks a particular training or evaluation metric and stops training if the metric does not improve within a given time interval. \n",
    "\n",
    "\n",
    "The callback takes the following parameters:\n",
    "\n",
    "-   `monitor`: The name of the metric to track\n",
    "-   `dataloader_label`: The `dataloader_label` identifies which dataloader the metric belongs to. By default, the train dataloader is labeled `train`, and the evaluation dataloader is labeled `eval`. (These names can be customized via the `train_dataloader_label` in the [Trainer](https://docs.mosaicml.com/en/stable/api_reference/composer.trainer.trainer.html#composer.trainer.trainer.Trainer) or the `label` argument of the [Evaluator](https://docs.mosaicml.com/en/stable/api_reference/composer.core.evaluator.html?highlight=Evaluator#composer.core.evaluator.Evaluator), respectively.)\n",
    "-   `patience`: The interval of the time that the callback will wait before stopping training if the metric is not improving. You can use integers to specify the number of epochs or provide a [Time](https://docs.mosaicml.com/en/stable/trainer/time.html) string -- e.g., \"50ba\" or \"2ep\" for 50 batches and 2 epochs, respectively.\n",
    "-   `min_delta`: If non-zero, the change in the tracked metric over the `patience` window must be at least this large.\n",
    "-   `comp`: A comparison operator can be provided to measure the change in the monitored metric. The comparison operator will be called like `comp(current_value, previous_best)`\n",
    "\n",
    "See the [API Reference](https://docs.mosaicml.com/en/stable/api_reference/composer.callbacks.early_stopper.html#composer.callbacks.early_stopper.EarlyStopper) for more information.\n",
    "\n",
    "Here, we'll use it track the Accuracy metric over one epoch on the test dataset:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from composer.callbacks import EarlyStopper\n",
    "\n",
    "early_stopper = EarlyStopper(monitor=\"Accuracy\", dataloader_label=\"eval\", patience=1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now that we have our callback, we can instantiate the Trainer and train:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from composer.trainer import Trainer\n",
    "\n",
    "# Early stopping should stop training before we reach 100 epochs!\n",
    "train_epochs = \"100ep\"\n",
    "\n",
    "trainer = Trainer(\n",
    "    model=model,\n",
    "    train_dataloader=train_dataloader,\n",
    "    eval_dataloader=evaluator,\n",
    "    max_duration=train_epochs,\n",
    "    optimizers=optimizer,\n",
    "    schedulers=lr_scheduler,\n",
    "    callbacks=[early_stopper],\n",
    "    train_subset_num_batches=10,  # only training on a subset of the data to trigger the callback sooner\n",
    ")\n",
    "\n",
    "# Train!\n",
    "trainer.fit()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## ThresholdStopper"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The [ThresholdStopper](https://docs.mosaicml.com/en/stable/api_reference/composer.callbacks.threshold_stopper.html#composer.callbacks.threshold_stopper.ThresholdStopper) callback is similar to the EarlyStopper, but it halts training when the metric crosses a threshold set in the ThresholdStopper callback.\n",
    "\n",
    "This callback takes the following parameters:\n",
    "-   `monitor`, `dataloader_label`, and `comp`: Same as the EarlyStopper callback\n",
    "-   `threshold`: The float threshold that dictates when the halt training.\n",
    "-   `stop_on_batch`: If True, training will halt in the middle of a batch if the training metrics satisfy the threshold."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We will reuse the same setup for the ThresholdStopper example."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from composer.callbacks import ThresholdStopper\n",
    "\n",
    "threshold_stopper = ThresholdStopper(\"Accuracy\", \"eval\", threshold=0.3)\n",
    "\n",
    "# Threshold stopping should stop training before we reach 100 epochs!\n",
    "train_epochs = \"100ep\"\n",
    "\n",
    "trainer = Trainer(\n",
    "    model=model,\n",
    "    train_dataloader=train_dataloader,\n",
    "    eval_dataloader=evaluator,\n",
    "    max_duration=train_epochs,\n",
    "    optimizers=optimizer,\n",
    "    schedulers=lr_scheduler,\n",
    "    callbacks=[threshold_stopper],\n",
    "    train_subset_num_batches=10,  # only training on a subset of the data to trigger the callback sooner\n",
    ")\n",
    "\n",
    "# Train!\n",
    "trainer.fit()"
   ]
  }
 ],
 "metadata": {
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}