chore: import upstream snapshot with attribution

doc/source/train/examples/pytorch/BUILD.bazel

@@ -0,0 +1,43 @@
+load("//bazel:python.bzl", "py_test_run_all_notebooks")
+
+filegroup(
+    name = "train_pytorch_examples",
+    srcs = glob(["*.ipynb"]),
+    visibility = ["//doc:__subpackages__"],
+)
+
+# GPU Tests
+py_test_run_all_notebooks(
+    size = "large",
+    include = ["*.ipynb"],
+    data = ["//doc/source/train/examples/pytorch:train_pytorch_examples"],
+    exclude = ["convert_existing_pytorch_code_to_ray_train.ipynb"],  # CPU test
+    tags = [
+        "exclusive",
+        "gpu",
+        "ray_air",
+        "team:ml",
+    ],
+)
+
+# CPU Tests
+py_test_run_all_notebooks(
+    size = "large",
+    include = ["convert_existing_pytorch_code_to_ray_train.ipynb"],
+    data = ["//doc/source/train/examples/pytorch:train_pytorch_examples"],
+    exclude = [],
+    tags = [
+        "exclusive",
+        "ray_air",
+        "team:ml",
+    ],
+)
+
+filegroup(
+    name = "pytorch_examples_ci_configs",
+    srcs = glob([
+        "**/ci/aws.yaml",
+        "**/ci/gce.yaml",
+    ]),
+    visibility = ["//doc/source/train/examples:__pkg__"],
+)

doc/source/train/examples/pytorch/dreambooth_finetuning.rst

