Platform For AI:Train and deploy a PyTorch model

Prerequisites

Before you begin, ensure that you have:

• An AccessKey pair. See Create an AccessKey pair.

• A PAI workspace. See Create and manage workspaces.

• An OSS bucket. See Get started with OSS.

• Python 3.7 or later.

Step 1: Install and configure PAI SDK for Python

Install the SDK:

python -m pip install "alipai>=0.4.0"

If you get a ModuleNotFoundError, run pip install --upgrade pip first.

Configure the SDK with your AccessKey pair, workspace, and OSS bucket:

python -m pai.toolkit.config

For setup details, see Install and configure PAI SDK for Python.

Step 2: Upload training data to OSS

This tutorial uses the MNIST handwritten digit dataset to train an image classification model. PAI training jobs read input data from OSS, so upload the dataset before submitting a job.

Download the MNIST dataset

Run the following shell script to download the dataset to a local data directory:

#!/bin/sh
set -e

url_prefix="https://ossci-datasets.s3.amazonaws.com/mnist/"
# Alternative mirror if the download is slow:
# url_prefix="http://yann.lecun.com/exdb/mnist/"

mkdir -p data/MNIST/raw/

wget -nv ${url_prefix}train-images-idx3-ubyte.gz -P data/MNIST/raw/
wget -nv ${url_prefix}train-labels-idx1-ubyte.gz -P data/MNIST/raw/
wget -nv ${url_prefix}t10k-images-idx3-ubyte.gz  -P data/MNIST/raw/
wget -nv ${url_prefix}t10k-labels-idx1-ubyte.gz  -P data/MNIST/raw/

Upload to OSS

Use either ossutil or PAI SDK for Python:

• With ossutil (see ossutil 1.0):

  ossutil cp -rf ./data oss://<YourOssBucket>/mnist/data/

• With PAI SDK for Python:

  from pai.common.oss_utils import upload
  from pai.session import get_default_session
  
  sess = get_default_session()
  data_uri = upload("./data/", oss_path="mnist/data/", bucket=sess.oss_bucket)
  print(data_uri)

If you upload with ossutil, set data_uri = "oss://<YourOssBucket>/mnist/data/" explicitly when submitting the training job in Step 4.

Step 3: Write a training script

PAI mounts input data and captures model output through environment variables. Your training script reads these variables instead of using hardcoded paths.

PAI environment variables for training jobs

This tutorial starts from the MNIST example in the PyTorch repository and makes two modifications:

Modification 1 — Read input data from the environment variable

- dataset1 = datasets.MNIST("../data", train=True, download=True, transform=transform)
- dataset2 = datasets.MNIST("../data", train=False, transform=transform)

+ # Read the mounted data path from the environment variable.
+ data_path = os.environ.get("PAI_INPUT_TRAIN_DATA", "../data")
+ dataset1 = datasets.MNIST(data_path, train=True, download=True, transform=transform)
+ dataset2 = datasets.MNIST(data_path, train=False, transform=transform)

Modification 2 — Save the model in TorchScript format to the output path

The built-in PyTorch processor requires models in TorchScript format.

- if args.save_model:
-     torch.save(model.state_dict(), "mnist_cnn.pt")

+ # Save the model.
+ save_model(model)

+ def save_model(model):
+     """Convert the model to TorchScript and save it to the output path."""
+     output_model_path = os.environ.get("PAI_OUTPUT_MODEL")
+     os.makedirs(output_model_path, exist_ok=True)
+
+     m = torch.jit.script(model)
+     m.save(os.path.join(output_model_path, "mnist_cnn.pt"))

Complete training script

Save the following script to train_src/train.py:

# source: https://github.com/pytorch/examples/blob/main/mnist/main.py
from __future__ import print_function

import argparse
import os

import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torch.optim.lr_scheduler import StepLR
from torchvision import datasets, transforms


