杂七杂八
2026-03-27
Quick Start 一个最简单的DDP Pytorch例子! 环境准备 PyTorch(gpu)>=1.5,python>=3.6 推荐使用官方打好的PyTorch docker,避免乱七八糟的环境问题影响心情。 # Dockerfile# Start FROM Nvidia PyTorch image https://ngc.nvidia.com/catalog/containers/nvidia:pytorch
# FROM nvcr.io/nvidia/pytorch:20.03-py3 代码 单GPU代码 ## main.py文件
import torch
# 构造模型
model = nn.Linear(10, 10).to(local_rank)
# 前向传播
outputs = model(torch.randn(20, 10).to(rank))
labels = torch.randn(20, 10).to(rank)
loss_fn = nn.MSELoss()
loss_fn(outputs, labels).backward()
#...