Skip to content

快速参考卡片

把最常用的命令和代码浓缩成一页,建议打印或收藏。所有命令默认在激活了 PyTorch 的虚拟环境的终端中执行。

安装与环境

场景命令
CPU 版(Win/mac/Linux)pip install torch torchvision
NVIDIA GPU(CUDA 12.1)pip install torch torchvision --index-url https://download.pytorch.org/whl/cu121
国内镜像(CPU 版)pip install torch torchvision -i https://pypi.tuna.tsinghua.edu.cn/simple
验证安装python -c "import torch; print(torch.__version__, torch.cuda.is_available())"
查驱动支持的最高 CUDAnvidia-smi

张量速查

python
torch.tensor([[1., 2.]])       # 从数据创建
torch.zeros(3, 4) / rand / ones / arange(0, 10, 2)
torch.zeros_like(x)            # 继承 x 的 dtype+device(推荐)
x.shape / x.dtype / x.device / x.numel()
x.view(2, -1) / x.reshape(...) / x.flatten(start_dim=1)
x.unsqueeze(0)                 # 单样本补 batch 维
x.permute(0, 2, 3, 1)          # 调维度顺序(NCHW↔NHWC)
x.to(device) / x.to(torch.float32)
x.detach().cpu().numpy()       # → NumPy 三步流水线
a @ b                          # 矩阵乘:[m,k]@[k,n]→[m,n]
torch.manual_seed(42)          # 固定随机,实验可复现

模型与层

python
class MyNet(nn.Module):
    def __init__(self):
        super().__init__()               # 铁律第一条
        self.net = nn.Sequential(
            nn.Flatten(), nn.Linear(784, 128), nn.ReLU(),
            nn.Dropout(0.2), nn.Linear(128, 10))
    def forward(self, x):
        return self.net(x)

model = MyNet().to(device)               # 先 to(device) 再建优化器
model.train() / model.eval()             # 模式切换,不是"训练/评估"本身
sum(p.numel() for p in model.parameters())   # 参数量

损失与优化器选择

任务loss标签 dtype输出层
回归nn.MSELoss()float321 神经元无激活
二分类nn.BCEWithLogitsLoss()float321 神经元无激活
多分类nn.CrossEntropyLoss()int64n 神经元无激活(logits)
python
opt = torch.optim.AdamW(model.parameters(), lr=1e-3)     # 默认首选
opt = torch.optim.SGD(model.parameters(), lr=0.01, momentum=0.9)
sched = torch.optim.lr_scheduler.StepLR(opt, step_size=3, gamma=0.5)

数据管道

python
datasets.MNIST("./data", train=True, download=True, transform=transform)
DataLoader(ds, batch_size=64, shuffle=True, num_workers=4, pin_memory=True)
random_split(full, [55000, 5000], generator=torch.Generator().manual_seed(42))
# transform 标配:ToTensor() → Normalize(mean, std)
# Windows:num_workers>0 时主逻辑必须包在 if __name__ == "__main__": 里

训练循环(心脏五行)

python
for epoch in range(EPOCHS):
    model.train()
    for x, y in train_loader:
        x, y = x.to(device), y.to(device)
        loss = loss_fn(model(x), y)   # ① 前向 ② 损失
        opt.zero_grad(set_to_none=True)  # ③ 清零
        loss.backward()               # ④ 求梯度
        opt.step()                    # ⑤ 更新
    sched.step()

评估与推理

python
model.eval()
with torch.no_grad():
    logits = model(x)
    pred = logits.argmax(dim=1)                    # 预测类别
    probs = torch.softmax(logits, dim=1)           # 展示概率(仅此处需要)
    acc = (pred == y).float().mean()               # 准确率

保存与加载

python
torch.save(model.state_dict(), "model.pt")                        # 存权重
model.load_state_dict(torch.load("model.pt", map_location=device,
                                 weights_only=True))               # 载权重
torch.save({"model": model.state_dict(), "opt": opt.state_dict(),
            "epoch": e}, "ckpt.pt")                                # 续训存档
# ⚠️ weights_only=False 仅用于自己生成的文件;来路不明的 .pt 不要加载

GPU 与性能

python
torch.cuda.is_available() / torch.cuda.device_count()
torch.cuda.memory_allocated() / torch.cuda.empty_cache()
nvidia-smi                                          # 终端实时监控
with torch.autocast("cuda", dtype=torch.float16): ...   # AMP 混合精度
model = torch.compile(model)                        # 2.x 编译加速

高频报错定位

报错关键词去这里查
mat1 and mat2 shapes cannot be multipliedFAQ Q1
Expected all tensors to be on the same deviceFAQ Q2
loss 不降(不报错)FAQ Q3
CUDA out of memoryFAQ Q4
Can't call numpy() on Tensor that requires gradFAQ Q5
Windows 卡死 / spawn 报错FAQ Q6
Error(s) in loading state_dictFAQ Q7
推理结果每次不同FAQ Q8