class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = nn.Conv2d(1, 32, 3, 1)
        self.conv2 = nn.Conv2d(32, 64, 3, 1)
        self.dropout1 = nn.Dropout(0.25)
        self.dropout2 = nn.Dropout(0.5)
        self.fc1 = nn.Linear(9216, 128)
        self.fc2 = nn.Linear(128, 10)

    def forward(self, x):
        x = self.conv1(x)
        x = F.relu(x)
        x = self.conv2(x)
        x = F.relu(x)
        x = F.max_pool2d(x, 2)
        x = self.dropout1(x)
        x = torch.flatten(x, 1)
        x = self.fc1(x)
        x = F.relu(x)
        x = self.dropout2(x)
        x = self.fc2(x)
        output = F.log_softmax(x, dim=1)
        return output


def train(args, model, device, train_loader, optimizer, epoch):
    model.train()
    for batch_idx, (data, target) in enumerate(train_loader):
        data, target = data.to(device), target.to(device)
        optimizer.zero_grad()
        output = model(data)
        loss = F.nll_loss(output, target)
        loss.backward()
        optimizer.step()
        if batch_idx % args.log_interval == 0:
            print(
                "Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}".format(
                    epoch,
                    batch_idx * len(data),
                    len(train_loader.dataset),
                    100.0 * batch_idx / len(train_loader),
                    loss.item(),
                )
            )
            if args.dry_run:
                break


def test(model, device, test_loader):
    model.eval()
    test_loss = 0
    correct = 0
    with torch.no_grad():
        for data, target in test_loader:
            data, target = data.to(device), target.to(device)
            output = model(data)
            test_loss += F.nll_loss(
                output, target, reduction="sum"
            ).item()  # sum up batch loss
            pred = output.argmax(
                dim=1, keepdim=True
            )  # get the index of the max log-probability
            correct += pred.eq(target.view_as(pred)).sum().item()

    test_loss /= len(test_loader.dataset)

    print(
        "\nTest set: Average loss: {:.4f}, Accuracy: {}/{} ({:.0f}%)\n".format(
            test_loss,
            correct,
            len(test_loader.dataset),
            100.0 * correct / len(test_loader.dataset),
        )
    )


def main():
    # Training settings
    parser = argparse.ArgumentParser(description="PyTorch MNIST Example")
    parser.add_argument(
        "--batch-size",
        type=int,
        default=64,
        metavar="N",
        help="input batch size for training (default: 64)",
    )
    parser.add_argument(
        "--test-batch-size",
        type=int,
        default=1000,
        metavar="N",
        help="input batch size for testing (default: 1000)",
    )
    parser.add_argument(
        "--epochs",
        type=int,
        default=14,
        metavar="N",
        help="number of epochs to train (default: 14)",
    )
    parser.add_argument(
        "--lr",
        type=float,
        default=1.0,
        metavar="LR",
        help="learning rate (default: 1.0)",
    )
    parser.add_argument(
        "--gamma",
        type=float,
        default=0.7,
        metavar="M",
        help="Learning rate step gamma (default: 0.7)",
    )
    parser.add_argument(
        "--no-cuda", action="store_true", default=False, help="disables CUDA training"
    )
    parser.add_argument(
        "--dry-run",
        action="store_true",
        default=False,
        help="quickly check a single pass",
    )
    parser.add_argument(
        "--seed", type=int, default=1, metavar="S", help="random seed (default: 1)"
    )
    parser.add_argument(
        "--log-interval",
        type=int,
        default=10,
        metavar="N",
        help="how many batches to wait before logging training status",
    )
    parser.add_argument(
        "--save-model",
        action="store_true",
        default=False,
        help="For Saving the current Model",
    )
    args = parser.parse_args()
    use_cuda = not args.no_cuda and torch.cuda.is_available()

    torch.manual_seed(args.seed)

    device = torch.device("cuda" if use_cuda else "cpu")

    train_kwargs = {"batch_size": args.batch_size}
    test_kwargs = {"batch_size": args.test_batch_size}
    if use_cuda:
        cuda_kwargs = {"num_workers": 1, "pin_memory": True, "shuffle": True}
        train_kwargs.update(cuda_kwargs)
        test_kwargs.update(cuda_kwargs)

    transform = transforms.Compose(
        [transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,))]
    )

    data_path = os.environ.get("PAI_INPUT_TRAIN_DATA", "../data")
    dataset1 = datasets.MNIST(data_path, train=True, download=True, transform=transform)
    dataset2 = datasets.MNIST(data_path, train=False, transform=transform)
    train_loader = torch.utils.data.DataLoader(dataset1, **train_kwargs)
    test_loader = torch.utils.data.DataLoader(dataset2, **test_kwargs)

    model = Net().to(device)
    optimizer = optim.Adadelta(model.parameters(), lr=args.lr)

    scheduler = StepLR(optimizer, step_size=1, gamma=args.gamma)
    for epoch in range(1, args.epochs + 1):
        train(args, model, device, train_loader, optimizer, epoch)
        test(model, device, test_loader)
        scheduler.step()

    # Save the model.
    save_model(model)


