Linux系统下PyTorch模型如何部署

2025年7月18日 互联网

在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 
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://: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等进行部署。

最新文章