PyTorch基本用法介绍：从零开始构建深度学习工作流

何伯特2026-02-14 10:26

引言：为什么选择PyTorch？

2017年诞生的PyTorch，如今已是深度学习研究与工业部署的事实标准。它的成功并非偶然------动态计算图、Pythonic的编程范式、完整的生态工具链，让研究人员能够"像写NumPy一样写神经网络"。

本文不追求覆盖全部API，而是构建一条从"张量操作"到"完整训练闭环"的认知路径。你将掌握：

  • PyTorch的核心抽象：张量、自动微分、模块化
  • 标准训练循环的工程化写法
  • 数据加载与预处理的最佳实践
  • 模型保存、部署与调试技巧

一、张量：PyTorch的"原子"

1.1 创建张量的五种姿势

python 复制代码 
import torch
import numpy as np

# 从数据直接创建
t1 = torch.tensor([[1, 2], [3, 4]])

# 从NumPy数组转换（共享内存！）
arr = np.array([[1, 2], [3, 4]])
t2 = torch.from_numpy(arr)

# 特殊初始化
t3 = torch.zeros(3, 4)      # 全零
t4 = torch.ones(2, 3)       # 全一
t5 = torch.randn(3, 5)      # 标准正态分布
t6 = torch.eye(4)           # 单位矩阵
t7 = torch.arange(0, 10, 2) # [0, 2, 4, 6, 8]

# 指定设备和数据类型
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
t8 = torch.tensor([1.0, 2.0], device=device, dtype=torch.float32)

关键认知 ：tensor和ndarray的API高度相似，差异在于设备位置（CPU/GPU）和自动求导能力。

1.2 张量的核心属性

python 复制代码 
x = torch.randn(2, 3, 4)
print(x.shape)      # torch.Size([2, 3, 4])
print(x.dtype)      # torch.float32
print(x.device)     # cpu / cuda:0
print(x.requires_grad) # 是否需要梯度（默认False）

1.3 索引、切片与变形

python 复制代码 
x = torch.randn(4, 5)

# 索引切片（与NumPy完全一致）
y1 = x[1:3, :]      # 第1-2行，所有列
y2 = x[:, -2:]      # 所有行，最后两列
y3 = x[x > 0]       # 条件索引（返回一维）

# 变形
z1 = x.view(2, 10)  # 改变形状（共享内存）
z2 = x.reshape(2, 10) # view的别名，但处理非连续内存更友好
z3 = x.permute(1, 0)  # 转置（维度交换）
z4 = x.flatten()      # 展平为一维

⚠️ 新手陷阱 ：view()要求原始张量在内存中连续，reshape()更鲁棒。不确定时用reshape()。

二、自动微分：PyTorch的"灵魂"

2.1 计算图与梯度

python 复制代码 
# 创建需要梯度的张量
x = torch.tensor([2.0], requires_grad=True)
y = torch.tensor([3.0], requires_grad=True)

# 执行计算
z = x**2 + y**3

# 反向传播
z.backward()

# 查看梯度
print(x.grad)  # dz/dx = 2x = 4.0
print(y.grad)  # dz/dy = 3y^2 = 27.0

可视化理解：

复制代码 
    x(2.0) ──┬── (x^2) ──┐
             │           ↓
             │          (+)── z
             │           ↑
    y(3.0) ──┴── (y^3) ──┘

2.2 梯度清零的艺术

python 复制代码 
# 梯度会累积！
for _ in range(3):
    z = x**2
    z.backward()
print(x.grad)  # 输出: tensor([12.])  每次累加2x=4，3次共12

# 必须手动清零
x.grad.zero_()
z = x**2
z.backward()
print(x.grad)  # tensor([4.])

工程实践 ：优化器提供了zero_grad()方法，推荐使用optimizer.zero_grad(set_to_none=True)获得更好性能。

2.3 关闭梯度追踪的三种场景

python 复制代码 
# 场景1：模型推理
with torch.no_grad():
    predictions = model(test_data)

# 场景2：张量操作后不需要梯度
x = torch.randn(3, requires_grad=True)
y = x.detach()  # 返回新的张量，切断梯度追踪

# 场景3：冻结预训练模型参数
for param in model.parameters():
    param.requires_grad = False

