3. PyTorch基础回顾：张量操作、自动求导、nn.Module构建模型

说实话，很多做工业缺陷检测的朋友，一上来就急着调模型、跑训练。结果呢？模型结构写错了都不知道，梯度传不下去也查不出原因。我个人习惯是，先把PyTorch这三块基本功打扎实——张量、自动求导、模型构建。这三样搞明白了，后面转ONNX、做推理，你心里才有底。

3.1 张量操作：一切数据的起点

张量（Tensor）就是PyTorch里的“数组”，但它比NumPy的ndarray多了GPU加速和自动求导的能力。我在做缺陷检测项目时，经常需要把工业相机采集的原始图像数据转成张量，再做归一化、裁剪、拼接。

3.1.1 创建张量的几种方式

方法 说明 常用场景
torch.tensor() 从列表或NumPy数组创建 加载单张图像数据
torch.zeros() 创建全零张量 初始化特征图或掩码
torch.ones() 创建全一张量 偏置初始化
torch.randn() 标准正态分布随机张量 权重初始化
torch.arange() 创建等差序列 生成标签索引 
import torch

# 从列表创建
data = torch.tensor([[1, 2], [3, 4]], dtype=torch.float32)
print(data.shape)  # torch.Size([2, 2])

# 创建全零张量，指定形状和设备
mask = torch.zeros((1, 256, 256), device='cuda' if torch.cuda.is_available() else 'cpu')
我的习惯：创建张量时尽量显式指定 dtype。工业模型推理时，ONNX默认用float32，你训练时如果用float64，转ONNX会出警告。我踩过这个坑，后来一律用 torch.float32。

3.1.2 张量的形状操作

缺陷检测模型里，特征图的形状变换特别频繁。比如从卷积层出来的特征图是 [B, C, H, W]，要送入全连接层就得先展平。

x = torch.randn(4, 3, 224, 224)  # 模拟一个batch的4张图像

# 展平：从第1维开始展平
x_flat = x.view(4, -1)           # 形状变为 [4, 3*224*224]
print(x_flat.shape)              # torch.Size([4, 150528])

# 转置：交换维度
x_transposed = x.transpose(1, 3) # 形状变为 [4, 224, 224, 3]

# 增加/减少维度
x_unsqueeze = x.unsqueeze(0)     # 在batch维前加一维，[1, 4, 3, 224, 224]
x_squeeze = x_unsqueeze.squeeze(0) # 去掉第0维，回到 [4, 3, 224, 224]
注意：view() 要求张量在内存中是连续的。如果你做了转置或切片操作，记得先调 .contiguous() 再 view()。我曾经在写自定义网络时，因为没加 contiguous()，模型跑着跑着就报错，查了半天才发现是形状不连续的问题。

3.1.3 张量的数学运算

这些跟NumPy几乎一样，但多了自动求导的追踪能力。

a = torch.tensor([1.0, 2.0, 3.0], requires_grad=True)
b = torch.tensor([4.0, 5.0, 6.0])

c = a + b        # 加法
d = a * b        # 逐元素乘
e = torch.matmul(a, b)  # 点积，结果是 1*4 + 2*5 + 3*6 = 32.0
f = torch.mean(a)       # 求均值

3.2 自动求导：梯度从哪来

自动求导（Autograd）是PyTorch训练模型的基石。说白了，你定义好前向计算图，PyTorch自动帮你算梯度。你想想看，如果没有这个机制，手写反向传播得有多痛苦？

3.2.1 计算图与梯度

每次你对一个张量做运算，PyTorch就会在后台构建一个计算图。这个图记录了数据流向和运算关系。当你调用 .backward() 时，它沿着图反向传播，计算每个节点的梯度。

x = torch.tensor(2.0, requires_grad=True)
y = x ** 2 + 3 * x + 1   # y = x^2 + 3x + 1
y.backward()              # 反向传播
print(x.grad)             # dy/dx = 2x + 3 = 2*2 + 3 = 7.0
关键点：只有设置了 requires_grad=True 的张量才会被追踪。模型参数默认就是 requires_grad=True，而输入数据通常不需要梯度。

3.2.2 梯度清零与detach

训练时，每次反向传播后梯度会累积。所以每个batch都要手动清零。另外，有时候你不想让某个操作影响梯度，就用 .detach() 把它从计算图中剥离。

