### 将 PyTorch 程序 Docker 化并支持命令行参数运行 #### 创建 Dockerfile 为了将基于 PyTorch 的程序封装到 Docker 中，首先需要编写一个 `Dockerfile` 文件。该文件描述了构建镜像所需的步骤和依赖项。以下是一个典型的用于配置 PyTorch 环境的 `Dockerfile` 示例： ```dockerfile # 使用官方的基础镜像 FROM nvidia/cuda:11.7-base-ubuntu20.04 # 设置工作目录 WORKDIR /app # 安装必要的软件包 RUN apt-get update && \ apt-get install -y python3-pip python3-dev git && \ rm -rf /var/lib/apt/lists/* # 升级 pip 并安装 PyTorch 和其他依赖库 RUN pip3 install --upgrade pip && \ pip3 install torch torchvision torchaudio # 复制当前项目的代码进入容器 COPY . . # 设置默认入口点 CMD ["python3", "main.py"] ``` 此脚本选择了 NVIDIA CUDA 驱动版本 11.7 的 Ubuntu 20.04 基础镜像，并在其基础上安装 Python、pip 工具链以及最新的 PyTorch 库[^1]。 --- #### 支持命令行参数传递 为了让容器内的应用程序能够接收外部传入的命令行参数，可以通过修改 `ENTRYPOINT` 或者保持现有的 `CMD` 指令不变，同时利用 Shell 执行动态解析逻辑。下面展示了一种灵活的方法来实现这一目标： 假设您的主程序名为 `main.py`，其内容可能类似于这样： ```python import argparse import torch def parse_args(): parser = argparse.ArgumentParser(description="PyTorch Example with Arguments") parser.add_argument('--epochs', type=int, default=10, help='number of epochs to train') parser.add_argument('--lr', type=float, default=0.01, help='learning rate') return parser.parse_args() def main(args): print(f"Training for {args.epochs} epochs with learning rate={args.lr}") device = 'cuda' if torch.cuda.is_available() else 'cpu' print(f"Using device: {device}") if __name__ == '__main__': args = parse_args() main(args) ``` 在此示例中，我们定义了一个接受两个参数 (`--epochs`, `--lr`) 的 CLI 接口[^3]。 随后更新 `Dockerfile` 的最后一行为允许自定义启动选项: ```dockerfile ENTRYPOINT ["sh", "-c", "python3 main.py $@"] # 允许通过附加参数覆盖默认设置 ``` 或者保留原始 CMD 行不动即可满足大部分需求。 --- #### 构建与运行容器 完成以上准备工作之后，按照常规流程生成新镜像并向其中注入额外变量： ```bash # 构建镜像 docker build -t pytorch-app . # 测试运行效果 docker run -it --rm --gpus all pytorch-app --epochs 5 --lr 0.001 ``` 这里 `-it` 参数表示交互模式；`--gpu all` 请求访问全部可用 GPU 资源；最后附带的具体数值即为实际调用时所期望的行为表现形式[^2]。 --- #### 注意事项 - 如果希望进一步增强安全性或隔离级别，可以考虑引入只读根文件系统等高级功能。 - 对于大型项目而言，建议采用多阶段构建技术减少最终产物体积。 - 记录日志信息以便后续排查问题所在位置。 ---