Platform For AI:Train and deploy a PyTorch model

python -m pip install "alipai>=0.4.0"

python -m pai.toolkit.config

#!/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/

- 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)

- 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"))

# 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()

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

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",
)

# 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())

Déployer avec un processeur intégré

PAI fournit un processeur PyTorch intégré qui charge et sert les modèles TorchScript. Utilisez cette méthode si votre modèle n'a pas d'exigences de service personnalisées.

Déployer le service

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)

Le déploiement prend généralement plusieurs minutes. Le SDK bloque l'exécution jusqu'à ce que le service soit prêt.

La méthode Model.deploy renvoie un objet Predictor. Utilisez Predictor.predict pour envoyer des requêtes et obtenir des prédictions.

Exécuter l'inférence

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))

Supprimer le service

p.delete_service()

Déployer avec une image personnalisée

Utilisez cette méthode si votre modèle nécessite un prétraitement ou un post-traitement personnalisé, ou s'il comporte des dépendances non couvertes par le processeur intégré.

Étape 1 : Rédiger le code du service d'inférence

Créez le fichier infer_src/run.py avec une application Flask qui charge le modèle et gère les requêtes HTTP :

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)))

Structure de répertoire attendue :

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

Étape 2 : Construire l'InferenceSpec

La fonction container_serving_spec regroupe votre code local avec une image PyTorch gérée par PAI. Le répertoire source_dir est chargé vers OSS et monté sur /ml/usercode dans le conteneur.

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())

Étape 3 : Déployer le service

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",
)

Le déploiement prend généralement plusieurs minutes. Le SDK bloque l'exécution jusqu'à ce que le service soit prêt.

Étape 4 : Exécuter l'inférence

Préparez une image de test MNIST :

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))

Envoyez l'image au service d'inférence. Le service accepte les octets d'image bruts dans le corps de la requête HTTP :

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()))

Étape 5 : Supprimer le service

predictor.delete_service()