我们在调试深度学习代码的时候，经常会因为环境的问题而头痛。
为了方便调试，把 pytorch cuda 等环境打成 docker image 是一个很不错的选择。
这是一个 Docker Image 示例，用于将深度学习环境打包：
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
| # 可以改成你需要的 pytorch 版本
FROM pytorch/pytorch:2.4.1-cuda11.8-cudnn9-devel
# 这些是常用的包
RUN apt-get update && apt-get install git zsh ffmpeg libsm6 libxext6 -y && apt-get clean && rm -rf /var/lib/apt/lists/*
WORKDIR /app
# 放在代码库的根目录下，可以安装 requirements.txt
COPY requirements.txt .
RUN pip install -r requirements.txt
# install jupyterlab
RUN pip install jupyterlab
# COPY . .
# 使用 jupyterlab 托管，可以快速启动，token 是`yourtoken`。如果你在公网使用，要考虑使用更复杂的 token。
CMD ["jupyter", "lab", "--ip=0.0.0.0", "--port=8888", "--no-browser", "--allow-root", "--NotebookApp.token=yourtoken"]
|
这个示例安装了一些基础的类库，通过 pip 即可安装 opencv-python。
将Dockerfile放在目录下，然后可以使用docker compose启动。
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
| services:
notebook:
build:
context: .
dockerfile: Dockerfile
volumes: # 还可以挂载自己需要的 dataset
- .:/app
- ~/.ssh:/root/.ssh # 支持 ssh
ports:
- 8888:8888
shm_size: '32gb'
deploy:
resources:
reservations:
devices:
- driver: nvidia
count: all
capabilities: [gpu]
|
启动命令是：docker compose up -d.
