AI开发平台ModelArts-PyTorch:推理代码

时间：2023-11-01 16:20:34

AI开发平台ModelArts

推理代码

在模型代码推理文件customize_service.py中，需要添加一个子类，该子类继承对应模型类型的父类，各模型类型的父类名称和导入语句如请参考表1。

from PIL import Imageimport logfrom model_service.pytorch_model_service import PTServingBaseServiceimport torch.nn.functional as Fimport torch.nn as nnimport torchimport jsonimport numpy as nplogger = log.getLogger(__name__)import torchvision.transforms as transforms# 定义模型预处理infer_transformation = transforms.Compose([    transforms.Resize((28,28)),    # 需要处理成pytorch tensor    transforms.ToTensor()])import osclass PTVisionService(PTServingBaseService):    def __init__(self, model_name, model_path):        # 调用父类构造方法        super(PTVisionService, self).__init__(model_name, model_path)        # 调用自定义函数加载模型        self.model = Mnist(model_path)        # 加载标签        self.label = [0,1,2,3,4,5,6,7,8,9]        # 亦可通过文件标签文件加载        # model目录下放置label.json文件，此处读取        dir_path = os.path.dirname(os.path.realpath(self.model_path))        with open(os.path.join(dir_path, 'label.json')) as f:            self.label = json.load(f)    def _preprocess(self, data):        preprocessed_data = {}        for k, v in data.items():            input_batch = []            for file_name, file_content in v.items():                with Image.open(file_content) as image1:                    # 灰度处理                    image1 = image1.convert("L")                    if torch.cuda.is_available():                        input_batch.append(infer_transformation(image1).cuda())                    else:                        input_batch.append(infer_transformation(image1))            input_batch_var = torch.autograd.Variable(torch.stack(input_batch, dim=0), volatile=True)            print(input_batch_var.shape)            preprocessed_data[k] = input_batch_var        return preprocessed_data    def _postprocess(self, data):        results = []        for k, v in data.items():            result = torch.argmax(v[0])            result = {k: self.label[result]}            results.append(result)        return results    def _inference(self, data):        result = {}        for k, v in data.items():            result[k] = self.model(v)        return resultclass Net(nn.Module):    def __init__(self):        super(Net, self).__init__()        self.hidden1 = nn.Linear(784, 5120, bias=False)        self.output = nn.Linear(5120, 10, bias=False)    def forward(self, x):        x = x.view(x.size()[0], -1)        x = F.relu((self.hidden1(x)))        x = F.dropout(x, 0.2)        x = self.output(x)        return F.log_softmax(x)def Mnist(model_path, **kwargs):    # 生成网络    model = Net()    # 加载模型    if torch.cuda.is_available():        device = torch.device('cuda')        model.load_state_dict(torch.load(model_path, map_location="cuda:0"))    else:        device = torch.device('cpu')        model.load_state_dict(torch.load(model_path, map_location=device))    # CPU或者GPU映射    model.to(device)    # 声明为推理模式    model.eval()    return model

上一篇：AI开发平台ModelArts-TensorFlow 2.1:训练并保存模型

下一篇：AI开发平台ModelArts-TensorFlow 2.1:训练并保存模型