Day 39 - 图像数据与显存

在深度学习的进阶之路上，从结构化数据转向图像数据是一个重要的里程碑。图像数据不仅在维度上更加复杂，对计算资源（尤其是显存）的需求也呈指数级增长。

本文将深入探讨图像数据的结构特性、PyTorch 中的模型定义差异，以及训练过程中显存占用的核心机制。

一、 图像数据的多维结构

与结构化表格数据（样本数 × 特征数）不同，图像数据保留了空间结构信息。理解图像的维度顺序是进行深度学习编程的基础。

1.1 灰度图像 vs 彩色图像

• 灰度图像（如 MNIST）
  • 通道数（Channels）：1（仅表示亮度，0为黑，255为白）。
  • 形状 ：(1, Height, Width)。例如 MNIST 为 (1, 28, 28)。
  • 数据类型 ：原始数据通常为 uint8 (0-255)，预处理后转为 float32 (0-1)。
• 彩色图像（如 CIFAR-10）
  • 通道数：3（RGB 红绿蓝）。
  • 形状 ：(3, Height, Width)。例如 CIFAR-10 为 (3, 32, 32)。

1.2 维度顺序的陷阱：Channel First vs Channel Last

这是一个初学者常踩的坑：

• PyTorch 格式 ：[Batch, Channel, Height, Width] (NCHW)。这是 PyTorch 的原生格式。
• NumPy / Matplotlib 格式 ：[Height, Width, Channel] (HWC)。这是 OpenCV 和绘图库的通用格式。

实战技巧：

在使用 matplotlib.pyplot.imshow() 显示 PyTorch 张量图像时，必须先进行维度转换：

# img 是 PyTorch Tensor: [3, 32, 32]
npimg = img.numpy()
# 转换为 [32, 32, 3] 以便显示
plt.imshow(np.transpose(npimg, (1, 2, 0)))

二、 图像分类模型的定义

处理图像数据时，模型定义与普通 MLP 有显著区别，主要体现在输入层的处理上。

2.1 核心差异：展平操作 (Flatten)

全连接层（Linear Layer）只能接受一维向量输入，而图像是三维张量（C, H, W）。因此，在输入第一层全连接层之前，必须将图像"拍扁"。

• nn.Flatten() ：这是 PyTorch 提供的标准层，它会将 [Batch, C, H, W] 展平为 [Batch, C*H*W]，保留 Batch 维度。

2.2 灰度图模型示例 (MNIST)

输入尺寸计算：1 * 28 * 28 = 784。

class MNIST_MLP(nn.Module):
    def __init__(self):
        super().__init__()
        self.flatten = nn.Flatten() # 展平层
        self.layer1 = nn.Linear(784, 128) # 输入 784
        # ... 后续层

2.3 彩色图模型示例 (CIFAR-10)

输入尺寸计算：3 * 32 * 32 = 3072。

class CIFAR_MLP(nn.Module):
    def __init__(self):
        super().__init__()
        self.flatten = nn.Flatten()
        self.layer1 = nn.Linear(3072, 128) # 输入 3072
        # ... 后续层

关键点 ：batch_size 不影响模型定义。无论 Batch 是 1 还是 1024，模型的权重矩阵形状是固定的，PyTorch 会自动广播处理 Batch 维度。

三、 显存占用深度剖析

训练深度学习模型时，"CUDA Out of Memory" (OOM) 是最令人头疼的报错。理解显存去哪了，是优化训练配置的前提。

3.1 显存占用的四大金刚

1. 模型参数 (Parameters) ：
   • 模型的权重（Weights）和偏置（Biases）。
   • 计算：参数数量 × 4 Bytes (float32)。
   • 特点：加载模型后立即占用，与 batch_size 无关。
2. 梯度 (Gradients) ：
   • 反向传播时计算的梯度值，用于更新参数。
   • 计算：通常与参数量相同，即 参数数量 × 4 Bytes。
   • 特点：反向传播开始后占用。
3. 优化器状态 (Optimizer States) ：
   • SGD：无额外状态，最省显存。
   • Adam：需要存储动量（Momentum）和方差（Variance），每个参数对应 2 个额外变量。
   • 计算：参数数量 × 8 Bytes (2 × float32)。
   • 特点：Adam 的显存占用是 SGD 的 3 倍（参数+梯度+2状态 vs 参数+梯度）。
4. 中间激活值 (Intermediate Activations) ：
   • 前向传播时每一层的输出结果，必须保存下来用于反向传播计算梯度。
   • 计算 ：Batch Size × 每层输出形状 × 4 Bytes。
   • 特点 ：显存杀手。它与 layers 数量和 batch_size 成正比。

3.2 显存优化策略

• 调整 Batch Size ：这是最直接的手段。显存不足时，优先减小 Batch Size。
  • 经验公式 ：Max Batch Size ≈ (显存容量 - 固定占用) / 单样本显存消耗。
• 混合精度训练 (AMP)：使用 float16 代替 float32，显存占用减半，计算速度翻倍。
• 梯度累积 (Gradient Accumulation)：如果显存只能跑 batch_size=16，但你想达到 batch_size=64 的效果，可以跑 4 次前向传播再更新一次参数。

3.3 Batch Size 对训练的影响

• 小 Batch Size ：
  • 显存占用低。
  • 梯度噪声大，训练震荡，有助于跳出局部最优，但收敛慢。
  • 无法充分利用 GPU 并行计算能力。
• 大 Batch Size ：
  • 显存占用高。
  • 梯度估计准确，训练稳定。
  • 计算效率高，但可能收敛到尖锐极小值，泛化能力稍弱。

总结：显存管理是一门平衡的艺术，需要在模型深度、Batch Size 和训练速度之间寻找最佳折衷点。