def save_model(model):
    """Convert the model to TorchScript and save it to the specified path."""
    output_model_path = os.environ.get("PAI_OUTPUT_MODEL")
    os.makedirs(output_model_path, exist_ok=True)

    m = torch.jit.script(model)
    m.save(os.path.join(output_model_path, "mnist_cnn.pt"))


if __name__ == "__main__":
    main()

Expected directory structure for the training code:

|-- train_src/           # Training script directory (uploaded to OSS before the job starts)
    |-- requirements.txt # Optional: third-party dependencies
    '-- train.py         # Training script

Step 4: Submit a training job

Estimator runs your local training script on a PAI-managed instance using a specified container image.

Key Estimator parameters

from pai.estimator import Estimator
from pai.image import retrieve

# Get a PAI-managed PyTorch 1.8 GPU image.
image_uri = retrieve(
    "PyTorch", framework_version="1.8PAI", accelerator_type="GPU"
).image_uri
print(image_uri)

est = Estimator(
    # Startup command; working directory is /ml/usercode.
    command="python train.py $PAI_USER_ARGS",
    # Local training script directory to upload.
    source_dir="./train_src/",
    image_uri=image_uri,
    # Instance type: 4 vCPU, 15 GB RAM, 1x NVIDIA T4.
    instance_type="ecs.gn6i-c4g1.xlarge",
    # Hyperparameters are passed to the script as CLI flags via $PAI_USER_ARGS.
    hyperparameters={
        "epochs": 5,
        "batch-size": 64 * 4,
        "lr": 0.5,
    },
    # Extract the loss metric from training logs.
    metric_definitions=[
        {
            "Name": "loss",
            "Regex": r".*loss=([-+]?[0-9]*.?[0-9]+(?:[eE][-+]?[0-9]+)?).*",
        },
    ],
    base_job_name="pytorch_mnist",
)

Submit the job and wait for it to complete:

# If you uploaded with ossutil, set data_uri explicitly:
# data_uri = "oss://<YourOssBucket>/mnist/data/"

est.fit(
    inputs={
        "train_data": data_uri,
    }
)

# OSS path of the trained model.
print("Training job output model path:")
print(est.model_data())

After submission, the SDK prints a link to the job details page and streams training logs until the job finishes. For more information, see Submit a training job.

Step 5: Deploy an inference service

After the training job finishes, use est.model_data() to get the OSS path of the trained model, then deploy it as an online inference service on EAS.

Deploying a service requires two components:

• Model assets: the trained model files from OSS.

• Runtime: instructions for loading the model and serving predictions — either a built-in processor or a custom container image.

Choose a deployment method based on your requirements:

from pai.model import Model, InferenceSpec
from pai.predictor import Predictor
from pai.common.utils import random_str

m = Model(
    model_data=est.model_data(),
    # Use the built-in PyTorch processor.
    inference_spec=InferenceSpec(processor="pytorch_cpu_1.10"),
)

p: Predictor = m.deploy(
    service_name="tutorial_pt_mnist_proc_{}".format(random_str(6)),
    instance_type="ecs.c6.xlarge",
)

print(p.service_name)
print(p.service_status)

import numpy as np

# Input shape: (batch_size, channels, height, width) as float32.
dummy_input = np.random.rand(2, 1, 28, 28).astype(np.float32)

res = p.predict(dummy_input)
print(res)
print(np.argmax(res, 1))

p.delete_service()

import json
from flask import Flask, request
from PIL import Image
import os
import torch
import torchvision.transforms as transforms
import numpy as np
import io

