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
    my_module = MyModule(10, 20)
    sm = torch.jit.script(my_module)

保存序列化模型：使用save方法将scriptmodule保存为文件。
```
sm.save("model.pt")
```

在C++中加载序列化的PyTorch模型：使用LibTorch库加载保存的模型文件并进行推理。

#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：调用module的forward方法进行推理。

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

使用Flask或FastAPI等Web框架部署

如果你想通过网络提供服务，可以使用Flask或FastAPI等Web框架来部署你的模型。

使用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()

if torch.cuda.is_available():
    model.cuda()

@app.route('/predict', methods=['POST'])
def predict():
    data = request.json['input_data']
    input_tensor = torch.tensor(data).unsqueeze(0)
    if torch.cuda.is_available():
        input_tensor = input_tensor.cuda()
    with torch.no_grad():
        output = model(input_tensor)
    return jsonify(output.tolist())

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

运行Flask应用：

python app.py

然后你可以通过HTTP请求发送数据到http://<your-server-ip>:5000/predict来获取预测结果。

使用Docker进行部署

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

安装Docker：

sudo apt update
sudo apt install docker.io

创建Dockerfile：

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
```

使用TorchServe进行部署

TorchServe是PyTorch的模型服务工具，可以部署TorchScript模型。

安装TorchServe：

pip install torchserve
torchserve --start --model_name=model --model_path=/path/to/model.pt --runtime_version=1.8

使用Flask或Django：你可以将模型部署为一个Web服务，使用Flask或Django框架。

以上步骤可以帮助你在Linux系统下成功部署PyTorch模型。根据你的具体需求，可以选择使用TorchServe、Flask或Django等进行部署。

0 赞

0 踩