模型封装进Docker需关注加载时机、路径固化和GPU/CPU自适应:模型须用绝对路径并显式COPY;FastAPI需全局加载而非每次请求;Dockerfile应按部署环境选用对应PyTorch镜像;优先保存state_dict而非完整模型。

直接说结论:模型封装进 Docker 不是“打包就完事”,关键在「加载时机」「路径固化」和「GPU/CPU 自适应」三处。做错任意一点,容器一跑就报 FileNotFoundError、AttributeError: 'NoneType' object has no attribute 'eval' 或 GPU 显存分配失败。
模型文件必须放在容器内固定路径,且不能依赖相对路径
本地开发时你可能写 torch.load("../models/asr.pt"),但镜像构建后工作目录变了,上级目录根本不存在。Docker 内部没有你的项目根目录概念。
实操建议:
- 所有模型路径统一用绝对路径,例如
/app/models/asr.pt,并在Dockerfile中显式COPY models/ /app/models/ - 在 Python 服务代码里,用
Path(__file__).parent / "models" / "asr.pt"这类基于当前文件位置的定位,比os.getcwd()可靠得多 - 启动前加校验逻辑:读取模型文件前先
assert model_path.exists(), f"Model not found: {model_path}"
FastAPI 启动时必须全局加载模型,禁止每次请求都 load
常见错误是把 torch.load() 放在路由函数里,结果每来一个请求就重新加载一次模型——100 并发就是 100 次磁盘 IO + 显存重复分配,延迟飙升,GPU 显存爆满。
立即学习“Python免费学习笔记(深入)”;
Docker Hub (hub.docker.com)。此技能适用于所有 Docker Hub 操作,包括读取、创建、更新和删除数据。只要任务涉及 Docker Hub,就使用此技能。
正确做法:
- 在
main.py文件顶层(import 之后、app = FastAPI()之前)完成模型加载和.eval()设置 - 显式指定
map_location:用torch.load(path, map_location=device),其中device = torch.device("cuda" if torch.cuda.is_available() else "cpu") - 如果模型较大,可加日志:
logger.info(f"Loaded model on {device}, size: {os.path.getsize(path)/1024/1024:.1f} MB")
Dockerfile 必须区分 CPU/GPU 镜像基础层,不能混用
用 FROM python:3.9-slim 构建的镜像,装了 torch==2.3.0+cu121 也跑不起来 CUDA —— 基础镜像没带 NVIDIA 驱动和 libcudnn。反过来,用 FROM nvidia/cuda:12.1.1-base-ubuntu22.04 却只跑 CPU 推理,又白白多出 1.2GB 镜像体积。
选择依据很实际:
- 若服务器有 GPU 且需启用 CUDA:用
FROM pytorch/pytorch:2.3.0-cuda12.1-cudnn8-runtime,再 pip install 其他依赖 - 若纯 CPU 部署:用
FROM python:3.9-slim,pip installtorch==2.3.0+cpu(注意带+cpu后缀) - 千万别在同一个
Dockerfile里写RUN pip install torch—— 默认装的是 CPU 版,GPU 机器上会静默失败
最易被忽略的一点:模型权重文件本身是否含绝对路径或环境相关 tensor device 信息。用 torch.save(model.state_dict(), ...) 是安全的;但若保存的是完整模型对象(torch.save(model, ...)),反序列化时可能因类定义路径变化而报 ModuleNotFoundError。永远优先保存 state_dict,加载时用代码重建模型结构。

















