47、生产环境部署：PyTorch 模型服务与导出

最新推荐文章于 2025-12-05 02:25:19 发布

onion

最新推荐文章于 2025-12-05 02:25:19 发布

阅读量61

点赞数

CC 4.0 BY-SA版权

分类专栏： PyTorch深度学习实战文章标签： PyTorch 模型部署模型导出

本文链接：https://blog.youkuaiyun.com/onion/article/details/151034199

PyTorch深度学习实战专栏收录该内容

51 篇文章 ¥499.90

订阅专栏¥69.90

会员秒杀 ¥9.9 重磅福利

超级会员免费看

生产环境部署：PyTorch 模型服务与导出

1. 初始服务器实现

我们首先有一个简单的服务器实现，用于加载模型并进行推理。以下是相关代码：

import sys
import torch
import numpy as np
from flask import Flask, request, jsonify

app = Flask(__name__)
model = ...  # 这里应该是模型的定义
model.load_state_dict(torch.load(sys.argv[1], map_location='cpu')['model_state'])
model.eval()

def run_inference(in_tensor):
    with torch.no_grad():
        # LunaModel takes a batch and outputs a tuple (scores, probs)
        out_tensor = model(in_tensor.unsqueeze(0))[1].squeeze(0)
        probs = out_tensor.tolist()
        out = {'prob_malignant': probs[1]}
    return out

@app.route("/predict", methods=["POST"])
def predict():
    meta = json.load(request.files['meta'])
    blob = request.files['blob'].read()