Ubuntu上如何共享PyTorch模型

Ubuntu上共享PyTorch模型的常用方式

在Ubuntu环境中，常见的模型共享方式包括：打包为Docker镜像进行分发与运行、用TorchServe发布成在线推理服务、通过NFS/Samba在局域网内共享模型文件，以及将模型导出为TorchScript .pt并集成到**C++共享库（.so）**供其他程序调用。下面给出可直接复用的操作要点与命令示例。

方式一 Docker镜像分发与运行

• 适用场景：跨机器、跨环境复现推理或部署；同时便于GPU/CPU环境统一。
• 步骤
  1. 准备模型文件（如：model.py、weights.pth）与推理脚本。
  2. 编写Dockerfile（示例基于官方PyTorch镜像）：

     FROM pytorch/pytorch:latest
     RUN pip install numpy torchvision matplotlib  # 按需增减依赖
     COPY model.py weights.pth /app/
     WORKDIR /app
     CMD ["python", "model.py"]
     

  3. 构建镜像并运行（CPU示例）：

     docker build -t pytorch-model .
     docker run --rm pytorch-model
     

  4. 如需GPU，安装nvidia-docker2后运行：

     docker run --gpus all --rm pytorch-model
     

  5. 如需挂载主机目录（便于热更新权重/数据）：

     docker run --gpus all --rm -v "$(pwd):/data:rw" pytorch-model /data/model.py /data/weights.pth
     

  以上流程与命令要点见Docker打包PyTorch模型与GPU挂载实践。

方式二 TorchServe在线发布模型

• 适用场景：需要对外提供REST/gRPC推理接口、支持模型版本管理与动态扩缩。
• 步骤
  1. 安装环境：建议基于官方镜像；TorchServe依赖Java 11（Ubuntu可安装openjdk-11-jdk）。
  2. 使用torch-model-archiver打包模型（示例以MNIST为例，替换为你自己的模型脚本、权重与handler）：

     torch-model-archiver \
       --model-name mnist \
       --version 1.0 \
       --model-file examples/image_classifier/mnist/mnist.py \
       --serialized-file examples/image_classifier/mnist/mnist_cnn.pt \
       --handler examples/image_classifier/mnist/mnist_handler.py
     # 生成 mnist.mar
     

  3. 启动TorchServe容器（映射端口：8080推理、8081管理）：

     docker run --rm -it \
       -p 3000:8080 -p 3001:8081 \
       pytorch/torchserve
     

  4. 访问管理接口（示例）：http://:3001，推理接口：http://:3000/predictions/mnist
     上述TorchServe环境、打包与端口说明可参考官方与实操教程。

方式三 局域网文件共享 NFS与Samba

• 适用场景：多台Ubuntu机器共享同一份模型目录，便于分布式训练/多进程读取。
• NFS（Linux间高性能共享）
  • 服务端：

    sudo apt install nfs-kernel-server
    sudo mkdir -p /path/share
    echo '/path/share *(rw,sync,no_root_squash,no_subtree_check)' | sudo tee -a /etc/exports
    sudo exportfs -a
    sudo systemctl restart nfs-kernel-server
    

  • 客户端：

    sudo apt install nfs-common
    sudo mkdir -p /path/share
    echo 'server_ip:/path/share /path/share nfs defaults 0 0' | sudo tee -a /etc/fstab
    sudo mount -a
    

• Samba（兼容Windows与Linux）
  • 服务端：

    sudo apt install samba
    sudo mkdir -p /home/share && sudo chmod 777 /home/share
    echo '[share]
    path=/home/share
    available=yes
    browseable=yes
    public=yes
    writable=yes
    guest-ok=yes' | sudo tee -a /etc/samba/smb.conf
    sudo systemctl restart smbd
    

  • 客户端挂载（Ubuntu）：

    sudo apt install cifs-utils
    sudo mount.cifs //SAMBA_IP/share /mnt/share -o user=username,pass=password,rw,file_mode=0777,dir_mode=0777
    

  以上NFS/Samba配置与挂载命令可直接用于模型目录共享。

方式四 C++共享库集成 TorchScript .so

• 适用场景：将模型嵌入C/C++服务或第三方系统，获得最小依赖与高性能加载。
• 步骤
  1. 将模型导出为TorchScript：

     import torch
     class MyModule(torch.nn.Module):
         def __init__(self):
             super().__init__()
             self.linear = torch.nn.Linear(10, 1)
         def forward(self, x):
             return self.linear(x)
     model = MyModule()
     model.eval()
     traced = torch.jit.trace(model, torch.randn(1, 10))
     traced.save("model.pt")
     

  2. 编写C++加载与推理代码（示例头/源略），CMakeLists.txt关键片段：

     find_package(Torch REQUIRED)
     add_library(pytorchtest SHARED test.cpp)
     target_link_libraries(pytorchtest "${TORCH_LIBRARIES}")
     set_property(TARGET pytorchtest PROPERTY CXX_STANDARD 11)
     

  3. 编译（指定libtorch路径）：

     mkdir build && cd build
     cmake -DCMAKE_PREFIX_PATH=/path/to/libtorch ..
     make
     

  4. 运行生成的可执行文件或.so将加载并执行model.pt。
     该流程涵盖模型保存、C++加载与CMake配置，适合发布为系统级推理组件。