PyTorch入门教程：下载、安装、配置、参数简介、DataLoader（数据迭代器）参数解析与用法合集

# 检查 GPU 驱动
nvidia-smi
# 检查已有 Python & pip
python -V && pip -V

pip3 install torch torchvision torchaudio

# 1⃣ 先安装 CUDA Runtime
sudo apt-get install cuda-12-8
# 2⃣ 安装 PyTorch
pip3 install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu128

import torch
print(torch.__version__)          # 2.7.0
print(torch.cuda.is_available())  # True
print(torch.cuda.get_device_name())

export CUDA_HOME=/usr/local/cuda-12.8
export LD_LIBRARY_PATH=$CUDA_HOME/lib64:$LD_LIBRARY_PATH
export PATH=$CUDA_HOME/bin:$PATH

cat $CUDA_HOME/include/cudnn_version.h | grep CUDNN_MAJOR -A 2

Dataset ⟶  Sampler / BatchSampler  ⟶  DataLoader(multiprocessing)  ⟶  Training Loop

# 自动寻找最优 num_workers
import os, time
from torch.utils.data import DataLoader
for nw in range(os.cpu_count()//2 + 1):
    loader = DataLoader(ds, batch_size=128, num_workers=nw, pin_memory=True)
    t0 = time.perf_counter()
    for _ in loader: pass
    print(f"{nw=}, {time.perf_counter()-t0:.2f}s")

from torch.utils.data import DataLoader, DistributedSampler
from torchvision import datasets, transforms

trans = transforms.Compose([
    transforms.RandomResizedCrop(224),
    transforms.ToTensor(),
    transforms.Normalize(mean=(0.5,)*3, std=(0.5,)*3)
])

dataset = datasets.ImageFolder("/data/imagenet-1k/", transform=trans)

sampler = DistributedSampler(dataset, shuffle=True)
loader = DataLoader(
    dataset,
    batch_size=256,
    sampler=sampler,
    num_workers=8,
    pin_memory=True,
    persistent_workers=True,
    prefetch_factor=4,
    collate_fn=lambda batch: tuple(zip(*batch))  # 返回 (imgs, labels)
)

import torch, torch.nn as nn, torch.nn.functional as F
from torch.optim import Adam
from torchvision import datasets, transforms
from torch.utils.data import DataLoader

device = 'cuda' if torch.cuda.is_available() else 'cpu'

# 1. 数据
train_tf = transforms.Compose([
    transforms.RandomCrop(32, padding=4),
    transforms.RandomHorizontalFlip(),
    transforms.ToTensor(),
    transforms.Normalize((0.5,)*3, (0.5,)*3)
])
test_tf  = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.5,)*3, (0.5,)*3)
])
train_ds = datasets.CIFAR10('.', train=True, download=True, transform=train_tf)
test_ds  = datasets.CIFAR10('.', train=False, download=True, transform=test_tf)

train_dl = DataLoader(train_ds, batch_size=128, shuffle=True,
                      num_workers=4, pin_memory=True, persistent_workers=True)
test_dl  = DataLoader(test_ds, batch_size=256, shuffle=False,
                      num_workers=2, pin_memory=True)

# 2. 模型
class SimpleCNN(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv = nn.Sequential(
            nn.Conv2d(3, 64, 3, padding=1), nn.ReLU(),
            nn.Conv2d(64, 128, 3, padding=1), nn.ReLU(),
            nn.MaxPool2d(2),
            nn.Conv2d(128, 256, 3, padding=1), nn.ReLU(),
            nn.MaxPool2d(2)
        )
        self.fc = nn.Sequential(
            nn.Flatten(),
            nn.Linear(256*8*8, 512), nn.ReLU(),
            nn.Linear(512, 10)
        )
    def forward(self, x): return self.fc(self.conv(x))

model = SimpleCNN().to(device)
opt = Adam(model.parameters(), lr=3e-4)
epochs = 20

# 3. 训练
for epoch in range(1, epochs+1):
    model.train()
    for imgs, labels in train_dl:
        imgs, labels = imgs.to(device, non_blocking=True), labels.to(device, non_blocking=True)
        logits = model(imgs)
        loss = F.cross_entropy(logits, labels)
        opt.zero_grad(set_to_none=True)
        loss.backward()
        opt.step()
    # 简易验证
    model.eval()
    correct = total = 0
    with torch.no_grad():
        for imgs, labels in test_dl:
            imgs, labels = imgs.to(device), labels.to(device)
            correct += (model(imgs).argmax(1) == labels).sum().item()
            total   += labels.size(0)
    print(f"Epoch {epoch:02d}: Acc = {correct/total*100:.2f}%")