三、模块化：从零构建神经网络

3.1 nn.Module：所有模型的基类

python 复制代码 
import torch.nn as nn
import torch.nn.functional as F

class TwoLayerNet(nn.Module):
    def __init__(self, input_dim, hidden_dim, output_dim):
        super().__init__()
        # 定义可学习参数层
        self.fc1 = nn.Linear(input_dim, hidden_dim)
        self.fc2 = nn.Linear(hidden_dim, output_dim)
        # 初始化（可选）
        nn.init.xavier_uniform_(self.fc1.weight)
        nn.init.zeros_(self.fc1.bias)
    
    def forward(self, x):
        # 定义前向传播逻辑
        x = self.fc1(x)
        x = F.relu(x)     # 无参数层使用functional
        x = self.fc2(x)
        return x

设计哲学 ：__init__注册参数，forward定义计算逻辑。不要手动调用forward()，应使用model(x)。

3.2 参数管理与设备迁移

python 复制代码 
model = TwoLayerNet(784, 256, 10)

# 查看所有参数
for name, param in model.named_parameters():
    print(f"{name}: {param.shape}, requires_grad={param.requires_grad}")

# 模型移到GPU
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = model.to(device)

# 数据也必须同步迁移
data = data.to(device)

3.3 常用模块速查表

层类型 PyTorch实现 说明
全连接 nn.Linear(in, out) 基础变换层
卷积 nn.Conv2d(in, out, kernel_size) 图像特征提取
池化 nn.MaxPool2d(kernel_size) 降维/平移不变性
RNN nn.LSTM(input_size, hidden_size) 序列建模
批归一化 nn.BatchNorm1d/2d(num_features) 加速收敛、稳定训练
Dropout nn.Dropout(p=0.5) 防止过拟合
Embedding nn.Embedding(num_embeddings, embedding_dim) 词/类别嵌入

四、训练循环：从脚本到工程

4.1 标准训练范式（模板级）

python 复制代码 
def train_one_epoch(model, dataloader, optimizer, criterion, device):
    model.train()  # 启用训练模式
    running_loss = 0.0
    correct = 0
    total = 0
    
    for inputs, targets in dataloader:
        inputs, targets = inputs.to(device), targets.to(device)
        
        # 1. 清零梯度
        optimizer.zero_grad(set_to_none=True)
        
        # 2. 前向传播
        outputs = model(inputs)
        loss = criterion(outputs, targets)
        
        # 3. 反向传播
        loss.backward()
        
        # 4. 参数更新
        optimizer.step()
        
        # 5. 统计指标
        running_loss += loss.item()
        _, predicted = outputs.max(1)
        total += targets.size(0)
        correct += predicted.eq(targets).sum().item()
    
    return running_loss / len(dataloader), 100. * correct / total

def evaluate(model, dataloader, criterion, device):
    model.eval()  # 启用评估模式
    running_loss = 0.0
    correct = 0
    total = 0
    
    with torch.no_grad():  # 关闭梯度追踪
        for inputs, targets in dataloader:
            inputs, targets = inputs.to(device), targets.to(device)
            outputs = model(inputs)
            loss = criterion(outputs, targets)
            
            running_loss += loss.item()
            _, predicted = outputs.max(1)
            total += targets.size(0)
            correct += predicted.eq(targets).sum().item()
    
    return running_loss / len(dataloader), 100. * correct / total

4.2 完整训练脚本骨架

python 复制代码 
def main():
    # 1. 超参数
    config = {
        'batch_size': 64,
        'lr': 0.001,
        'epochs': 10,
        'device': 'cuda' if torch.cuda.is_available() else 'cpu'
    }
    
    # 2. 数据准备
    train_loader, val_loader = get_data_loaders(batch_size=config['batch_size'])
    
    # 3. 模型、损失函数、优化器
    model = TwoLayerNet(784, 256, 10).to(config['device'])
    criterion = nn.CrossEntropyLoss()
    optimizer = torch.optim.Adam(model.parameters(), lr=config['lr'])
    scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=5, gamma=0.1)
    
    # 4. 训练循环
    for epoch in range(config['epochs']):
        train_loss, train_acc = train_one_epoch(
            model, train_loader, optimizer, criterion, config['device']
        )
        val_loss, val_acc = evaluate(
            model, val_loader, criterion, config['device']
        )
        scheduler.step()  # 调整学习率
        
        print(f"Epoch {epoch+1}: "
              f"Train Loss: {train_loss:.4f}, Acc: {train_acc:.2f}% | "
              f"Val Loss: {val_loss:.4f}, Acc: {val_acc:.2f}%")
    
    # 5. 保存模型
    torch.save({
        'model_state_dict': model.state_dict(),
        'optimizer_state_dict': optimizer.state_dict(),
        'epoch': epoch,
        'config': config
    }, 'checkpoint.pth')

