import numpy as np
import os
import torch
from torch import nn
import io
from PIL import Image
import base64
from random import randint
Image.MAX_IMAGE_PIXELS = None
model_name = 'PyTorch_ResNet152.pth'
# get an available GPU device
def get_torch_device():
num_devices = torch.cuda.device_count()
if num_devices == 0:
return "cpu"
if num_devices == 1:
return "cuda:0"
else:
return f"cuda:{randint(0, num_devices-1)}"
print("Device selected for inference", get_torch_device())
device = torch.device(get_torch_device())
def load_model(model_file_name=model_name):
"""
Loads model from the serialized format
Returns
-------
model: Pytorch model instance
"""
print(f"Devcie {device}")
model_dir = os.path.dirname(os.path.realpath(__file__))
contents = os.listdir(model_dir)
if model_file_name in contents:
model.load_state_dict(torch.load(os.path.abspath(model_file_name)))
model = model.to(device)
print(f"model saved to {model.get_device()}")
return model
else:
raise FileNotFoundError(f'{model_file_name} is not found in model directory {model_dir}.')
def predict(data, model=load_model()):
"""
Returns prediction given the model and data to predict
Parameters
----------
model: Model instance returned by load_model API
data: Data format in json
Returns
-------
predictions: Output from scoring server
Format: {'prediction':output from model.predict method}
"""
img_bytes = io.BytesIO(base64.b64decode(data.encode('utf-8')))
image = Image.open(img_bytes).resize((224, 224))
arr = np.array(image)
X = torch.FloatTensor(np.transpose(arr, axes=(2, 0, 1))).unsqueeze(0)
X = X.to(device)
with torch.no_grad():
Y = model(X).to("cpu")
pred = torch.nn.functional.softmax(Y[0], dim=0).argmax().item()
return {'prediction': pred}