@@ -0,0 +1,304 @@
+:orphan:
+
+Fine-tuning of Stable Diffusion with DreamBooth and Ray Train
+=============================================================
+
+.. raw:: html
+
+    <a id="try-anyscale-quickstart-dreambooth_finetuning" target="_blank" href="https://console.anyscale.com/register/ha?render_flow=ray&utm_source=ray_docs&utm_medium=docs&utm_campaign=dreambooth_finetuning">
+      <img src="../../../_static/img/run-on-anyscale.svg" alt="Run on Anyscale" />
+      <br/><br/>
+    </a>
+
+This is an intermediate example that shows how to do DreamBooth fine-tuning of a Stable Diffusion model using Ray Train.
+It demonstrates how to use :ref:`Ray Data <data>` with PyTorch Lightning in Ray Train.
+
+
+See the original `DreamBooth project homepage <https://dreambooth.github.io/>`_ for more details on what this fine-tuning method achieves.
+
+.. image:: https://dreambooth.github.io/DreamBooth_files/high_level.png
+  :target: https://dreambooth.github.io
+  :alt: DreamBooth fine-tuning overview
+
+This example builds on `this Hugging Face 🤗 tutorial <https://huggingface.co/docs/diffusers/training/dreambooth>`_.
+See the Hugging Face tutorial for useful explanations and suggestions on hyperparameters.
+**Adapting this example to Ray Train allows you to easily scale up the fine-tuning to an arbitrary number of distributed training workers.**
+
+**Compute requirements:**
+
+* Because of the large model sizes, you need a machine with at least 1 A10G GPU.
+* Each training worker uses 1 GPU. You can use multiple GPUs or workers to leverage data-parallel training to speed up training time.
+
+This example fine-tunes both the ``text_encoder`` and ``unet`` models used in the stable diffusion process, with respect to a prior preserving loss.
+
+
+.. image:: /templates/05_dreambooth_finetuning/dreambooth/images/dreambooth_example.png
+   :alt: DreamBooth overview
+
+Find the full code repository at `https://github.com/ray-project/ray/tree/master/doc/source/templates/05_dreambooth_finetuning <https://github.com/ray-project/ray/tree/master/doc/source/templates/05_dreambooth_finetuning>`_
+
+
+How it works
+------------
+
+This example uses Ray Data for data loading and Ray Train for distributed training.
+
+Data loading
+^^^^^^^^^^^^
+
+.. note::
+    Find the latest version of the code at `dataset.py <https://github.com/ray-project/ray/blob/master/doc/source/templates/05_dreambooth_finetuning/dreambooth/dataset.py>`_
+
+    The latest version might differ slightly from the code presented here.
+
+
+Use Ray Data for data loading. The code has three interesting parts.
+
+First, load two datasets using :func:`ray.data.read_images`:
+
+.. literalinclude:: /templates/05_dreambooth_finetuning/dreambooth/dataset.py
+  :language: python
+  :start-at: instance_dataset = read
+  :end-at: class_dataset = read
+  :dedent: 4
+
+Then, tokenize the prompt that generated these images:
+
+.. literalinclude:: /templates/05_dreambooth_finetuning/dreambooth/dataset.py
+  :language: python
+  :start-at: tokenizer = AutoTokenizer
+  :end-at: instance_prompt_ids = _tokenize
+  :dedent: 4
+
+
+And lastly, apply a ``torchvision`` preprocessing pipeline to the images:
+
+.. literalinclude:: /templates/05_dreambooth_finetuning/dreambooth/dataset.py
+  :language: python
+  :start-after: START: image preprocessing
+  :end-before: END: image preprocessing
+  :dedent: 4
+
+Apply all three parts in a final step:
+
+.. literalinclude:: /templates/05_dreambooth_finetuning/dreambooth/dataset.py
+  :language: python
+  :start-after: START: Apply preprocessing
+  :end-before: END: Apply preprocessing
+  :dedent: 4
+
+
+Distributed training
+^^^^^^^^^^^^^^^^^^^^
+
+
+.. note::
+    Find the latest version of the code at `train.py <https://github.com/ray-project/ray/blob/master/doc/source/templates/05_dreambooth_finetuning/dreambooth/train.py>`_
+
+    The latest version might differ slightly from the code presented here.
+
+
+The central part of the training code is the :ref:`training function <train-overview-training-function>`. This function accepts a configuration dict that contains the hyperparameters. It then defines a regular PyTorch training loop.
+
+You interact with the Ray Train API in only a few locations, which follow in-line comments in the snippet below.
+
+Remember that you want to do data-parallel training for all the models.
+
+
+#. Load the data shard for each worker with `session.get_dataset_shard("train")``
+#. Iterate over the dataset with `train_dataset.iter_torch_batches()``
+#. Report results to Ray Train with `session.report(results)``
+
+The code is compacted for brevity. The `full code <https://github.com/ray-project/ray/blob/master/doc/source/templates/05_dreambooth_finetuning/dreambooth/train.py>`_ is more thoroughly annotated.
+
+
+.. literalinclude:: /templates/05_dreambooth_finetuning/dreambooth/train.py
+  :language: python
+  :start-at: def train_fn(config)
+  :end-before: END: Training loop
+
+You can then run this training function with Ray Train's TorchTrainer:
+
+
+.. literalinclude:: /templates/05_dreambooth_finetuning/dreambooth/train.py
+  :language: python
+  :start-at: args = train_arguments
+  :end-at: trainer.fit()
+  :dedent: 4
+
+Configure the scale
+^^^^^^^^^^^^^^^^^^^
+
+In the TorchTrainer, you can easily configure the scale.
+The preceding example uses the ``num_workers`` argument to specify the number
+of workers. This argument defaults to 2 workers with 1 GPU each, totalling to 2 GPUs.
+
+To run the example on 4 GPUs, set the number of workers to 4 using ``--num-workers=4``.
+Or you can change the scaling config directly:
+
+.. code-block:: diff
+
+     scaling_config=ScalingConfig(
+         use_gpu=True,
+    -    num_workers=args.num_workers,
+    +    num_workers=4,
+     )
+
+If you're running multi-node training, make sure that all nodes have access to a shared
+storage like NFS or EFS. In the following example script, you can adjust the location with the
+``DATA_PREFIX`` environment variable.
+
+Training throughput
+~~~~~~~~~~~~~~~~~~~
+
+Compare throughput of the preceding training runs that used 1,  2, and 4 workers or GPUs.
+
+Consider the following setup:
+
+* 1 GCE g2-standard-48-nvidia-l4-4 instance with 4 GPUs
+* Model as configured below
+* Data from this example
+* 200 regularization images
+* Training for 4 epochs (local batch size = 2)
+* 3 runs per configuration
+
+You expect that the training time should benefit from scale and decreases when running with
+more workers and GPUs.
+
+.. image:: /templates/05_dreambooth_finetuning/dreambooth/images/dreambooth_training.png
+   :alt: DreamBooth training times
+
+.. list-table::
+   :header-rows: 1
+
+   * - Number of workers/GPUs
+     - Training time (seconds)
+   * - 1
+     - 802.14
+   * - 2
+     - 487.82
+   * - 4
+     - 313.25
+
+
+While the training time decreases linearly with the amount of workers/GPUs, you can observe some penalty.
+Specifically, with double the amount of workers you don't get half of the training time.
+
+This penalty is most likely due to additional communication between processes and the transfer of large model
+weights. You are also only training with a batch size of one because of the GPU memory limitation. On larger
+GPUs with higher batch sizes you would expect a greater benefit from scaling out.
+
+
+Run the example
+---------------
+
+First, download the pre-trained Stable Diffusion model as a starting point.
+
+Then train this model with a few images of a subject.
+
+To achieve this, choose a non-word as an identifier, such as ``unqtkn``. When fine-tuning the model with this subject, you teach the model that the prompt is ``A photo of a unqtkn <class>``.
+
+After fine-tuning you can run inference with this specific prompt.
+For instance: ``A photo of a unqtkn <class>`` creates an image of the subject.
+Similarly, ``A photo of a unqtkn <class> at the beach`` creates an image of the subject at the beach.
+
+Step 0: Preparation
+^^^^^^^^^^^^^^^^^^^
+
+Clone the Ray repository, go to the example directory, and install dependencies.
+
+.. code-block:: bash
+
+   git clone https://github.com/ray-project/ray.git
+   cd doc/source/templates/05_dreambooth_finetuning
+   pip install -Ur dreambooth/requirements.txt
+
+Prepare some directories and environment variables.
+
+.. literalinclude:: /templates/05_dreambooth_finetuning/dreambooth_run.sh
+  :language: bash
+  :start-after: __preparation_start__
+  :end-before: __preparation_end__
+
+
+Step 1: Download the pre-trained model
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+Download and cache a pre-trained Stable Diffusion model locally.
+
+.. literalinclude:: /templates/05_dreambooth_finetuning/dreambooth_run.sh
+  :language: bash
+  :start-after: __cache_model_start__
+  :end-before: __cache_model_end__
+
+You can access the downloaded model checkpoint at the ``$ORIG_MODEL_PATH``.
+
+Step 2: Supply images of your subject
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+Use one of the sample datasets, like `dog` or `lego car`, or provide your own directory
+of images, and specify the directory with the ``$INSTANCE_DIR`` environment variable.
+
+Then, copy these images to ``$IMAGES_OWN_DIR``.
+
+.. literalinclude:: /templates/05_dreambooth_finetuning/dreambooth_run.sh
+  :language: bash
+  :start-after: __supply_own_images_start__
+  :end-before: __supply_own_images_end__
+
+The ``$CLASS_NAME`` should be the general category of your subject.
+The images produced by the prompt ``photo of a unqtkn <class>`` should be diverse images
+that are different enough from the subject in order for generated images to clearly
+show the effect of fine-tuning.
+
+Step 3: Create the regularization images
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+Create a regularization image set for a class of subjects using the pre-trained
+Stable Diffusion model. This regularization set ensures that
+the model still produces decent images for random images of the same class,
+rather than just optimize for producing good images of the subject.
+
+.. literalinclude:: /templates/05_dreambooth_finetuning/dreambooth_run.sh
+  :language: bash
+  :start-after: Step 3: START
+  :end-before: Step 3: END
+
+Use Ray Data to do batch inference with 4 workers, to generate more images in parallel.
+
+Step 4: Fine-tune the model
+^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+Save a few, like 4 to 5, images of the subject being fine-tuned
+in a local directory. Then launch the training job with:
+
+.. literalinclude:: /templates/05_dreambooth_finetuning/dreambooth_run.sh
+  :language: bash
+  :start-after: Step 4: START
+  :end-before: Step 4: END
+
+Step 5: Generate images of the subject
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+Try your model with the same command line as Step 2, but point
+to your own model this time.
+
+.. literalinclude:: /templates/05_dreambooth_finetuning/dreambooth_run.sh
+  :language: bash
+  :start-after: Step 5: START
+  :end-before: Step 5: END
+
+Next, try replacing the prompt with something more interesting.
+
+For example, for the dog subject, you can try:
+
+- "photo of a unqtkn dog in a bucket"
+- "photo of a unqtkn dog sleeping"
+- "photo of a unqtkn dog in a doghouse"
+
+See also
+--------
+
+* :doc:`Ray Train Examples <../../examples>` for more use cases
+
+* :ref:`Ray Train User Guides <train-user-guides>` for how-to guides