五、数据加载：被低估的性能瓶颈

5.1 Dataset与DataLoader解耦

python 复制代码 
from torch.utils.data import Dataset, DataLoader

class CustomDataset(Dataset):
    def __init__(self, images, labels, transform=None):
        self.images = images
        self.labels = labels
        self.transform = transform
    
    def __len__(self):
        return len(self.images)
    
    def __getitem__(self, idx):
        image = self.images[idx]
        label = self.labels[idx]
        
        if self.transform:
            image = self.transform(image)
        
        return image, label

5.2 DataLoader黄金参数

python 复制代码 
dataloader = DataLoader(
    dataset,
    batch_size=32,
    shuffle=True,              # 训练集必须打乱
    num_workers=4,            # 多进程加载（瓶颈定位）
    pin_memory=True,          # GPU训练时启用（锁页内存）
    persistent_workers=True,  # epoch间不销毁进程
    drop_last=True           # 丢弃不完整批次
)

性能诊断 ：若GPU利用率低于80%，优先增加num_workers和pin_memory。

5.3 torchvision数据增强

python 复制代码 
from torchvision import transforms

train_transform = transforms.Compose([
    transforms.RandomResizedCrop(224),
    transforms.RandomHorizontalFlip(),
    transforms.ColorJitter(brightness=0.2, contrast=0.2),
    transforms.ToTensor(),  # 必须！将PIL图像转为张量
    transforms.Normalize(mean=[0.485, 0.456, 0.406], 
                        std=[0.229, 0.224, 0.225])
])

六、模型保存与加载：不止于torch.save

6.1 推荐的保存方式（完整检查点）

python 复制代码 
# 保存
checkpoint = {
    'epoch': epoch,
    'model_state_dict': model.state_dict(),
    'optimizer_state_dict': optimizer.state_dict(),
    'loss': loss,
    'config': config
}
torch.save(checkpoint, 'model_checkpoint.pth')

# 加载
checkpoint = torch.load('model_checkpoint.pth')
model.load_state_dict(checkpoint['model_state_dict'])
optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
epoch = checkpoint['epoch']

6.2 推理部署的保存方式

python 复制代码 
# 仅保存模型参数（推荐推理）
torch.save(model.state_dict(), 'model_weights.pth')

# 保存完整模型（包含结构，不推荐跨版本）
torch.save(model, 'full_model.pth')

# 加载权重
model = TwoLayerNet(784, 256, 10)
model.load_state_dict(torch.load('model_weights.pth'))
model.eval()  # 重要！

七、调试技巧：告别玄学调参

7.1 过拟合单批次------快速验证

python 复制代码 
def verify_model(model, dataloader, criterion, optimizer, device):
    """在小数据集上验证模型能否收敛"""
    sample_batch = next(iter(dataloader))
    inputs, targets = [x.to(device) for x in sample_batch]
    
    for _ in range(50):  # 反复拟合同一批数据
        optimizer.zero_grad()
        outputs = model(inputs)
        loss = criterion(outputs, targets)
        loss.backward()
        optimizer.step()
    
    print(f"Final loss on single batch: {loss.item():.6f}")
    # 正常情况应接近0

7.2 梯度检查与NaN追踪

python 复制代码 
# 开启异常检测
torch.autograd.set_detect_anomaly(True)

# 查看梯度分布
for name, param in model.named_parameters():
    if param.grad is not None:
        print(f"{name}: grad_mean={param.grad.mean():.3e}, "
              f"grad_std={param.grad.std():.3e}")

7.3 TensorBoard可视化

