PyTorch在Linux上如何进行模型部署 - 问答

一、环境准备：安装PyTorch及依赖

在Linux系统（如Ubuntu）上部署PyTorch模型前，需先配置基础环境：

安装Python基础工具：通过包管理器安装Python3及pip（若未安装）：
```
sudo apt update && sudo apt install python3 python3-pip
```
安装PyTorch：根据是否需要GPU加速，从PyTorch官网获取对应安装命令。例如，使用CUDA 11.8时，可通过pip安装：
```
pip install torch torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/cu118
```
若无需GPU，直接安装CPU版本：
```
pip install torch torchvision torchaudio
```

验证安装：运行Python检查PyTorch版本及CUDA可用性：

import torch
print(torch.__version__)  # 查看PyTorch版本
print(torch.cuda.is_available())  # 检查GPU支持

二、准备模型：保存或导出模型文件

模型部署需加载已训练好的权重，常见格式包括PyTorch原生.pth/.pt、TorchScript（.pt）或ONNX（.onnx）：

保存原生PyTorch模型：训练完成后，保存模型架构与权重：

# 假设模型定义为MyModel（需提前导入）
model = MyModel()
torch.save(model.state_dict(), 'model.pth')  # 仅保存权重
# 或保存完整模型（包含架构）
torch.save(model, 'full_model.pth')

转换为TorchScript格式（推荐用于生产环境，支持无Python依赖部署）：

model = MyModel()
model.load_state_dict(torch.load('model.pth'))
model.eval()  # 切换至评估模式
# 方法1：通过trace生成脚本
example_input = torch.randn(1, 3, 224, 224)  # 示例输入（需匹配模型输入维度）
traced_script = torch.jit.trace(model, example_input)
traced_script.save('model_torchscript.pt')
# 方法2：通过script生成脚本（支持动态控制流）
scripted_module = torch.jit.script(model)
scripted_module.save('model_scripted.pt')

转换为ONNX格式（跨框架部署，如TensorFlow Serving）：

import torchvision.models as models
model = models.resnet18(pretrained=True)
model.eval()
dummy_input = torch.randn(1, 3, 224, 224)  # 示例输入
torch.onnx.export(
    model, 
    dummy_input, 
    "model_onnx.onnx", 
    verbose=True,
    input_names=['input'],  # 输入节点名称
    output_names=['output'],  # 输出节点名称
    dynamic_axes={'input': {0: 'batch_size'}, 'output': {0: 'batch_size'}}  # 动态批次
)

三、本地测试：验证模型推理功能

部署前需通过脚本验证模型能否正确加载并进行推理：

import torch
from model import MyModel  # 假设模型定义在model.py中

# 加载模型
model = MyModel()
model.load_state_dict(torch.load('model.pth'))
model.eval()

# 准备输入数据（需与训练时一致，如归一化、维度调整）
input_data = torch.randn(1, 3, 224, 224)  # 示例输入

# 推理（禁用梯度计算以提升性能）
with torch.no_grad():
    output = model(input_data)

print("推理结果：", output)

四、部署方式：选择适合的场景

1. 轻量级Web服务（Flask/FastAPI）

适合快速搭建HTTP API接口，适用于小规模或原型部署：

Flask示例：

from flask import Flask, request, jsonify
import torch
from model import MyModel

app = Flask(__name__)
model = MyModel()
model.load_state_dict(torch.load('model.pth'))
model.eval()
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model.to(device)

@app.route('/predict', methods=['POST'])
def predict():
    # 获取JSON输入（假设输入为列表格式）
    data = request.json['input']
    input_tensor = torch.tensor(data).float().unsqueeze(0).to(device)  # 添加批次维度
    with torch.no_grad():
        output = model(input_tensor)
    return jsonify({'output': output.tolist()})

if __name__ == '__main__':
    app.run(host='0.0.0.0', port=5000)  # 监听所有IP的5000端口

运行后，通过curl或Postman测试：

curl -X POST http://localhost:5000/predict -H "Content-Type: application/json" -d '{"input": [[0.1, 0.2, ...], ...]}'

FastAPI示例（高性能异步框架，推荐用于生产）：
参考FastAPI文档，可实现自动API文档（Swagger UI）、异步推理等功能。

2. 官方模型服务（TorchServe）

适合大规模、高并发的生产环境，支持模型版本管理、自动扩缩容：

安装TorchServe：

pip install torchserve torch-model-archiver

打包模型：将模型文件打包为.mar（模型存档）：
```
torch-model-archiver --model-name my_model --version 1.0 --serialized-file model_torchscript.pt --handler image_classifier  # 若为图像分类模型
```
（--handler指定模型处理程序，可根据模型类型选择内置处理器或自定义）

启动TorchServe：

torchserve --start --model-store ./model_store --models my_model=my_model.mar --ncs  # --ncs启用NVIDIA Triton加速（若有GPU）

发送推理请求：

curl -X POST http://localhost:8080/predictions/my_model -T input.json  # input.json需符合模型输入格式

3. 容器化部署（Docker）

适合隔离环境、简化部署流程，尤其适合云服务器或多容器编排（如Kubernetes）：

创建Dockerfile：

FROM python:3.8-slim  # 使用轻量级镜像
WORKDIR /app
# 安装PyTorch及依赖
RUN pip install torch torchvision torchaudio flask
# 复制应用代码
COPY . /app
# 暴露端口
EXPOSE 5000
# 启动应用
CMD ["python3", "flask_app.py"]  # 替换为你的应用入口

构建并运行镜像：

docker build -t pytorch-model-server .
docker run -d -p 5000:5000 --gpus all pytorch-model-server  # 若需GPU，添加--gpus all

（--gpus all允许容器访问宿主机GPU，需安装NVIDIA Container Toolkit）

五、进阶优化：提升部署性能

模型量化：通过降低模型精度（如FP32→INT8）减少内存占用和推理延迟：

quantized_model = torch.quantization.quantize_dynamic(
    model, 
    {torch.nn.Linear},  # 量化层类型
    dtype=torch.qint8
)
quantized_model.save('quantized_model.pt')

异步推理：使用FastAPI的异步特性或消息队列（如RabbitMQ）处理高并发请求。
负载均衡：通过Nginx反向代理多个模型服务实例，提升吞吐量。

0 赞

0 踩