首页
学习
活动
专区
圈层
工具
发布
首页
学习
活动
专区
圈层
工具
MCP广场
社区首页 >专栏 >pytorch容器化环境

pytorch容器化环境

原创
作者头像
努力的gopher
发布于 2025-06-10 06:31:31
发布于 2025-06-10 06:31:31
1260
举报
文章被收录于专栏:AIOPSAIOPS

背景介绍:简化 PyTorch 训练环境的初衷

在传统的 PyTorch 开发流程中,环境配置始终是令人头疼的一环。开发者通常需要手动完成以下步骤:

  1. 安装 Miniforge 以便于 Conda 环境管理和依赖隔离;
  2. 创建虚拟环境并安装符合自身 CUDA 驱动版本的 pytorchtorchvisiontorchaudio 等依赖;
  3. 避免 numpy 等库与 PyTorch 编译版本不兼容(如 numpy 2.x 导致的初始化失败);
  4. 搭配 VS Code 的 Remote-SSH / Remote-Containers 等插件远程连接进行训练调试;
  5. 配置 GPU 驱动、CUDA 工具链、TensorBoard 和 MLflow 等训练辅助组件;
  6. 最后还要保证整个环境能够在 KubernetesDocker容器化平台中无缝运行。

上述过程不仅冗长繁琐,而且在多台机器间复现性差、兼容性容易出错。


项目目标:构建一个真正“开箱即用”的 PyTorch GPU 训练环境

为了解决上述痛点,我们基于如下需求构建了一个可高度复用的容器化训练环境:

  • 基于官方 CUDA 镜像构建,确保 GPU 加速能力
  • 预装 Miniforge,支持 Conda 虚拟环境隔离
  • 默认集成 PyTorch(支持 GPU)、TorchVision、TorchAudio 等主流框架
  • 内置 TensorBoard 与 MLflow,便于训练日志与模型管理
  • 集成 code-server,实现 Web IDE,免去 VS Code 本地配置
  • 支持 Kubernetes 部署,自动挂载共享训练目录与实验数据

通过本镜像,用户无需再手动配置复杂依赖,即可在容器启动后立即开始 GPU 加速训练,同时享受完整的 IDE 开发体验。

dockerfile如下