python 复制代码 
from torch.utils.tensorboard import SummaryWriter

writer = SummaryWriter('runs/experiment_1')

# 记录标量
writer.add_scalar('Loss/train', loss, epoch)
writer.add_scalar('Accuracy/val', acc, epoch)

# 记录模型计算图
writer.add_graph(model, inputs)

# 记录直方图（参数分布）
writer.add_histogram('fc1/weights', model.fc1.weight, epoch)

八、PyTorch 2.x：编译时代

8.1 torch.compile入门

python 复制代码 
model = torch.compile(model, 
                      dynamic=False,      # 是否动态形状
                      fullgraph=False,   # 开发时设为True捕获图断裂
                      mode="reduce-overhead")  # 或"max-autotune"

# 其余代码完全不变！

适用场景 ：GPU密集型运算、固定输入形状。图断裂 （如if x.sum() > 0）会限制优化效果。

8.2 编译后验证

python 复制代码 
# 验证编译是否生效
print(model)  # 显示OptimizedModule说明成功

九、从"跑通"到"跑赢"：实用者清单

✅ 基础三要素

  • 指定随机种子（torch.manual_seed(42)）
  • 启用pin_memory=True（GPU训练）
  • 使用set_to_none=True清梯度

✅ 训练稳定化

  • 输入归一化（零均值单位方差）
  • 学习率预热与衰减
  • 梯度裁剪（torch.nn.utils.clip_grad_norm_）

✅ 性能榨取

  • 尝试torch.compile
  • 启用自动混合精度（AMP）
  • num_workers调优（CPU核数/2）

✅ 可复现性

  • 固定CUDNN确定性算法
python 复制代码 
torch.backends.cudnn.deterministic = True
torch.backends.cudnn.benchmark = False

结语：PyTorch学习的三重境界

第一重：把PyTorch当NumPy用，拼出别人的模型。

第二重：理解自动微分原理，能够Debug梯度消失与爆炸。

第三重：形成自己的工程范式------知道何时用DDP、何时用FSDP、何时用ONNX导出，并在代码层面固化这些最佳实践。

本文的目标是帮助你跨越第一重，窥见第二重。真正的掌握，始于你开始质疑官方教程、修改源码、为开源项目贡献PR的那一刻。

记住：PyTorch只是工具，你对深度学习问题的认知深度，才是真正的护城河。

附录：必读资源

相关推荐
宁远x1 小时前
【万字长文】PyTorch FSDP 设计解读与性能分析
人工智能·pytorch·深度学习·云计算
赛博鲁迅1 小时前
dify添加中转站模型教程
人工智能·gpt·aigc·ai编程·dify·ai-native
还在忙碌的吴小二1 小时前
飞牛NAS ARM版升级全指南+性能深度解析｜低成本盘活闲置设备，功耗与体验双突破
arm开发·人工智能
what丶k1 小时前
AI 中的向量详解：从原理到实战
人工智能
hans汉斯1 小时前
基于联邦学习的隐私保护和抗投毒攻击方法研究
网络·人工智能·算法·yolo·数据挖掘·聚类·汉斯出版社
MicrosoftReactor1 小时前
技术速递｜社区驱动的 AI 安全：一个面向安全研究的开源框架
人工智能·安全·开源
yong99901 小时前
基于SIFT的MATLAB图像拼接实现
人工智能·计算机视觉·matlab
雨大王5121 小时前
广域铭岛构建全链路智能制造新引擎
人工智能
ActionTech2 小时前
数据集推荐 06 | 首款 NL2GeoSQL 的测试基准和数据集来了！
数据库·人工智能·sql
热门推荐
01GitHub 镜像站点02Claude Code + GLM4.7 避坑指南：解决 Unable to connect to Anthropic services03爬虫逆向之观安（观镜WEB应用安全防护系统）04openClaw安装飞书插件｜核心踩坑：spawn EINVAL 错误终极解决指南05BongoCat - 跨平台键盘猫动画工具06HTB 赛季10 - Pterodactyl - user07UV安装并设置国内源08243 行 microGPT：把“训练 + 推理”拆到骨头里09openclaw配置教程（linux+局域网ollama）10openclaw 配置飞书 报错 应用未建立长连接