# 梯度清零
optimizer.zero_grad()   # 或者手动：model.zero_grad()

# detach示例：比如你想看中间特征图的值，但不希望它影响梯度
features = model(x)
features_detached = features.detach()  # 从计算图中剥离
print(features_detached.numpy())       # 可以转成numpy，不影响原图
避坑指南：我曾经在写一个多任务损失时，忘了对某个分支的梯度做detach，结果两个任务的梯度互相干扰，模型怎么训都不收敛。后来加上detach，问题就解决了。

3.3 nn.Module构建模型：搭积木的艺术

PyTorch用 nn.Module 来组织模型。你只需要继承这个类，在 __init__ 里定义层，在 forward 里定义前向传播逻辑。我习惯把模型拆成几个子模块，每个子模块负责一个功能，这样调试起来特别方便。

3.3.1 一个简单的缺陷检测分类模型

import torch.nn as nn
import torch.nn.functional as F

class DefectClassifier(nn.Module):
    def __init__(self, num_classes=5):
        super(DefectClassifier, self).__init__()
        # 卷积层：输入3通道，输出16通道，卷积核3x3
        self.conv1 = nn.Conv2d(3, 16, kernel_size=3, padding=1)
        self.bn1 = nn.BatchNorm2d(16)
        self.pool = nn.MaxPool2d(2, 2)

        self.conv2 = nn.Conv2d(16, 32, kernel_size=3, padding=1)
        self.bn2 = nn.BatchNorm2d(32)

        # 全连接层
        self.fc = nn.Linear(32 * 56 * 56, num_classes)  # 假设输入224x224，两次池化后56x56

    def forward(self, x):
        x = self.pool(F.relu(self.bn1(self.conv1(x))))  # [B, 16, 112, 112]
        x = self.pool(F.relu(self.bn2(self.conv2(x))))  # [B, 32, 56, 56]
        x = x.view(x.size(0), -1)                       # 展平
        x = self.fc(x)                                  # [B, num_classes]
        return x

3.3.2 模型参数管理与设备迁移

模型建好后，你需要管理它的参数，并把它放到GPU上。这个操作在工业部署时特别重要——你总不能在CPU上跑推理吧？

model = DefectClassifier(num_classes=5)

# 查看模型参数
for name, param in model.named_parameters():
    print(f"{name}: {param.shape}")

# 迁移到GPU
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model = model.to(device)

# 训练时，输入数据也要迁移
inputs = torch.randn(4, 3, 224, 224).to(device)
labels = torch.randint(0, 5, (4,)).to(device)
注意：模型和数据必须在同一个设备上。我见过有人把模型放GPU，数据忘在CPU，结果报错 Expected all tensors to be on the same device。这个错误很常见，排查时先检查设备一致性。

3.3.3 训练循环模板

下面是我常用的训练循环模板，简洁但完整。工业项目里，你只需要在这个框架里加数据加载、验证逻辑和模型保存就行。

import torch.optim as optim

model = DefectClassifier(num_classes=5).to(device)
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)

for epoch in range(10):
    model.train()
    running_loss = 0.0
    for inputs, labels in train_loader:
        inputs, labels = inputs.to(device), labels.to(device)

        optimizer.zero_grad()
        outputs = model(inputs)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()

        running_loss += loss.item()

    print(f"Epoch {epoch+1}, Loss: {running_loss/len(train_loader):.4f}")

3.4 本章知识结构图

下面这张图帮你理清PyTorch三块核心内容的关系。记住：张量是数据载体，自动求导是训练引擎，nn.Module是模型骨架。三者缺一不可。

PyTorch基础三件套 张量操作 自动求导 nn.Module 创建 形状变换 数学运算 计算图 backward() 梯度清零 定义层 forward() 参数管理 核心关系 张量是数据载体 → 自动求导提供梯度 → nn.Module组织模型结构 三者结合，才能完成训练 → 导出ONNX → 工业部署

3.5 小结

这一章我们回顾了PyTorch最核心的三个基础：张量操作、自动求导和nn.Module模型构建。说白了，这些就是你在工业缺陷检测项目里每天都要打交道的东西。我个人建议，别急着往下赶，先把这几个代码示例自己敲一遍，跑通。后面转ONNX的时候，你会感谢现在打下的基础。

← 上一章 📖 返回目录 下一章 →