notebooks/official/training/pytorch_gcs_data_training.ipynb

{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "q7r5vp1ZZROn"
      },
      "outputs": [],
      "source": [
        "# Copyright 2023 Google LLC\n",
        "#\n",
        "# Licensed under the Apache License, Version 2.0 (the \"License\");\n",
        "# you may not use this file except in compliance with the License.\n",
        "# You may obtain a copy of the License at\n",
        "#\n",
        "#     https://www.apache.org/licenses/LICENSE-2.0\n",
        "#\n",
        "# Unless required by applicable law or agreed to in writing, software\n",
        "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
        "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
        "# See the License for the specific language governing permissions and\n",
        "# limitations under the License."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "JAPoU8Sm5E6e"
      },
      "source": [
        "# Train PyTorch model on Vertex AI with data from Cloud Storage\n",
        "\n",
        "<table align=\"left\">\n",
        "\n",
        "  <td>\n",
        "    <a href=\"https://colab.research.google.com/github/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/official/training/pytorch_gcs_data_training.ipynb\">\n",
        "      <img src=\"https://cloud.google.com/ml-engine/images/colab-logo-32px.png\" alt=\"Colab logo\"> Run in Colab\n",
        "    </a>\n",
        "  </td>\n",
        "  <td>\n",
        "    <a href=\"https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/official/training/pytorch_gcs_data_training.ipynb\">\n",
        "      <img src=\"https://cloud.google.com/ml-engine/images/github-logo-32px.png\" alt=\"GitHub logo\">\n",
        "      View on GitHub\n",
        "    </a>\n",
        "  </td>\n",
        "  <td>\n",
        "    <a href=\"https://console.cloud.google.com/vertex-ai/workbench/deploy-notebook?download_url=https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/main/notebooks/official/training/pytorch_gcs_data_training.ipynb\">\n",
        "      <img src=\"https://lh3.googleusercontent.com/UiNooY4LUgW_oTvpsNhPpQzsstV5W8F7rYgxgGBD85cWJoLmrOzhVs_ksK_vgx40SHs7jCqkTkCk=e14-rj-sc0xffffff-h130-w32\" alt=\"Vertex AI logo\">\n",
        "      Open in Vertex AI Workbench\n",
        "    </a>\n",
        "  </td>                                                                                               \n",
        "</table>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "tvgnzT1CKxrO"
      },
      "source": [
        "## Overview\n",
        "\n",
        "This tutorial shows you how to create a custom training job using PyTorch and a dataset stored on Cloud Storage.\n",
        "\n",
        "Learn more about [PyTorch integration in Vertex AI](https://cloud.google.com/vertex-ai/docs/start/pytorch)."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "d975e698c9a4"
      },
      "source": [
        "### Objective\n",
        "\n",
        "In this tutorial, you learn how to create a training job using PyTorch and a dataset stored on Cloud Storage. You build a custom training script that uses GCSFuse to load data from a bucket. The custom training script creates a simple neural network and saves the model artifact to a bucket on Cloud Storage.\n",
        "\n",
        "This tutorial uses the following Google Cloud ML services and resources:\n",
        "\n",
        "- Vertex AI Training\n",
        "- Cloud Storage\n",
        "\n",
        "The steps performed include:\n",
        "\n",
        "- Write a custom training script that creates your train & test datasets and trains the model.\n",
        "- Run a Vertex AI SDK `CustomTrainingJob`"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "08d289fa873f"
      },
      "source": [
        "### Dataset\n",
        "\n",
        "This tutorial uses the [MNIST handwriting sample](https://en.wikipedia.org/wiki/MNIST_database) that classifies hand-written digits. For this tutorial, a CSV version of the dataset available on [Kaggle](https://www.kaggle.com/datasets/oddrationale/mnist-in-csv?select=mnist_train.csv) has been uploaded to a Cloud Storage bucket for your use."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "aed92deeb4a0"
      },
      "source": [
        "### Costs \n",
        "\n",
        "This tutorial uses billable components of Google Cloud:\n",
        "\n",
        "* Vertex AI\n",
        "* Cloud Storage\n",
        "\n",
        "Learn about [Vertex AI pricing](https://cloud.google.com/vertex-ai/pricing),\n",
        "and [Cloud Storage pricing](https://cloud.google.com/storage/pricing), \n",
        "and use the [Pricing Calculator](https://cloud.google.com/products/calculator/)\n",
        "to generate a cost estimate based on your projected usage."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "1566ddfb0d0d"
      },
      "source": [
        "## Requirements\n",
        "\n",
        "This tutorial requires using a notebook that is optimized for use with PyTorch. If you run this notebook in Vertex AI Workbench, ensure that you notebook image meets the following requirements:\n",
        "\n",
        "+ PyTorch 1.13 notebook\n",
        "+ 1 NVIDIA T4 GPU\n",
        "\n",
        "Colab notebooks meet the requirements (after installation and authentication). You might need to switch to a GPU-enabled runtime."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "i7EUnXsZhAGF"
      },
      "source": [
        "## Installation\n",
        "\n",
        "Install the following packages required to execute this notebook. "
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "b4LrAIokarTK"
      },
      "outputs": [],
      "source": [
        "%%writefile requirements.txt\n",
        "\n",
        "google-cloud-aiplatform\n",
        "google-cloud-storage\n",
        "torch==1.12"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "WfE-8h14a1za"
      },
      "outputs": [],
      "source": [
        "%pip install -q -r requirements.txt"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "67e101d52a4c"
      },
      "source": [
        "### Colab only: Uncomment the following cell to restart the kernel"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "bfcf20418562"
      },
      "outputs": [],
      "source": [
        "# Automatically restart kernel after installs so that your environment can access the new packages\n",
        "# import IPython\n",
        "\n",
        "# app = IPython.Application.instance()\n",
        "# app.kernel.do_shutdown(True)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "BF1j6f9HApxa"
      },
      "source": [
        "## Before you begin\n",
        "\n",
        "### Set up your Google Cloud project\n",
        "\n",
        "**The following steps are required, regardless of your notebook environment.**\n",
        "\n",
        "1. [Select or create a Google Cloud project](https://console.cloud.google.com/cloud-resource-manager). When you first create an account, you get a $300 free credit towards your compute/storage costs.\n",
        "\n",
        "2. [Make sure that billing is enabled for your project](https://cloud.google.com/billing/docs/how-to/modify-project).\n",
        "\n",
        "3. [Enable the Vertex AI API](https://console.cloud.google.com/flows/enableapi?apiid=aiplatform.googleapis.com). \n",
        "\n",
        "4. [Enable the Storage API](https://console.cloud.google.com/flows/enableapi?apiid=storage.googleapis.com).\n",
        "\n",
        "5. If you are running this notebook locally, you need to install the [Cloud SDK](https://cloud.google.com/sdk)."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "WReHDGG5g0XY"
      },
      "source": [
        "#### Set your project ID\n",
        "\n",
        "**If you don't know your project ID**, try the following:\n",
        "* Run `gcloud config list`.\n",
        "* Run `gcloud projects list`.\n",
        "* See the support page: [Locate the project ID](https://support.google.com/googleapi/answer/7014113)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "3f4a7fc51051"
      },
      "outputs": [],
      "source": [
        "PROJECT_ID = \"[your-project-id]\"  # @param\n",
        "\n",
        "# Set the project id\n",
        "! gcloud config set project {PROJECT_ID}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "region"
      },
      "source": [
        "#### Region\n",
        "\n",
        "You can also change the `REGION` variable used by Vertex AI. Learn more about [Vertex AI regions](https://cloud.google.com/vertex-ai/docs/general/locations)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "0gtuztzPg0x8"
      },
      "outputs": [],
      "source": [
        "REGION = \"us-central1\"  # @param"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "sBCra4QMA2wR"
      },
      "source": [
        "### Authenticate your Google Cloud account\n",
        "\n",
        "Depending on your Jupyter environment, you may have to manually authenticate. Follow the relevant instructions below."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "74ccc9e52986"
      },
      "source": [
        "**1. Vertex AI Workbench**\n",
        "* Do nothing as you are already authenticated."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "de775a3773ba"
      },
      "source": [
        "**2. Local JupyterLab instance, uncomment and run:**"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "254614fa0c46"
      },
      "outputs": [],
      "source": [
        "# ! gcloud auth login"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ef21552ccea8"
      },
      "source": [
        "**3. Colab, uncomment and run:**"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "603adbbf0532"
      },
      "outputs": [],
      "source": [
        "# from google.colab import auth\n",
        "# auth.authenticate_user()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "f6b2ccc891ed"
      },
      "source": [
        "**4. Service account or other**\n",
        "* See how to grant Cloud Storage permissions to your service account at https://cloud.google.com/storage/docs/gsutil/commands/iam#ch-examples."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "zgPO1eR3CYjk"
      },
      "source": [
        "### Create a Cloud Storage bucket\n",
        "\n",
        "Create a storage bucket to store intermediate artifacts such as datasets."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "MzGDU7TWdts_"
      },
      "outputs": [],
      "source": [
        "BUCKET_URI = f\"gs://your-bucket-name-{PROJECT_ID}-unique\"  # @param {type:\"string\"}\n",
        "BUCKET_PREFIX = \"pytorch-on-gcs\"  # @param"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "-EcIXiGsCePi"
      },
      "source": [
        "**Only if your bucket doesn't already exist**: Run the following cell to create your Cloud Storage bucket."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "NIq7R4HZCfIc"
      },
      "outputs": [],
      "source": [
        "! gsutil mb -l {REGION} -p {PROJECT_ID} {BUCKET_URI}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "960505627ddf"
      },
      "source": [
        "### Import libraries"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Py_tlZxxacOq"
      },
      "outputs": [],
      "source": [
        "import os\n",
        "from datetime import datetime\n",
        "\n",
        "import numpy as np\n",
        "import pandas as pd\n",
        "import torch\n",
        "import torch.nn as nn\n",
        "from google.cloud import aiplatform\n",
        "from matplotlib import pyplot as plt\n",
        "from torch.utils.data import DataLoader, Dataset"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "init_aip:mbsdk,all"
      },
      "source": [
        "### Initialize Vertex AI SDK for Python\n",
        "\n",
        "Initialize the Vertex AI SDK for Python for your project."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "init_aip:mbsdk,all"
      },
      "outputs": [],
      "source": [
        "aiplatform.init(project=PROJECT_ID, location=REGION, staging_bucket=BUCKET_URI)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "e54d6bccf451"
      },
      "source": [
        "### Provide the URIs for the data\n",
        "\n",
        "As mentioned before, this tutorial uses the classic MNIST handwritten digits dataset as an input. The dataset has been stored for your use on a publicly available Cloud Storage location. You can use these CSV files directly in your training script.\n",
        "\n",
        "**Note**: You can use the [PyTorch datasets library](https://pytorch.org/vision/stable/generated/torchvision.datasets.MNIST.html#torchvision.datasets.MNIST) to download a this dataset. For learning purposes, this tutorial uses a copy of the dataset on Cloud Storage."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "9b5044d31ed9"
      },
      "outputs": [],
      "source": [
        "TRAIN_URI = \"gs://cloud-samples-data/vertex-ai/training/pytorch/mnist_train.csv\"\n",
        "TEST_URI = \"gs://cloud-samples-data/vertex-ai/training/pytorch/mnist_test.csv\"\n",
        "\n",
        "print(TRAIN_URI)\n",
        "print(TEST_URI)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "_7zjKtBsLOrL"
      },
      "source": [
        "## [OPTIONAL] Inspect the dataset from GCS\n",
        "\n",
        "Before creating the training script, take a quick look at the data contained in the CSV files on Cloud Storage. You can use the PyTorch [`Dataset`](https://pytorch.org/docs/stable/data.html?highlight=torch+utils+data+dataset#torch.utils.data.Dataset) and [`DataLoader`](https://pytorch.org/docs/stable/data.html?highlight=torch+utils+data+dataset#torch.utils.data.DataLoader) classes to instantiate a dataset and then plot the data using [matplotlib](https://matplotlib.org/stable/index.html).\n",
        "\n",
        "Start by downloading the CSV files into your local development environment."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "K1MQpp0QLRWK"
      },
      "outputs": [],
      "source": [
        "! gsutil -m cp -r $TRAIN_URI .\n",
        "! gsutil -m cp -r $TEST_URI ."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "365a4f9c588b"
      },
      "source": [
        "Next you need to define a custom image dataset that inherits from the base `Dataset` class. \n",
        "\n",
        "Note that your custom `Dataset` class must override the `__init__`, `__len__`, and `__getitem__` methods. These methods are used by the `DataLoader` class to iterate through your dataset.\n",
        "\n",
        "The following `CustomImageDataset` class has a several notable features. First, the dimensions of images are hardcoded as 28 pixels high by 28 pixels wide. This corresponds to the dimensions of the images in the MNIST dataset.\n",
        "\n",
        "Second, this dataset uses [`pandas.Dataframe`](https://pandas.pydata.org/docs/reference/api/pandas.DataFrame.html) objects for reading the CSV files and accessing the data inside of them.\n",
        "\n",
        "Finally, the `__getitem__` method needs to extract the image and the label out of each row in the CSV file. These two values are provided to the caller as a tuple.\n",
        "\n",
        "The image itself needs to be converted from a one-dimensional vector value (a list) into a 2-dimensional matrix (a list of lists). In addition, the grayscale values in the CSV files, stored as integers, need to be converted into a float value between 0.0 and 1.0. To do this conversion, you multiple the grayscale value by the decimal equivalent of 1/255."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "A9Z2j58NlIu7"
      },
      "outputs": [],
      "source": [
        "class CustomImageDataset(Dataset):\n",
        "    width = 28  # hard-coded width & height of image matrix\n",
        "    height = 28\n",
        "\n",
        "    def __init__(self, data_file, transform=None, target_transform=None):\n",
        "        self.dataset = pd.read_csv(data_file)\n",
        "        self.transform = (\n",
        "            transform  # We would use ToTensor() if we were taking in raw images\n",
        "        )\n",
        "        self.target_transform = target_transform\n",
        "\n",
        "    def __len__(self):\n",
        "        return self.dataset.shape[0]\n",
        "\n",
        "    def __getitem__(self, idx):\n",
        "        label = self.dataset.at[idx, \"label\"]\n",
        "        image = self.dataset.iloc[idx, 1:]\n",
        "\n",
        "        # Create a matrix from the pandas.Series\n",
        "        image = image.to_numpy() * 0.00392156862745098  # 1 / 255\n",
        "        image = image.reshape(self.width, self.height)\n",
        "        image = image.astype(float)\n",
        "        image = torch.Tensor(image)\n",
        "\n",
        "        if self.target_transform:\n",
        "            label = self.target_transform(label)\n",
        "        if self.transform:\n",
        "            image = self.transform(image)\n",
        "        return image, label"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Wv4TdH9_otsu"
      },
      "outputs": [],
      "source": [
        "train_set = CustomImageDataset(\"mnist_train.csv\")\n",
        "test_set = CustomImageDataset(\"mnist_test.csv\")\n",
        "\n",
        "batch_size = 64\n",
        "shuffle = False\n",
        "\n",
        "train_dataloader = DataLoader(train_set, batch_size=batch_size, shuffle=shuffle)\n",
        "test_dataloader = DataLoader(test_set, batch_size=batch_size, shuffle=shuffle)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "7bc84344443d"
      },
      "source": [
        "With the dataset loaded into `DataLoader` objects, you can inspect them. For this tutorial, the datasets are provided to the training application in batches of 64 dataset rows. Each dataset row contains a 28x28 image and a label (a value between 0 and 9)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "qFzN3mS-Tlhl"
      },
      "outputs": [],
      "source": [
        "for batch, (X, y) in enumerate(train_dataloader):\n",
        "    print(len(X))\n",
        "    print(len(y))\n",
        "\n",
        "    first_image = X[0]\n",
        "    first_label = y[0]\n",
        "\n",
        "    print(len(first_image))\n",
        "    print(first_label)  # This will be a Tensor object with a single scalar value, 5\n",
        "    break"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "1479af04fb9e"
      },
      "source": [
        "### Plot an image from the dataset.\n",
        "\n",
        "For verify the data quality, plot the first item in the dataset to validate that it renders an image that matches the label."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "0HaLdouYmTBq"
      },
      "outputs": [],
      "source": [
        "first_image, label = (None, None)\n",
        "for i in range(len(train_set)):\n",
        "    sample = train_set[i]\n",
        "    sample, label = sample\n",
        "    first_image = sample.numpy()\n",
        "    break"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "AUJowAEVrI9w"
      },
      "outputs": [],
      "source": [
        "np.shape(first_image)\n",
        "imgplot = plt.imshow(first_image, cmap=\"gray\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "WBdwLPRQK4lS"
      },
      "source": [
        "## [OPTIONAL] Train the neural network locally\n",
        "\n",
        "Although not required for training on Vertex AI, you can train the model locally using PyTorch. This tutorial declares a `NeuralNetwork` class that inherits from PyTorch's [`torch.nn.Module`](https://pytorch.org/docs/stable/generated/torch.nn.Module.html?highlight=nn+module#torch.nn.Module) class. The `nn.Module` class provides a base class for all neural network modules."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "bxHfdDaVK2_p"
      },
      "outputs": [],
      "source": [
        "# Get cpu or gpu device for training\n",
        "device = (\n",
        "    \"cuda\"\n",
        "    if torch.cuda.is_available()\n",
        "    else \"mps\"\n",
        "    if torch.backends.mps.is_available()\n",
        "    else \"cpu\"\n",
        ")\n",
        "print(f\"Using {device} device\")\n",
        "\n",
        "\n",
        "# Define model\n",
        "class NeuralNetwork(nn.Module):\n",
        "    def __init__(self):\n",
        "        super().__init__()\n",
        "        self.flatten = nn.Flatten()\n",
        "        self.linear_relu_stack = nn.Sequential(\n",
        "            nn.Linear(28 * 28, 512, dtype=torch.float),\n",
        "            nn.ReLU(),\n",
        "            nn.Linear(512, 512),\n",
        "            nn.ReLU(),\n",
        "            nn.Linear(512, 10),\n",
        "        )\n",
        "\n",
        "    def forward(self, x):\n",
        "        x = self.flatten(x)\n",
        "        logits = self.linear_relu_stack(x)\n",
        "        return logits\n",
        "\n",
        "\n",
        "model = NeuralNetwork().to(device)\n",
        "print(model)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "cxWBegh4NRYl"
      },
      "outputs": [],
      "source": [
        "def train(dataloader, model, loss_fn, optimizer):\n",
        "    size = len(dataloader.dataset)\n",
        "    model.train()\n",
        "    for batch, (X, y) in enumerate(dataloader):\n",
        "        X, y = X.to(device), y.to(device)\n",
        "\n",
        "        pred = model(X)\n",
        "        loss = loss_fn(pred, y)\n",
        "\n",
        "        optimizer.zero_grad()\n",
        "        loss.backward()\n",
        "        optimizer.step()\n",
        "\n",
        "        if batch % 100 == 0:\n",
        "            loss, current = loss.item(), batch * len(X)\n",
        "            print(f\"loss: {loss:>7f}  [{current:>5d}/{size:>5d}]\")\n",
        "\n",
        "\n",
        "def test(dataloader, model, loss_fn) -> bool:\n",
        "    size = len(dataloader.dataset)\n",
        "    num_batches = len(dataloader)\n",
        "    model.eval()\n",
        "    test_loss, correct = 0, 0\n",
        "    with torch.no_grad():\n",
        "        for X, y in dataloader:\n",
        "            X, y = X.to(device), y.to(device)\n",
        "            pred = model(X)\n",
        "            test_loss += loss_fn(pred, y).item()\n",
        "            correct += (pred.argmax(1) == y).type(torch.float).sum().item()\n",
        "    test_loss /= num_batches\n",
        "    correct /= size\n",
        "    accuracy = 100 * correct\n",
        "    print(f\"Test Error: \\n Accuracy: {(accuracy):>0.1f}%, Avg loss: {test_loss:>8f} \\n\")\n",
        "    if accuracy <= 0.0:\n",
        "        return False\n",
        "    return True\n",
        "\n",
        "\n",
        "# Define a loss function and an optimizer.\n",
        "loss_fn = nn.CrossEntropyLoss()\n",
        "optimizer = torch.optim.SGD(model.parameters(), lr=1e-3)\n",
        "\n",
        "epochs = 5\n",
        "for t in range(epochs):\n",
        "    print(f\"Epoch {t+1}\\n-------------------------------\")\n",
        "    train(train_dataloader, model, loss_fn, optimizer)\n",
        "    is_going_well = test(test_dataloader, model, loss_fn)\n",
        "    if not is_going_well:\n",
        "        print(\"unacceptable accuracy\")\n",
        "        break\n",
        "print(\"Done!\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "F_xoA1zDzany"
      },
      "source": [
        "## Create training script\n",
        "\n",
        "Setting aside all else, the main task for training a custom PyTorch model on Vertex AI is creating a training script. This script is loaded into a [pre-built container for PyTorch training](https://cloud.google.com/vertex-ai/docs/training/pre-built-containers#pytorch) that is then run as a [custom training job on Vertex AI training service](https://cloud.google.com/vertex-ai/docs/training/create-custom-job). \n",
        "\n",
        "The first step is to select a compatible set of accelerators and training images for your custom training job."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "AeXG25Lv0hIv"
      },
      "outputs": [],
      "source": [
        "TRAIN_GPU, TEST_GPU = (aiplatform.gapic.AcceleratorType.NVIDIA_TESLA_K80, 1)\n",
        "if TRAIN_GPU:\n",
        "    TRAIN_VERSION = \"pytorch-gpu.1-9\"\n",
        "else:\n",
        "    TRAIN_VERSION = \"pytorch-xla.1-9\"\n",
        "\n",
        "TRAIN_IMAGE = \"{}-docker.pkg.dev/vertex-ai/training/{}:latest\".format(\n",
        "    REGION.split(\"-\")[0], TRAIN_VERSION\n",
        ")\n",
        "\n",
        "MACHINE_TYPE = \"n1-standard\"\n",
        "\n",
        "VCPU = \"4\"\n",
        "TRAIN_COMPUTE = MACHINE_TYPE + \"-\" + VCPU\n",
        "print(\"Train machine type\", TRAIN_COMPUTE)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "d0EqfgBk2JbD"
      },
      "source": [
        "### Writing out the training script\n",
        "\n",
        "Next, you write the training script, 'task.py', to file before creating the training job. Note that the script includes the dataset, dataloader, and neural net module that you inspected previously.\n",
        "\n",
        "In the training script, the training scripts are loaded from Cloud Storage using [Storage FUSE](https://cloud.google.com/storage/docs/gcs-fuse). FUSE mounts Cloud Storage buckets as folders in the training container's file system. This allows the training script to load files storaged in the bucket as a dataset. FUSE also allows the training script to store the output of training--the model artifact--in a Cloud Storage bucket.\n",
        "\n",
        "To use a bucket mounted to the container using FUSE, you replace the `gs://` portion of the bucket's URI with the folder path `/gcs/`."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ziL1zM9E1Ptx"
      },
      "outputs": [],
      "source": [
        "# Make folder for Python training script\n",
        "if not os.path.exists(\"trainer\"):\n",
        "    os.mkdir(\"trainer\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "uIvoEYeB1sJe"
      },
      "outputs": [],
      "source": [
        "%%writefile trainer/task.py\n",
        "import sys\n",
        "import os\n",
        "import argparse\n",
        "import logging\n",
        "import hypertune\n",
        "\n",
        "import torch\n",
        "import torch.nn as nn\n",
        "import torch.optim as optim\n",
        "import torch.nn.functional as F\n",
        "import numpy as np\n",
        "import pandas as pd\n",
        "from torchvision.io import read_image\n",
        "from torch.utils.data import Dataset\n",
        "from torchvision.transforms import ToTensor\n",
        "from torch.utils.data import DataLoader\n",
        "\n",
        "parser = argparse.ArgumentParser(description='PyTorch CNN Training')\n",
        "parser.add_argument('--train_uri', dest='train_uri',\n",
        "                    type=str, help='Storage location of training CSV')\n",
        "parser.add_argument('--test_uri', dest='test_uri',\n",
        "                    type=str, help='Storage location of test CSV')\n",
        "parser.add_argument('--model-dir', dest='model_dir',\n",
        "                    default=os.getenv('AIP_MODEL_DIR'), type=str, help='Model directory')\n",
        "parser.add_argument('--batch_size', dest='batch_size',\n",
        "                    type=int, default=16, help='Batch size')\n",
        "parser.add_argument('--epochs', dest='epochs',\n",
        "                    type=int, default=20, help='Number of epochs')\n",
        "parser.add_argument('--lr', dest='lr',\n",
        "                    type=int, default=20, help='Learning rate')\n",
        "args = parser.parse_args()\n",
        "\n",
        "logging.getLogger().setLevel(logging.INFO)\n",
        "\n",
        "class CustomImageDataset(Dataset):\n",
        "    width = 28\n",
        "    height = 28\n",
        "\n",
        "    def __init__(self, data_file, transform=None, target_transform=None):\n",
        "        self.dataset = pd.read_csv(data_file)\n",
        "        self.transform = transform\n",
        "        self.target_transform = target_transform\n",
        "\n",
        "    def __len__(self):\n",
        "        return self.dataset.shape[0]\n",
        "\n",
        "    def __getitem__(self, idx):\n",
        "        label = self.dataset.at[idx, \"label\"]\n",
        "        image = self.dataset.iloc[idx,1:]\n",
        "\n",
        "        # Create a matrix from the pandas.Series\n",
        "        image = image.to_numpy() * 0.00392156862745098 # 1 / 255\n",
        "        image = image.reshape(self.width, self.height)\n",
        "        image = image.astype(float)\n",
        "        image = torch.Tensor(image)\n",
        "\n",
        "        if self.target_transform:\n",
        "            label = self.target_transform(label)\n",
        "        if self.transform:\n",
        "            image = self.transform(image)\n",
        "        return image, label\n",
        "\n",
        "# Define model\n",
        "class NeuralNetwork(nn.Module):\n",
        "    def __init__(self):\n",
        "        super().__init__()\n",
        "        self.flatten = nn.Flatten()\n",
        "        self.linear_relu_stack = nn.Sequential(\n",
        "            nn.Linear(28*28, 512, dtype=torch.float),\n",
        "            nn.ReLU(),\n",
        "            nn.Linear(512, 512),\n",
        "            nn.ReLU(),\n",
        "            nn.Linear(512, 10)\n",
        "        )\n",
        "\n",
        "    def forward(self, x):\n",
        "        x = self.flatten(x)\n",
        "        logits = self.linear_relu_stack(x)\n",
        "        return logits\n",
        "\n",
        "def get_data(train_gcs_uri, test_gcs_uri):\n",
        "\n",
        "    train_set = CustomImageDataset(train_gcs_uri)\n",
        "    test_set = CustomImageDataset(test_gcs_uri)\n",
        "\n",
        "    # HARDCODED batch_size and shuffle-can customize\n",
        "    batch_size = 64\n",
        "    shuffle = False\n",
        "\n",
        "    train_dataloader = DataLoader(train_set, batch_size=batch_size, shuffle=shuffle)\n",
        "    test_dataloader = DataLoader(test_set, batch_size=batch_size, shuffle=shuffle)\n",
        "\n",
        "    return train_dataloader, test_dataloader\n",
        "\n",
        "def get_model():\n",
        "    logging.info(\"Get model architecture\")\n",
        "    device = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")\n",
        "    gpu_id = \"0\" if torch.cuda.is_available() else None\n",
        "    logging.info(f\"Device: {device}\")\n",
        "\n",
        "    model = NeuralNetwork()\n",
        "    model.to(device)\n",
        "\n",
        "    loss = nn.CrossEntropyLoss()\n",
        "    optimizer = torch.optim.SGD(model.parameters(), lr=1e-3)\n",
        "    return model, loss, optimizer, device\n",
        "\n",
        "def train_model(model, loss_func, optimizer, train_loader, test_loader, device):\n",
        "    def train(dataloader, model, loss_fn, optimizer):\n",
        "        size = len(dataloader.dataset)\n",
        "        model.train()\n",
        "        for batch, (X, y) in enumerate(dataloader):\n",
        "            X, y = X.to(device), y.to(device)\n",
        "\n",
        "            pred = model(X)\n",
        "            loss = loss_fn(pred, y)\n",
        "\n",
        "            optimizer.zero_grad()\n",
        "            loss.backward()\n",
        "            optimizer.step()\n",
        "\n",
        "            if batch % 100 == 0:\n",
        "                loss, current = loss.item(), batch * len(X)\n",
        "                print(f\"loss: {loss:>7f}  [{current:>5d}/{size:>5d}]\")\n",
        "\n",
        "    def test(dataloader, model, loss_fn):\n",
        "        size = len(dataloader.dataset)\n",
        "        num_batches = len(dataloader)\n",
        "        model.eval()\n",
        "        test_loss, correct = 0, 0\n",
        "        with torch.no_grad():\n",
        "            for X, y in dataloader:\n",
        "                X, y = X.to(device), y.to(device)\n",
        "                pred = model(X)\n",
        "                test_loss += loss_fn(pred, y).item()\n",
        "                correct += (pred.argmax(1) == y).type(torch.float).sum().item()\n",
        "        test_loss /= num_batches\n",
        "        correct /= size\n",
        "        accuracy = 100 * correct\n",
        "        print(f\"Test Error: \\n Accuracy: {(accuracy):>0.1f}%, Avg loss: {test_loss:>8f} \\n\")\n",
        "\n",
        "    # Define a loss function and an optimizer.\n",
        "    loss_fn = nn.CrossEntropyLoss()\n",
        "    optimizer = torch.optim.SGD(model.parameters(), lr=1e-3)\n",
        "\n",
        "    epochs = 5\n",
        "    for t in range(epochs):\n",
        "        print(f\"Epoch {t+1}\\n-------------------------------\")\n",
        "        train(train_loader, model, loss_fn, optimizer)\n",
        "        test(test_loader, model, loss_fn)\n",
        "\n",
        "    # Done training\n",
        "    return model\n",
        "\n",
        "# import data from Cloud Storage\n",
        "logging.info('importing training data')\n",
        "gs_prefix = 'gs://'\n",
        "gcsfuse_prefix = '/gcs/'\n",
        "\n",
        "if args.train_uri.startswith(gs_prefix):\n",
        "    args.train_uri.replace(gs_prefix, gcsfuse_prefix)\n",
        "\n",
        "if args.test_uri.startswith(gs_prefix):\n",
        "    args.test_uri.replace(gs_prefix, gcsfuse_prefix)\n",
        "\n",
        "train_dataset, test_dataset = get_data(train_gcs_uri=args.train_uri,\n",
        "                                      test_gcs_uri=args.test_uri)\n",
        "\n",
        "logging.info('starting training')\n",
        "model, loss, optimizer, device = get_model()\n",
        "train_model(model, loss, optimizer, train_dataset, test_dataset, device)\n",
        "\n",
        "\n",
        "# export model to gcs using GCSFuse\n",
        "logging.info('start saving')\n",
        "logging.info(\"Exporting model artifacts ...\")\n",
        "gs_prefix = 'gs://'\n",
        "gcsfuse_prefix = '/gcs/'\n",
        "if args.model_dir.startswith(gs_prefix):\n",
        "    args.model_dir = args.model_dir.replace(gs_prefix, gcsfuse_prefix)\n",
        "    dirpath = os.path.split(args.model_dir)[0]\n",
        "    if not os.path.isdir(dirpath):\n",
        "        os.makedirs(dirpath)\n",
        "\n",
        "gcs_model_path = os.path.join(os.path.join(args.model_dir, 'model.pth'))\n",
        "torch.save(model.state_dict(), gcs_model_path)\n",
        "logging.info(f'Model is saved to {args.model_dir}')"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "97695e9c8655"
      },
      "source": [
        "### Create the training job\n",
        "\n",
        "Once you have written the training script to file, you now train the model. For this model, the following parameters are provided in the call to `CustomTrainingJob.run()`. Also note that the strings provided in the `args` list are defined arguments in the training script.\n",
        "\n",
        "+ The `--train_uri` and `--test_uri` arguments point towards CSV files available on a publicly available Cloud Storage bucket. The training script accesses the files using Storage FUSE.\n",
        "+ The `--model_dir` argument point towards a Storage bucket that you must provide to the script. The training script creates a new folder on the bucket to store your model artifacts."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "GDaFQpMEMeVb"
      },
      "outputs": [],
      "source": [
        "TIMESTAMP = datetime.now().strftime(\"%Y%m%d-%H%M%S\")\n",
        "JOB_DISPLAY_NAME = \"pytorch-custom-job\"\n",
        "TRAIN_IMAGE_URI = \"us-docker.pkg.dev/vertex-ai/training/pytorch-gpu.1-13:latest\"\n",
        "\n",
        "job = aiplatform.CustomTrainingJob(\n",
        "    display_name=JOB_DISPLAY_NAME,\n",
        "    script_path=\"trainer/task.py\",\n",
        "    container_uri=TRAIN_IMAGE_URI,\n",
        ")\n",
        "\n",
        "job.run(\n",
        "    replica_count=1,\n",
        "    machine_type=\"n1-standard-4\",\n",
        "    args=[\n",
        "        \"--train_uri\",\n",
        "        TRAIN_URI,\n",
        "        \"--test_uri\",\n",
        "        TEST_URI,\n",
        "        \"--model-dir\",\n",
        "        f\"{BUCKET_URI}/{BUCKET_PREFIX}/{TIMESTAMP}/\",\n",
        "    ],\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "TpV-iwP9qw9c"
      },
      "source": [
        "## Cleaning up\n",
        "\n",
        "To clean up all Google Cloud resources used in this project, you can [delete the Google Cloud\n",
        "project](https://cloud.google.com/resource-manager/docs/creating-managing-projects#shutting_down_projects) you used for the tutorial.\n",
        "\n",
        "Otherwise, you can delete the individual resources you created in this tutorial:\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "sx_vKniMq9ZX"
      },
      "outputs": [],
      "source": [
        "import os\n",
        "\n",
        "# Delete training job created\n",
        "job.delete(sync=False)\n",
        "\n",
        "# Delete Cloud Storage objects that were created\n",
        "delete_bucket = False\n",
        "if delete_bucket or os.getenv(\"IS_TESTING\"):\n",
        "    ! gsutil -m rm -r $BUCKET_URI"
      ]
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "name": "pytorch_gcs_data_training.ipynb",
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}