app = Flask(__name__)
# The model is loaded from this path by default.
MODEL_PATH = "/eas/workspace/model/"

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
model = torch.jit.load(os.path.join(MODEL_PATH, "mnist_cnn.pt"), map_location=device).to(device)
transform = transforms.Compose(
    [transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,))]
)


@app.route("/", methods=["POST"])
def predict():
    # Preprocess the image.
    im = Image.open(io.BytesIO(request.data))
    input_tensor = transform(im).to(device)
    input_tensor.unsqueeze_(0)
    # Run inference.
    output_tensor = model(input_tensor)
    pred_res = output_tensor.detach().cpu().numpy()[0]
    return json.dumps(pred_res.tolist())


if __name__ == '__main__':
    app.run(host="0.0.0.0", port=int(os.environ.get("LISTENING_PORT", 8000)))

|-- infer_src/           # Inference service code directory (uploaded to OSS before deployment)
    |-- requirements.txt # Optional: third-party dependencies
    '-- run.py           # Inference server script

from pai.model import InferenceSpec, container_serving_spec
from pai.image import retrieve, ImageScope

torch_image_uri = retrieve(
    "PyTorch", framework_version="latest", image_scope=ImageScope.INFERENCE
).image_uri

inf_spec = container_serving_spec(
    command="python run.py",
    source_dir="./infer_src/",
    image_uri=torch_image_uri,
    requirements=["flask==2.0.0", "Werkzeug==2.2.2", "pillow", "torchvision"],
)
print(inf_spec.to_dict())

from pai.model import Model
from pai.common.utils import random_str

m = Model(
    model_data=est.model_data(),
    inference_spec=inf_spec,
)

predictor = m.deploy(
    service_name="torch_mnist_script_container_{}".format(random_str(6)),
    instance_type="ecs.c6.xlarge",
)

import base64
from PIL import Image
import io

# raw_data is an MNIST image of the digit 9.
raw_data = base64.b64decode(b"/9j/4AAQSkZJRgABAQAAAQABAAD/2wBDAAgGBgcGBQgHBwcJCQgKDBQNDAsLDBkSEw8UHRofHh0aHBwgJC4nICIsIxwcKDcpLDAxNDQ0Hyc5PTgyPC4zNDL/wAALCAAcABwBAREA/8QAHwAAAQUBAQEBAQEAAAAAAAAAAAECAwQFBgcICQoL/8QAtRAAAgEDAwIEAwUFBAQAAAF9AQIDAAQRBRIhMUEGE1FhByJxFDKBkaEII0KxwRVS0fAkM2JyggkKFhcYGRolJicoKSo0NTY3ODk6Q0RFRkdISUpTVFVWV1hZWmNkZWZnaGlqc3R1dnd4eXqDhIWGh4iJipKTlJWWl5iZmqKjpKWmp6ipqrKztLW2t7i5usLDxMXGx8jJytLT1NXW19jZ2uHi4+Tl5ufo6erx8vP09fb3+Pn6/9oACAEBAAA/APn+rVhpmoarP5GnWNzeTYz5dvE0jfkoJovNMv8ATmK3tjc2zByhE8TIQw6jkdR6VVq9oumPrWuWGlxyLG95cRwK7dFLMFyfzr3aXwp4ltAfB3gWwudI01JNuoa7eZhku5AMHafvFOw2Dn6ZJ4z4yeLk1HUbXwrZSSy2Oh5heeaQu88wG1mLHk4wR9c+1eXUqsVYMpIIOQR2r1D4QazqOs/FnSG1fVLi9ZI5vL+2TNKc+U2ApYnB7/hXml5LLNfXEsxLSvIzOSMEsTk1DRVnT7+60vULe/spmhureQSRSL1Vh0NWNd1mXX9ZuNUuLe2gmuCGkS2QohbABbBJwTjJ9yelZ1f/2Q==")

im = Image.open(io.BytesIO(raw_data))

from pai.predictor import RawResponse
import numpy as np

resp: RawResponse = predictor.raw_predict(data=raw_data)
print(resp.json())
print(np.argmax(resp.json()))

predictor.delete_service()

What's next

• Submit a training job — learn more about Estimator parameters, instance types, and job monitoring.

• Deploy inference services — explore advanced deployment options and autoscaling.

• PyTorch processor — supported model formats and processor configurations.