Ubuntu上如何共享PyTorch模型

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

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

方式一 Docker镜像分发与运行

适用场景：跨机器、跨环境复现推理或部署；同时便于GPU/CPU环境统一。

步骤

准备模型文件（如：model.py、weights.pth）与推理脚本。

编写Dockerfile（示例基于官方PyTorch镜像）：

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

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

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

如需GPU，安装nvidia-docker2后运行：
```
docker run --gpus all --rm pytorch-model
```

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

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++服务或第三方系统，获得最小依赖与高性能加载。

步骤

将模型导出为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")

编写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)

编译（指定libtorch路径）：

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

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

最新问答

相关标签