pytorch模型部署

最新推荐文章于 2024-07-20 18:06:55 发布

一步一步HH

最新推荐文章于 2024-07-20 18:06:55 发布

阅读量3.4k

点赞数 2

文章标签： python pytorch flask

本文链接：https://blog.csdn.net/m0_44982764/article/details/121427011

版权

本文介绍了如何在无法使用torchserve的内网环境中，利用Flask框架部署PyTorch模型。首先将模型转换为ONNX格式，然后通过ONNX运行时在Flask API中调用模型进行预测。

摘要由CSDN通过智能技术生成

Pytorch使用torchserve部署模型比较方便和简单，也方便管理。但是由于内网服务器系统的原因，无法使用torchserve。所以选择flask框架写webapi的方式，来调用模型。

1.这里首先将模型保存未onnx格式，然后使用onnx运行时调用。

import json
import re
import logging
import cv2
import torchvision.transforms as T
import numpy as np
import torch
import os
import onnxruntime
from flask import Flask,jsonify,abort,request
app = Flask(__name__)
app.config.update(RESTIFUL_JSON=dict(ensure_ascii=False))
imageDir = "E:\\data\\test\\"
ortSession = onnxruntime.InferenceSession("./resnet34.onnx")
softmax = torch.nn.Softmax()
logging.basicConfig(level=logging.DEBUG,
                    format='%(asctime)s %(filename)s[line:%(lineno)d] %(levelname)s %(message)s',
                    datefmt='%a, %d %b %Y %H:%M:%S',
                    filename='./record.log',
                    filemode='w')
with open("./index.json", "