BentoML - Tích hợp và triển khai PyTorch model ở local

from torch import nn

import torch
import torch.nn.functional as F


class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = nn.Conv2d(3, 64, 3, padding=1)
        self.norm1 = nn.BatchNorm2d(64)
        self.conv2 = nn.Conv2d(64, 64, 3, padding=1)
        self.norm2 = nn.BatchNorm2d(64)
        self.pool1 = nn.MaxPool2d(2, 2)
        self.dropout1 = nn.Dropout2d(0.2)

        self.conv3 = nn.Conv2d(64, 128, 3, padding=1)
        self.norm3 = nn.BatchNorm2d(128)
        self.conv4 = nn.Conv2d(128, 128, 3, padding=1)
        self.norm4 = nn.BatchNorm2d(128)
        self.pool2 = nn.MaxPool2d(2, 2)
        self.dropout2 = nn.Dropout2d(0.3)

        self.conv5 = nn.Conv2d(128, 256, 3, padding=1)
        self.norm5 = nn.BatchNorm2d(256)
        self.conv6 = nn.Conv2d(256, 256, 3, padding=1)
        self.norm6 = nn.BatchNorm2d(256)
        self.pool3 = nn.MaxPool2d(2, 2)
        self.dropout3 = nn.Dropout2d(0.4)

        self.fc1 = nn.Linear(256 * 4 * 4, 256)
        self.norm7 = nn.BatchNorm1d(256)
        self.dropout4 = nn.Dropout(0.5)
        self.fc2 = nn.Linear(256, 10)

    def forward(self, x):
        x = F.relu(self.conv1(x))
        x = self.norm1(x)
        x = F.relu(self.conv2(x))
        x = self.norm2(x)
        x = self.pool1(x)
        x = self.dropout1(x)

        x = F.relu(self.conv3(x))
        x = self.norm3(x)
        x = F.relu(self.conv4(x))
        x = self.norm4(x)
        x = self.pool2(x)
        x = self.dropout2(x)

        x = F.relu(self.conv5(x))
        x = self.norm5(x)
        x = F.relu(self.conv6(x))
        x = self.norm6(x)
        x = self.pool3(x)
        x = self.dropout3(x)

        x = x.view(-1, 256 * 4 * 4)

        x = F.relu(self.fc1(x))
        x = self.norm7(x)
        x = self.dropout4(x)
        x = self.fc2(x)
        return x

import pytorch_image_classifier as PIC

if __name__ == '__main__':
    PATH = 'torch_scratch_cifar10_net.pth'

    # Tôi đang chạy bằng Macbook pro nên chỉ có CPU thôi :D
    net = Net()
    net.load_state_dict(torch.load(PATH, map_location=torch.device('cpu')))

    # Khởi tạo BentoML
    bento_svc = PIC.PytorchImageClassifier()
    bento_svc.pack('net', net)

    # Tạo Bento
    saved_path = bento_svc.save()
    print('Finished Training:', saved_path)

from PIL import Image
import os
import numpy as np

import torch
from torch.autograd import Variable
from torchvision import transforms

import bentoml
from bentoml.artifact import PytorchModelArtifact
from bentoml.adapters import ImageInput

# Tôi không chắc về dòng lệnh bên dưới, bạn có thể test xem nó có chạy hay không. Tuy nhiên dòng này là bắt buộc đối với Macbook, thiết bị mà tôi đang sử dụng
os.environ['KMP_DUPLICATE_LIB_OK'] = 'True'

# Đây là những lớp trong CIFAR-10, các lớp này sẽ được trả về sau khi gọi API hoặc chạy command line
classes = ('plane', 'car', 'bird', 'cat',
           'deer', 'dog', 'frog', 'horse', 'ship', 'truck')

device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')

# Định nghĩa hàm tăng cường dữ liệu cho ảnh trong giai đoạn test
transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
])

# Định nghĩa thư viện cần thiết cho model này
@bentoml.env(pip_dependencies=['torch', 'Pillow'])
# Định nghĩa thuộc tính để truy cập tới model, thường ta sẽ dùng "net" cho PyTorch
@bentoml.artifacts([PytorchModelArtifact('net')])
class PytorchImageClassifier(bentoml.BentoService):
    # Định nghĩa loại dữ liệu bạn sẽ truyền lên, trong trường hợp này đầu vào sẽ là một ảnh
    @bentoml.api(ImageInput)
    def predict(self, img):
        # Truy cập thuộc tính "net" và thiết lập trạng thái đánh giá (evaluation mode)
        self.artifacts.net.eval()
        input_data = []

        img = Image.fromarray(img).resize((32, 32))
        input_data.append(transform(img))

        outputs = self.artifacts.net(Variable(torch.stack(input_data)).to(device))
        _, output_classes = outputs.max(dim=1)

        return [classes[output_class] for output_class in output_classes]

import cv2

bento_svc = bentoml.load('/my/path/PytorchImageClassifier/20200708170229_EE710E')
image = cv2.imread('dog-puppy-on-garden-royalty.jpg')

result = bento_svc.predict(image)