4.8 Pytorch框架Flask部署例子

老张谋

已于 2023-08-06 18:02:06 修改

阅读量46

点赞数

分类专栏： AI--深度学习文章标签： pytorch flask 人工智能

于 2023-07-29 21:30:52 首次发布

本文链接：https://blog.csdn.net/wgf5845201314/article/details/132000251

版权

深度学习同时被 2 个专栏收录

12 篇文章 2 订阅 ¥9.90 ¥99.00

订阅专栏

超级会员免费看

AI--深度学习

26 篇文章 0 订阅

订阅专栏

本文提供了一个使用Pytorch框架的模型在Flask服务端部署的实例，包括加载模型、数据预处理及启动服务的三个步骤。同时，还介绍了与之交互的客户端flask_predict.py。

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

服务端 flask_server.py

三步走：

加载模型
数据预处理
开启服务

import io
import json
import flask
import torch
import torch
import torch.nn.functional as F
from PIL import Image
from torch import nn
#from torchvision import transforms as T
from torchvision import transforms, models, datasets
from torch.autograd import Variable

# 初始化Flask app
app = flask.Flask(__name__)
model = None
use_gpu = False

# 加载模型进来
def load_model():
    """Load the pre-trained model, you can use your model just as easily.
    """
    global model
    #这里我们直接加载官方工具包里提供的训练好的模型（代码会自动下载）括号内参数为是否下载模型对应的配置信息
    model = models.resnet18()
    num_ftrs = model.fc.in_features
    model.fc = nn.Sequential(nn.Linear(num_ftrs, 102))  # 类别数自己根据自己任务来

    #print(model)
    checkpoint = torch.load('best.pth')
    model.load_state_dict(checkpoint['state_dict'])
    #将