NexonTalk 20180605: DeepLearning with CNN on AWS Lambda

lambda_function.py

import torch #파이토치
import torchvision.datasets as vdatasets #데이터셋 모음
import torchvision.utils as vutils #유틸리티함수
import torchvision.transforms as vtransform #변환함수
import numpy as np #수치연산
import torch.nn as nn
import torch.nn.functional as F
from PIL import Image
from io import BytesIO
import base64

transform = vtransform.Compose(
    [vtransform.Resize((28,28))
    ,vtransform.ToTensor()]
)

device = torch.device('cpu')

class CNN(nn.Module):
    def __init__(self):
        super(CNN, self).__init__()
        self.conv1 = nn.Conv2d(1, 32, 3)
        self.pool = nn.MaxPool2d(2, 2)
        self.conv2 = nn.Conv2d(32, 64, 3)
        self.fc1 = nn.Linear(1600, 1024)
        self.fc2 = nn.Linear(1024, 10)

    def forward(self, inputs):
        
        # 1st conv 
        x = self.conv1(inputs)
        x = F.relu(x)
        x = self.pool(x)
        
        #TODO: 2nd conv
        x = self.conv2(x)
        x = F.relu(x)
        x = self.pool(x)
        
        x = x.view(inputs.size()[0], -1)
        x = F.relu(self.fc1(x))
        x = self.fc2(x)
        
        return x

def predict(model, img):
    model.eval()
    my_image_tensor = transform(img)
    tensor = my_image_tensor.unsqueeze(0)
    # 모델이 무슨 숫자로 읽었는지 반환합니다.
    return model(tensor).argmax(dim=1).item()

def lambda_handler(event=None, context=None):
    binary = event['base64Image']
    img = Image.open(BytesIO(base64.b64decode(binary))).convert("L")
    model = CNN()
    model.load_state_dict(
        torch.load(
            "cnn.pth.tar", 
            map_location=lambda storage, loc: storage
            )
        )
    return {
        'result': predict(model, img)
    }