背景介绍:简化 PyTorch 训练环境的初衷
在传统的 PyTorch 开发流程中,环境配置始终是令人头疼的一环。开发者通常需要手动完成以下步骤:
pytorch
、torchvision
、torchaudio
等依赖;numpy
等库与 PyTorch 编译版本不兼容(如 numpy 2.x 导致的初始化失败);上述过程不仅冗长繁琐,而且在多台机器间复现性差、兼容性容易出错。
为了解决上述痛点,我们基于如下需求构建了一个可高度复用的容器化训练环境:
通过本镜像,用户无需再手动配置复杂依赖,即可在容器启动后立即开始 GPU 加速训练,同时享受完整的 IDE 开发体验。
dockerfile如下
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 编译和启动命令
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,该步骤可选,已经配置好了则无需配置
# 安装方式(例如 Ubuntu)
sudo apt install nvidia-driver-535
reboot#
# 验证是否安装成功:
nvidia-smi
nvidia-docker
插件)安装步骤(Ubuntu 示例):
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 删除。