代码语言:dockerfile
AI代码解释
复制
 FROM nvidia/cuda:12.2.2-cudnn8-devel-ubuntu22.04# ---------- 环境变量 ----------
 ENV DEBIAN_FRONTEND=noninteractive
 ENV TZ=Asia/Shanghai
 ENV LANG=C.UTF-8
 ENV LC_ALL=C.UTF-8
 ENV PATH=/opt/conda/bin:$PATH# ---------- 安装系统依赖 + Miniforge + Git ----------
 RUN apt-get update && apt-get install -y \
     wget git curl bzip2 ca-certificates libglib2.0-0 libxext6 libsm6 libxrender1 \
     python3 python3-pip gnupg lsb-release \
     && wget https://github.com/conda-forge/miniforge/releases/latest/download/Miniforge3-Linux-x86_64.sh -O /tmp/miniforge.sh \
     && bash /tmp/miniforge.sh -b -p /opt/conda \
     && rm /tmp/miniforge.sh && apt-get clean && rm -rf /var/lib/apt/lists/*# ---------- 配置 Git (可选) ----------
 #ARG GIT_USERNAME=xxx
 #ARG GIT_PASSWORD=xxx
 #ARG GIT_REPO_DEV=gitlab.xxx.com:8864
 #ARG GIT_REPO_PRO=gitlab.xxx.com:30016#RUN git config --global user.email "pytorch@pytorch.com" && \
 #    git config --global user.name "pytorch" && \
 #    git config --global credential.helper store && \
 #    git config --global http.sslVerify false && \
 #    git config --global url."https://${GIT_USERNAME}:${GIT_PASSWORD}@".insteadOf "https://" && \
 #    git config --global url."http://${GIT_USERNAME}:${GIT_PASSWORD}@".insteadOf "http://" && \
 #    echo "http://${GIT_USERNAME}:${GIT_PASSWORD}@${GIT_REPO_DEV}" >> /root/.git-credentials && \
 #    echo "http://${GIT_USERNAME}:${GIT_PASSWORD}@${GIT_REPO_PRO}" >> /root/.git-credentials# ---------- 安装 Code-Server ----------
 RUN curl -fsSL https://code-server.dev/install.sh | bash -s -- --version 4.96.4# ---------- Conda 环境安装 PyTorch 等依赖 ----------
 RUN conda create -n ml python=3.10 -y && \
     conda run -n ml conda install -y \
         pytorch=2.2.0 torchvision=0.17.0 torchaudio=2.2.0 \
         pytorch-cuda=12.1 -c pytorch -c nvidia && \
     conda run -n ml pip install tensorboard==2.19.0 mlflow==2.22.0 && \
     conda run -n ml pip install "numpy<2.0"# ---------- 自动激活环境 ----------
 SHELL ["/bin/bash", "-c"]
 RUN echo "source /opt/conda/etc/profile.d/conda.sh && conda activate ml" >> ~/.bashrc# ---------- 默认工作目录 ----------
 RUN mkdir -p /workspace /shared/logs /shared/mlruns
 WORKDIR /workspace# ---------- 暴露 code-server 端口 ----------
 EXPOSE 8080# ---------- 容器启动时默认运行 code-server ----------
 CMD ["bash", "-c", "source /opt/conda/etc/profile.d/conda.sh && conda activate ml && /usr/bin/code-server --bind-addr 0.0.0.0:8080 --auth none /workspace"]

docker 编译和启动命令

代码语言:bash
AI代码解释
复制
 docker build -t code-server:ml-cuda  .# 宿主机需要安装nvidia驱动和NVIDIA Container Toolkit
 docker run --rm \
   --gpus all -p 8070:8080 \
   -v /data/localpv/cuda:/workspace \
   --name train-env \
   code-server:ml-cuda

安装的必要的宿主机组件(可选)

宿主机环境配置,如果没有配置好环境,无法在容器内直接使用gpu,该步骤可选,已经配置好了则无需配置

1. NVIDIA 驱动(NVIDIA Driver)

  • 这是最关键的依赖,必须安装并与显卡匹配。
  • 不需要安装 CUDA toolkit 到宿主机,驱动即可
代码语言:bash
AI代码解释
复制
# 安装方式(例如 Ubuntu)
sudo apt install nvidia-driver-535
reboot#
代码语言:bash
AI代码解释
复制
# 验证是否安装成功:
nvidia-smi

2. NVIDIA Container Toolkit(nvidia-docker 插件)

  • 这个工具用于在容器中挂载 GPU、传递 CUDA 能力。

安装步骤(Ubuntu 示例):

代码语言:bash
AI代码解释
复制
distribution=$(. /etc/os-release;echo $ID$VERSION_ID)
curl -s -L https://nvidia.github.io/nvidia-docker/gpgkey | sudo apt-key add -
curl -s -L https://nvidia.github.io/nvidia-docker/$distribution/nvidia-docker.list | sudo tee /etc/apt/sources.list.d/nvidia-docker.list
sudo apt-get update && sudo apt-get install -y nvidia-container-toolkit
sudo systemctl restart docker

原创声明:本文系作者授权腾讯云开发者社区发表,未经许可,不得转载。

如有侵权,请联系 cloudcommunity@tencent.com 删除。

原创声明:本文系作者授权腾讯云开发者社区发表,未经许可,不得转载。

如有侵权,请联系 cloudcommunity@tencent.com 删除。

评论
登录后参与评论
0 条评论
热度
最新
推荐阅读
领券
问题归档专栏文章快讯文章归档关键词归档开发者手册归档开发者手册 Section 归档