GPU推論

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}