doc/source/train/examples/pytorch/pytorch_resnet_finetune.ipynb

@@ -0,0 +1,519 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# Finetuning a Pytorch Image Classifier with Ray Train\n",
+    "\n",
+    "<a id=\"try-anyscale-quickstart-pytorch_resnet_finetune\" href=\"https://console.anyscale.com/register/ha?render_flow=ray&utm_source=ray_docs&utm_medium=docs&utm_campaign=pytorch_resnet_finetune\">\n",
+    "    <img src=\"../../../_static/img/run-on-anyscale.svg\" alt=\"try-anyscale-quickstart\">\n",
+    "</a>\n",
+    "<br></br>\n",
+    "\n",
+    "This example fine-tunes a pre-trained ResNet model with Ray Train. \n",
+    "\n",
+    "For this example, the network architecture consists of the intermediate layer output of a pre-trained ResNet model, which feeds into a randomly initialized linear layer that outputs classification logits for our new task.\n",
+    "\n",
+    "\n"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Load and preprocess finetuning dataset\n",
+    "This example is adapted from Pytorch's [Transfer Learning for Computer Vision](https://pytorch.org/tutorials/beginner/transfer_learning_tutorial.html) tutorial.\n",
+    "We will use *hymenoptera_data* as the finetuning dataset, which contains two classes (bees and ants) and 397 total images (across training and validation). This is a quite small dataset and used only for demonstration purposes. "
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 15,
+   "metadata": {
+    "tags": [
+     "remove-cell"
+    ]
+   },
+   "outputs": [],
+   "source": [
+    "# To run full example, set SMOKE_TEST as False\n",
+    "SMOKE_TEST = True\n"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "The dataset is publicly available [here](https://www.kaggle.com/datasets/ajayrana/hymenoptera-data). Note that it is structured with directory names as the labels. Use `torchvision.datasets.ImageFolder()` to load the images and their corresponding labels."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import os\n",
+    "import torch\n",
+    "import torch.nn as nn\n",
+    "import torch.optim as optim\n",
+    "from torch.utils.data import DataLoader\n",
+    "from torchvision import datasets, models, transforms\n",
+    "import numpy as np\n",
+    "\n",
+    "# Data augmentation and normalization for training\n",
+    "# Just normalization for validation\n",
+    "data_transforms = {\n",
+    "    \"train\": transforms.Compose(\n",
+    "        [\n",
+    "            transforms.RandomResizedCrop(224),\n",
+    "            transforms.RandomHorizontalFlip(),\n",
+    "            transforms.ToTensor(),\n",
+    "            transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),\n",
+    "        ]\n",
+    "    ),\n",
+    "    \"val\": transforms.Compose(\n",
+    "        [\n",
+    "            transforms.Resize(224),\n",
+    "            transforms.CenterCrop(224),\n",
+    "            transforms.ToTensor(),\n",
+    "            transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),\n",
+    "        ]\n",
+    "    ),\n",
+    "}\n",
+    "\n",
+    "def download_datasets():\n",
+    "    os.system(\n",
+    "        \"wget https://download.pytorch.org/tutorial/hymenoptera_data.zip >/dev/null 2>&1\"\n",
+    "    )\n",
+    "    os.system(\"unzip hymenoptera_data.zip >/dev/null 2>&1\")\n",
+    "\n",
+    "# Download and build torch datasets\n",
+    "def build_datasets():\n",
+    "    torch_datasets = {}\n",
+    "    for split in [\"train\", \"val\"]:\n",
+    "        torch_datasets[split] = datasets.ImageFolder(\n",
+    "            os.path.join(\"./hymenoptera_data\", split), data_transforms[split]\n",
+    "        )\n",
+    "    return torch_datasets\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 16,
+   "metadata": {
+    "tags": [
+     "remove-cell"
+    ]
+   },
+   "outputs": [],
+   "source": [
+    "if SMOKE_TEST:\n",
+    "    from torch.utils.data import Subset\n",
+    "\n",
+    "    def build_datasets():\n",
+    "        torch_datasets = {}\n",
+    "        for split in [\"train\", \"val\"]:\n",
+    "            torch_datasets[split] = datasets.ImageFolder(\n",
+    "                os.path.join(\"./hymenoptera_data\", split), data_transforms[split]\n",
+    "            )\n",
+    "            \n",
+    "        # Only take a subset for smoke test\n",
+    "        for split in [\"train\", \"val\"]:\n",
+    "            indices = list(range(100))\n",
+    "            torch_datasets[split] = Subset(torch_datasets[split], indices)\n",
+    "        return torch_datasets\n"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Initialize Model and Fine-tuning configs"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Next, let's define the training configuration that will be passed into the training loop function later."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 4,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "train_loop_config = {\n",
+    "    \"input_size\": 224,  # Input image size (224 x 224)\n",
+    "    \"batch_size\": 32,  # Batch size for training\n",
+    "    \"num_epochs\": 10,  # Number of epochs to train for\n",
+    "    \"lr\": 0.001,  # Learning Rate\n",
+    "    \"momentum\": 0.9,  # SGD optimizer momentum\n",
+    "}\n"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Next, let's define our model. You can either create a model from pre-trained weights or reload the model checkpoint from a previous run."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 5,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import os\n",
+    "import torch\n",
+    "from ray.train import Checkpoint\n",
+    "\n",
+    "# Option 1: Initialize model with pretrained weights\n",
+    "def initialize_model():\n",
+    "    # Load pretrained model params\n",
+    "    model = models.resnet50(pretrained=True)\n",
+    "\n",
+    "    # Replace the original classifier with a new Linear layer\n",
+    "    num_features = model.fc.in_features\n",
+    "    model.fc = nn.Linear(num_features, 2)\n",
+    "\n",
+    "    # Ensure all params get updated during finetuning\n",
+    "    for param in model.parameters():\n",
+    "        param.requires_grad = True\n",
+    "    return model\n",
+    "\n",
+    "\n",
+    "# Option 2: Initialize model with an Train checkpoint\n",
+    "# Replace this with your own uri\n",
+    "CHECKPOINT_FROM_S3 = Checkpoint(\n",
+    "    path=\"s3://air-example-data/finetune-resnet-checkpoint/TorchTrainer_4f69f_00000_0_2023-02-14_14-04-09/checkpoint_000001/\"\n",
+    ")\n",
+    "\n",
+    "\n",
+    "def initialize_model_from_checkpoint(checkpoint: Checkpoint):\n",
+    "    with checkpoint.as_directory() as tmpdir:\n",
+    "        state_dict = torch.load(os.path.join(tmpdir, \"checkpoint.pt\"))\n",
+    "    resnet50 = initialize_model()\n",
+    "    resnet50.load_state_dict(state_dict[\"model\"])\n",
+    "    return resnet50\n"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Define the Training Loop\n",
+    "\n",
+    "The `train_loop_per_worker` function defines the fine-tuning procedure for each worker.\n",
+    "\n",
+    "**1. Prepare dataloaders for each worker**:\n",
+    "- This tutorial assumes you are using PyTorch's native `torch.utils.data.Dataset` for data input. {meth}`train.torch.prepare_data_loader() <ray.train.torch.prepare_data_loader>`  prepares your dataLoader for distributed execution. You can also use Ray Data for more efficient preprocessing. For more details on using Ray Data for images, see the {doc}`Working with Images </data/working-with-images>` Ray Data user guide.\n",
+    "\n",
+    "**2. Prepare your model**:\n",
+    "- {meth}`train.torch.prepare_model() <ray.train.torch.prepare_model>` prepares the model for distributed training. Under the hood, it converts your torch model to `DistributedDataParallel` model, which synchronize its weights across all workers.\n",
+    "\n",
+    "**3. Report metrics and checkpoint**:\n",
+    "- {meth}`train.report() <ray.train.report>` will report metrics and checkpoints to Ray Train.\n",
+    "- Saving checkpoints through {meth}`train.report(metrics, checkpoint=...) <ray.train.report>` will automatically [upload checkpoints to cloud storage](tune-cloud-checkpointing) (if configured), and allow you to easily enable Ray Train worker fault tolerance in the future."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 6,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import os\n",
+    "from tempfile import TemporaryDirectory\n",
+    "\n",
+    "import ray.train as train\n",
+    "from ray.train import Checkpoint\n",
+    "\n",
+    "\n",
+    "\n",
+    "def evaluate(logits, labels):\n",
+    "    _, preds = torch.max(logits, 1)\n",
+    "    corrects = torch.sum(preds == labels).item()\n",
+    "    return corrects\n",
+    "\n",
+    "\n",
+    "def train_loop_per_worker(configs):\n",
+    "    import warnings\n",
+    "\n",
+    "    warnings.filterwarnings(\"ignore\")\n",
+    "\n",
+    "    # Calculate the batch size for a single worker\n",
+    "    worker_batch_size = configs[\"batch_size\"] // train.get_context().get_world_size()\n",
+    "\n",
+    "    # Download dataset once on local rank 0 worker\n",
+    "    if train.get_context().get_local_rank() == 0:\n",
+    "        download_datasets()\n",
+    "    torch.distributed.barrier()\n",
+    "\n",
+    "    # Build datasets on each worker\n",
+    "    torch_datasets = build_datasets()\n",
+    "\n",
+    "    # Prepare dataloader for each worker\n",
+    "    dataloaders = dict()\n",
+    "    dataloaders[\"train\"] = DataLoader(\n",
+    "        torch_datasets[\"train\"], batch_size=worker_batch_size, shuffle=True\n",
+    "    )\n",
+    "    dataloaders[\"val\"] = DataLoader(\n",
+    "        torch_datasets[\"val\"], batch_size=worker_batch_size, shuffle=False\n",
+    "    )\n",
+    "\n",
+    "    # Distribute\n",
+    "    dataloaders[\"train\"] = train.torch.prepare_data_loader(dataloaders[\"train\"])\n",
+    "    dataloaders[\"val\"] = train.torch.prepare_data_loader(dataloaders[\"val\"])\n",
+    "\n",
+    "    device = train.torch.get_device()\n",
+    "\n",
+    "    # Prepare DDP Model, optimizer, and loss function\n",
+    "    model = initialize_model()\n",
+    "    model = train.torch.prepare_model(model)\n",
+    "\n",
+    "    optimizer = optim.SGD(\n",
+    "        model.parameters(), lr=configs[\"lr\"], momentum=configs[\"momentum\"]\n",
+    "    )\n",
+    "    criterion = nn.CrossEntropyLoss()\n",
+    "\n",
+    "    # Start training loops\n",
+    "    for epoch in range(configs[\"num_epochs\"]):\n",
+    "        # Each epoch has a training and validation phase\n",
+    "        for phase in [\"train\", \"val\"]:\n",
+    "            if phase == \"train\":\n",
+    "                model.train()  # Set model to training mode\n",
+    "            else:\n",
+    "                model.eval()  # Set model to evaluate mode\n",
+    "\n",
+    "            running_loss = 0.0\n",
+    "            running_corrects = 0\n",
+    "\n",
+    "            if train.get_context().get_world_size() > 1:\n",
+    "                dataloaders[phase].sampler.set_epoch(epoch)\n",
+    "\n",
+    "            for inputs, labels in dataloaders[phase]:\n",
+    "                inputs = inputs.to(device)\n",
+    "                labels = labels.to(device)\n",
+    "\n",
+    "                # zero the parameter gradients\n",
+    "                optimizer.zero_grad()\n",
+    "\n",
+    "                # forward\n",
+    "                with torch.set_grad_enabled(phase == \"train\"):\n",
+    "                    # Get model outputs and calculate loss\n",
+    "                    outputs = model(inputs)\n",
+    "                    loss = criterion(outputs, labels)\n",
+    "\n",
+    "                    # backward + optimize only if in training phase\n",
+    "                    if phase == \"train\":\n",
+    "                        loss.backward()\n",
+    "                        optimizer.step()\n",
+    "\n",
+    "                # calculate statistics\n",
+    "                running_loss += loss.item() * inputs.size(0)\n",
+    "                running_corrects += evaluate(outputs, labels)\n",
+    "\n",
+    "            size = len(torch_datasets[phase]) // train.get_context().get_world_size()\n",
+    "            epoch_loss = running_loss / size\n",
+    "            epoch_acc = running_corrects / size\n",
+    "\n",
+    "            if train.get_context().get_world_rank() == 0:\n",
+    "                print(\n",
+    "                    \"Epoch {}-{} Loss: {:.4f} Acc: {:.4f}\".format(\n",
+    "                        epoch, phase, epoch_loss, epoch_acc\n",
+    "                    )\n",
+    "                )\n",
+    "\n",
+    "            # Report metrics and checkpoint every epoch\n",
+    "            if phase == \"val\":\n",
+    "                with TemporaryDirectory() as tmpdir:\n",
+    "                    state_dict = {\n",
+    "                        \"epoch\": epoch,\n",
+    "                        \"model\": model.module.state_dict(),\n",
+    "                        \"optimizer_state_dict\": optimizer.state_dict(),\n",
+    "                    }\n",
+    "                    torch.save(state_dict, os.path.join(tmpdir, \"checkpoint.pt\"))\n",
+    "                    train.report(\n",
+    "                        metrics={\"loss\": epoch_loss, \"acc\": epoch_acc},\n",
+    "                        checkpoint=Checkpoint.from_directory(tmpdir),\n",
+    "                    )\n"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Next, setup the TorchTrainer:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 7,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from ray.train.torch import TorchTrainer\n",
+    "from ray.train import ScalingConfig, RunConfig, CheckpointConfig\n",
+    "\n",
+    "# Scale out model training across 4 GPUs.\n",
+    "scaling_config = ScalingConfig(\n",
+    "    num_workers=4, use_gpu=True, resources_per_worker={\"CPU\": 1, \"GPU\": 1}\n",
+    ")\n",
+    "\n",
+    "# Save the latest checkpoint\n",
+    "checkpoint_config = CheckpointConfig(num_to_keep=1)\n",
+    "\n",
+    "# Set experiment name and checkpoint configs\n",
+    "run_config = RunConfig(\n",
+    "    name=\"finetune-resnet\",\n",
+    "    storage_path=\"/tmp/ray_results\",\n",
+    "    checkpoint_config=checkpoint_config,\n",
+    ")\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 8,
+   "metadata": {
+    "tags": [
+     "remove-cell"
+    ]
+   },
+   "outputs": [],
+   "source": [
+    "if SMOKE_TEST:\n",
+    "    scaling_config = ScalingConfig(\n",
+    "        num_workers=8, use_gpu=False, resources_per_worker={\"CPU\": 1}\n",
+    "    )\n",
+    "    train_loop_config[\"num_epochs\"] = 1\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "trainer = TorchTrainer(\n",
+    "    train_loop_per_worker=train_loop_per_worker,\n",
+    "    train_loop_config=train_loop_config,\n",
+    "    scaling_config=scaling_config,\n",
+    "    run_config=run_config,\n",
+    ")\n",
+    "\n",
+    "result = trainer.fit()\n",
+    "print(result)\n"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Load the checkpoint for prediction:\n",
+    "\n",
+    " \n",
+    " The metadata and checkpoints have already been saved in the `storage_path` specified in TorchTrainer:"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "We now need to load the trained model and evaluate it on test data. The best model parameters have been saved in `log_dir`. We can load the resulting checkpoint from our fine-tuning run using the previously defined `initialize_model_from_checkpoint()` function."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "model = initialize_model_from_checkpoint(result.checkpoint)\n",
+    "device = torch.device(\"cuda\")\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 12,
+   "metadata": {
+    "tags": [
+     "remove-cell"
+    ]
+   },
+   "outputs": [],
+   "source": [
+    "if SMOKE_TEST:\n",
+    "    device = torch.device(\"cpu\")\n"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Finally, define a simple evaluation loop and check the performance of the checkpoint model."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 14,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Accuracy:  0.934640522875817\n"
+     ]
+    }
+   ],
+   "source": [
+    "model = model.to(device)\n",
+    "model.eval()\n",
+    "\n",
+    "download_datasets()\n",
+    "torch_datasets = build_datasets()\n",
+    "dataloader = DataLoader(torch_datasets[\"val\"], batch_size=32, num_workers=4)\n",
+    "corrects = 0\n",
+    "for inputs, labels in dataloader:\n",
+    "    inputs = inputs.to(device)\n",
+    "    labels = labels.to(device)\n",
+    "    preds = model(inputs)\n",
+    "    corrects += evaluate(preds, labels)\n",
+    "\n",
+    "print(\"Accuracy: \", corrects / len(dataloader.dataset))\n"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3 (ipykernel)",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.8.13"
+  },
+  "orphan": true,
+  "vscode": {
+   "interpreter": {
+    "hash": "a8c1140d108077f4faeb76b2438f85e4ed675f93d004359552883616a1acd54c"
+   }
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 4
+}

doc/source/train/examples/pytorch/torch_fashion_mnist_example.rst

@@ -0,0 +1,27 @@
+:orphan:
+
+.. _train-pytorch-fashion-mnist:
+
+Train a PyTorch model on Fashion MNIST
+======================================
+
+.. raw:: html
+
+    <a id="try-anyscale-quickstart-torch_fashion_mnist_example" target="_blank" href="https://console.anyscale.com/register/ha?render_flow=ray&utm_source=ray_docs&utm_medium=docs&utm_campaign=torch_fashion_mnist_example">
+      <img src="../../../_static/img/run-on-anyscale.svg" alt="Run on Anyscale" />
+      <br/><br/>
+    </a>
+
+This example runs distributed training of a PyTorch model on Fashion MNIST with Ray Train.
+
+Code example
+------------
+
+.. literalinclude:: /../../python/ray/train/examples/pytorch/torch_fashion_mnist_example.py
+
+See also
+--------
+
+* :ref:`Get Started with PyTorch <train-pytorch>` for a tutorial on using Ray Train and PyTorch
+
+* :doc:`Ray Train Examples <../../examples>` for more use cases

doc/source/train/examples/pytorch/torch_regression_example.rst

@@ -0,0 +1,14 @@
+
+:orphan:
+
+torch_regression_example
+========================
+
+.. raw:: html
+
+    <a id="try-anyscale-quickstart-torch_regression_example" target="_blank" href="https://console.anyscale.com/register/ha?render_flow=ray&utm_source=ray_docs&utm_medium=docs&utm_campaign=torch_regression_example">
+      <img src="../../../_static/img/run-on-anyscale.svg" alt="Run on Anyscale" />
+      <br/><br/>
+    </a>
+
+.. literalinclude:: /../../python/ray/train/examples/pytorch/torch_regression_example.py