소개

Docker 이미지를 Lambda에서 실행할 수 있게 되었고, 이미지도 10GB까지가 되었기 때문에, 기계 학습의 추론 API의 로직을 Lambda로 실행할 수 있을까를 시도했습니다.
이번에는 PyTorch로 만든 MNIST 모델을 넣어 보겠습니다.

환경

Python 3.8

PyTorch 1.6.0 CPU

PyTorch 1.1.8

API Gateway + Lambda에서 이미지 수신

우선은, API이므로 이미지를 받아 Lambda 로 처리할 수 있는 형태까지 가져갈 필요가 있습니다. 이쪽은 길어졌으므로 다른 기사에 기재합니다. 다음 기사의 내용이 완료되었다고 가정합니다.
API Gateway+Lambda에서 이미지를 받고 Pillow로 처리(AWS CDK) - Qiita

모델링

이 기사에서는 자세한 내용을 다루지 않지만, 다음과 같이 별도의 디렉토리를 작성하여 PyTorch 모델의 가중치가 저장된 파일을 작성합니다.

$ mkdir models
$ cd models
$ touch train.py

다음 패키지를 설치합니다.

torch

torchvision

학습용 Python 파일을 만듭니다.

models/train.py

from pytorch_lightning.metrics.functional import accuracy
import pytorch_lightning as pl
from torchvision import transforms, datasets
import torch.nn.functional as F
import torch.nn as nn
from torchvision import datasets
import torch

import torchvision
from torchvision import transforms


transform = transforms.Compose([
    transforms.ToTensor()
])
train_val = datasets.MNIST(
    './', train=True, download=True, transform=transform)
test = datasets.MNIST('./', train=False, download=True, transform=transform)

n_train, n_val = 50000, 10000
train, val = torch.utils.data.random_split(train_val, [n_train, n_val])

batch_size = 1028

train_loader = torch.utils.data.DataLoader(
    train, batch_size, shuffle=True, drop_last=True)
val_loader = torch.utils.data.DataLoader(val, batch_size)
test_loader = torch.utils.data.DataLoader(test, batch_size)


class Net(pl.LightningModule):

    def __init__(self):
        super().__init__()

        self.conv = nn.Conv2d(in_channels=1, out_channels=3,
                              kernel_size=3, padding=1)
        self.bn = nn.BatchNorm2d(3)
        self.fc = nn.Linear(588, 10)

    def forward(self, x):
        h = self.conv(x)
        h = F.relu(h)
        h = self.bn(h)
        h = F.max_pool2d(h, kernel_size=2, stride=2)
        h = h.view(-1, 588)
        h = self.fc(h)
        return h

    def training_step(self, batch, batch_idx):
        x, t = batch
        y = self(x)
        loss = F.cross_entropy(y, t)
        self.log('train_loss', loss, on_step=True,
                 on_epoch=True, prog_bar=True)
        self.log('train_acc', accuracy(y, t), on_step=True,
                 on_epoch=True, prog_bar=True)
        return loss

    def validation_step(self, batch, batch_idx):
        x, t = batch
        y = self(x)
        loss = F.cross_entropy(y, t)
        self.log('val_loss', loss, on_step=False, on_epoch=True)
        self.log('val_acc', accuracy(y, t), on_step=False, on_epoch=True)
        return loss

    def test_step(self, batch, batch_idx):
        x, t = batch
        y = self(x)
        loss = F.cross_entropy(y, t)
        self.log('test_loss', loss, on_step=False, on_epoch=True)
        self.log('test_acc', accuracy(y, t), on_step=False, on_epoch=True)
        return loss

    def configure_optimizers(self):
        optimizer = torch.optim.SGD(self.parameters(), lr=0.01)
        return optimizer


net = Net()
trainer = pl.Trainer(max_epochs=5, gpus=0, deterministic=True)
trainer.fit(net, train_loader, val_loader)
results = trainer.test(test_dataloaders=test_loader)
torch.save(net.state_dict(), 'mnist.pt')

방금 만든 학습용 Python 파일을 실행합니다.

$ python train.py
$ ls
MNIST           lightning_logs  mnist.pt       train.py

학습이 성공하면 mnist.pt 파일이 생성됩니다.

Lambda 구축

다른 기사에서 해설한 내용과 거의 같기 때문에 차이를 해설합니다.

Dockerfile 변경

필요한 패키지를 설치해야 하므로 Dockerfile을 변경합니다.

src/Dockerfile

FROM public.ecr.aws/lambda/python:3.8
RUN pip install torch==1.6.0+cpu torchvision==0.7.0+cpu -f https://download.pytorch.org/whl/torch_stable.html && pip install pillow && pip install pytorch-lightning
COPY mnist.pt ./
COPY app.py ./
CMD [ "app.handler" ]

모델 배치

학습하고 작성한 모델 파일을 src 디렉토리에 배치하십시오.

# プロジェクトの Root ディレクトリーに戻る
$ cd ..
$ cp models/mnist.pt src/mnist.pt

Lambda 코드 변경

app.py 를 다음과 같이 변경합니다.
Pillow로 로드한 파일을 PyTorch 모델 + 미리 학습한 가중치로 추론하여 결과를 반환하도록 합니다.

src/app.py

import base64
from io import BytesIO
import torch
from torchvision import transforms
import pytorch_lightning as pl
import torch.nn as nn
import torch.nn.functional as F
import json
from PIL import Image


class Net(pl.LightningModule):

    def __init__(self):
        super().__init__()

        self.conv = nn.Conv2d(in_channels=1, out_channels=3,
                              kernel_size=3, padding=1)
        self.bn = nn.BatchNorm2d(3)
        self.fc = nn.Linear(588, 10)

    def forward(self, x):
        h = self.conv(x)
        h = F.relu(h)
        h = self.bn(h)
        h = F.max_pool2d(h, kernel_size=2, stride=2)
        h = h.view(-1, 588)
        h = self.fc(h)
        return h


def handler(event, context):
    data = event.get('body', '')
    data = BytesIO(base64.b64decode(data))
    image = Image.open(data)

    net = Net().cpu().eval()
    net.load_state_dict(torch.load(
        'mnist.pt', map_location=torch.device('cpu')))

    transform = transforms.Compose([
        transforms.ToTensor()
    ])

    x = transform(image)

    y = net(x.unsqueeze(0))
    y = F.softmax(y)
    y = torch.argmax(y)

    return {
        'statusCode': 200,
        'body': json.dumps({
            'number': '{}'.format(y.item()),
        }),
    }

배포

배포는 cdk를 사용하여 수행합니다.

# プロジェクトの Root ディレクトリーに戻ります
$ cd ..
$ cdk deploy

실행

배포에 성공하면 Insomnia에서 숫자로 작성된 이미지 파일을 전송해 봅니다.

처음은 이쪽의 1 의 화상을 송신해 보겠습니다.


처음 시작하는 데 시간이 걸리지 만 두 번째 이후에는 예상보다 빠른 속도로 응답이 반환됩니다.

5

8

결론

PyTorch + API Gateway + Lambda로 추론 API를 구축하는 방법을 설명했습니다.

mnist라면 충분히 동작하는 API가 된 것이 아닐까 생각합니다.mnist.pt 는 26KB 로 작았기 때문에, 좀 더 큰 모델을 구축했을 경우에 어떻게 되는지는 별도 시험해 보고 싶습니다.