Linux下PyTorch模型如何部署

在Linux下部署PyTorch模型有多种方法，以下是一些常见的部署流程：

使用TorchScript进行部署

TorchScript是PyTorch的一种序列化格式，可用于在C++、Python等平台上进行部署。

模型转换：

Tracing：适用于模型中没有控制流的情况。通过给模型一组输入，走一遍推理网络，然后用torch.jit.trace记录路径上的信息并保存为torch.jit.scriptmodule。
```
import torch
import torchvision

model = torchvision.models.resnet18()
example = torch.rand(1, 3, 224, 224)
traced_script_module = torch.jit.trace(model, example)
```

Scripting：适用于有控制流的模型。在Torch脚本中编写模型并相应地注释模型，通过torch.jit.script编译模块，将其转换为scriptmodule。

class MyModule(torch.nn.Module):
    def __init__(self, n, m):
        super(MyModule, self).__init__()
        self.weight = torch.nn.Parameter(torch.rand(n, m))

    def forward(self, input):
        if input.sum() > 0:
            output = self.weight.mv(input)
        else:
            output = self.weight + input
        return output

my_module = MyModule(10, 20)
sm = torch.jit.script(my_module)

保存序列化模型：
```
sm.save("model.pt")
```

在C++中加载序列化的PyTorch模型：

#include <torch/torch.h>

int main() {
    torch::jit::script::Module module;
    try {
        module = torch::jit::load("model.pt");
    }
    catch (const c10::Error& e) {
        std::cerr << "error loading the model\n";
        return -1;
    }
    // Use the module for inference
}

执行Script Module：

at::Tensor input = torch::rand({1, 3, 224, 224});
at::Tensor output = module.forward({input}).toTensor();

使用Web框架部署（可选）

如果你想通过网络提供服务，可以使用Flask或FastAPI等Web框架来部署你的模型。以下是一个使用Flask的简单示例：

安装Flask：
```
pip install flask
```

创建Flask应用：

from flask import Flask, request, jsonify
import torch
import torchvision.transforms as transforms
from PIL import Image

app = Flask(__name__)
# 加载模型
model = torch.load('path_to_your_model.pt', map_location=torch.device('cpu'))
model.eval()
# 定义图像预处理
transform = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
])

@app.route('/predict', methods=['POST'])
def predict():
    if request.method == 'POST':
        file = request.files['image']
        image = Image.open(file.stream)
        image = transform(image).unsqueeze(0)  # 添加batch维度
        with torch.no_grad():
            output = model(image)
        _, predicted_idx = torch.max(output, 1)
        return jsonify({'prediction': int(predicted_idx.item())})

if __name__ == '__main__':
    app.run(host='0.0.0.0', port=5000)

运行Flask应用：
```
python app.py
```

测试模型：使用curl或Postman来测试你的模型。

curl -X POST -F "image=@path_to_your_image.jpg" http://localhost:5000/predict

使用Docker进行部署

Docker是一个容器化平台，可以让你轻松地打包和部署应用程序。

安装Docker：

sudo apt update
sudo apt install docker.io

创建Dockerfile：

# 使用官方的PyTorch镜像作为基础镜像
FROM pytorch/pytorch:latest
# 设置工作目录
WORKDIR /app
# 复制当前目录下的所有文件到工作目录
COPY . /app
# 安装依赖项
RUN pip install --no-cache-dir -r requirements.txt
# 暴露端口(如果需要)
EXPOSE 80
# 运行应用
CMD ["python", "your_script.py"]

创建requirements.txt：
```
torch
torchvision
numpy
```
构建Docker镜像：
```
docker build -t your-image-name .
```
运行Docker容器：
```
docker run -p 80:80 your-image-name
```

以上是Linux下部署PyTorch模型的几种常见方法，你可以根据自己的需求选择合适的方式进行部署。

使用TorchScript进行部署

使用Web框架部署（可选）

使用Docker进行部署

最新问答

相关标签