PyTorch Docker

Contents

PyTorch Docker#

Tensors and Dynamic neural networks in Python with strong GPU acceleration

Docker#

docker run --gpus all --rm -ti --ipc=host pytorch/pytorch

jupyter/pytorch-notebook#

# CPU
docker run -d --name pytorch-notebook -p 8888:8888 -v $PWD/pytorch-notebook:/home/jovyan/work quay.io/jupyter/pytorch-notebook
# GPU
docker run -d --name pytorch-notebook -p 8888:8888 -v $PWD/pytorch-notebook:/home/jovyan/work quay.io/jupyter/pytorch-notebook:cuda12-latest
docker exec -it pytorch-notebook pip show torch

pip#

pip install torch==2.6.0 torchvision==0.21.0 torchaudio==2.6.0 --index-url https://download.pytorch.org/whl/cpu
python -c "import torch; print(torch.rand(3,3))"

Poetry#

Poetry 是一个用于 Python 中 依赖项管理 和打包的工具

pipx install poetry
poetry new poetry-demo
# 初始化现有项目
# poetry init
cd poetry-demo
vi pyproject.toml # 放在[tool.poetry]下方

[[tool.poetry.source]]
name = "pytorch-cpu"
url = "https://download.pytorch.org/whl/cpu"
priority = "supplemental"

poetry add torch==2.6.0+cpu torchvision==0.21.0+cpu torchaudio==2.6.0+cpu
# poetry add ultralytics fastapi pydantic uvicorn requests onnxruntime
# 安装插件
poetry self add poetry-plugin-shell
poetry shell
python -c "import torch; print(f'PyTorch: {torch.__version__}, GPU: {torch.cuda.is_available()}')"
# python -c "import onnxruntime"

TensorBoard#

TensorBoard 是一个用于 TensorFlow 的可视化工具包

pip install tensorboard
tensorboard --logdir=runs --host=0.0.0.0

http://localhost:6006/
可视化类型: torch.utils.tensorboard 教程

Example#

import torch
from torch.utils.tensorboard import SummaryWriter
writer = SummaryWriter()
x = torch.arange(-5, 5, 0.1).view(-1, 1)
y = -5 * x + 0.1 * torch.randn(x.size())

model = torch.nn.Linear(1, 1)
criterion = torch.nn.MSELoss()
optimizer = torch.optim.SGD(model.parameters(), lr = 0.1)

def train_model(iter):
    for epoch in range(iter):
        y1 = model(x)
        loss = criterion(y1, y)
        writer.add_scalar("Loss/train", loss, epoch)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

train_model(10)
writer.flush()
writer.close()

Tips#

DLL load failed while importing onnxruntime_pybind11_state: 动态链接库(DLL)初始化例程失败。

下载并安装 Visual Studio 2017–2026 运行库
安装后重启电脑，再测试推理代码

Runtime Environment#

Screenshots#

References#