📋 前言

各位伙伴们，大家好！从今天起，我们正式踏入计算机视觉（CV）的奇妙世界。告别了整洁的表格数据，我们迎来了色彩斑斓但结构复杂的图像数据。随之而来的，是所有深度学习从业者都必须面对的灵魂拷问：“我的 GPU 显存又爆了（Out of Memory）！”

Day 39 是理论与工程实践紧密结合的一天。我们将学习图像数据在 PyTorch 中是如何被“表达”的，并深入剖析 GPU 显存这份“预算报告”，看看每一分显存都“花”在了哪里。掌握了今天的知识，你将从一个对 OOM 错误束手无策的新手，成长为一名懂得精打细算、合理分配计算资源的“炼丹工程师”。

---

一、图像数据的“三维”世界：从灰度到彩色

与结构化数据（一个样本就是一个一维向量）不同，图像数据是“立体”的，它包含着丰富的空间信息。

1.1 灰度图像 (以 MNIST 为例)

灰度图只有一个颜色通道，代表像素的明暗程度。

• 形状：在 PyTorch 中，一张 MNIST 图片的形状是 [1, 28, 28]。
  • 1: 通道 (Channel) - 代表这是灰度图。
  • 28: 高度 (Height) - 图像有28个像素高。
  • 28: 宽度 (Width) - 图像有28个像素宽。

这种 (通道, 高, 宽) 的格式被称为 Channel First，是 PyTorch 的标准格式。

1.2 彩色图像 (以 CIFAR-10 为例)

彩色图像通常由红（R）、绿(G)、蓝(B)三个颜色通道叠加而成。

• 形状：一张 CIFAR-10 图片的形状是 [3, 32, 32]。
  • 3: 通道 (Channel) - 分别代表 R, G, B 三个通道。
  • 32: 高度 (Height)。
  • 32: 宽度 (Width)。

⚠️ 注意：可视化时的“陷阱”！

很多可视化库（如 matplotlib）习惯的图像格式是 (高, 宽, 通道) (Channel Last)。因此，当你需要显示一个 PyTorch 张量时，必须进行维度转换：np.transpose(tensor.numpy(), (1, 2, 0))。这是新手最常遇到的问题之一！

---

二、揭秘GPU显存的“四大家族”：钱都花哪儿了？

当模型和数据被加载到 GPU 上时，它们会占用宝贵的显存。显存主要被以下四个部分（我称之为“四大家族”）瓜分：

1. 模型参数与梯度 (Model & Gradients)

   • 是什么：模型自身的权重（weights）和偏置（biases）。在反向传播时，还会产生一份与参数同样大小的梯度（gradients）。
   • 特点：这是“固定资产”，一旦模型确定，这部分开销基本就固定了。

2. 优化器状态 (Optimizer States)

   • 是什么：一些高级优化器（如 Adam）需要为每个参数存储额外的状态信息，比如动量（momentum）和二阶矩估计（variance）。
   • 特点：这是“额外配置”的开销。使用简单的 SGD 优化器，这部分开销几乎为零；而使用 Adam，则会使参数相关的显存占用增加约2倍。

3. 数据批量 (Batch Data)

   • 是什么：每个训练步中，输入到模型的一个批次（batch）的数据。
   • 特点：这是最主要的“流动资金”，也是我们最能直接控制的部分。其大小由 batch_size 决定。batch_size 越大，这部分开销越大。

4. 中间激活值 (Intermediate Activations)

   • 是什么：在前向传播过程中，每一层网络计算出的输出结果。这些结果需要被保存下来，以便在反向传播时计算梯度。
   • 特点：这是“临时工作区”，其大小与 batch_size 和模型深度、宽度正相关。

显存占用部分	特点	控制方式
1. 模型参数与梯度	固定开销，与模型复杂度正相关	简化模型结构
2. 优化器状态	Adam ≈ 2倍参数大小，SGD ≈ 0	更换优化器 (如 SGD)
3. 数据批量	主要可变开销，与 batch_size 正相关	调整 batch_size
4. 中间激活值	可变开销，与 batch_size 和模型深度正相关	调整 batch_size、使用梯度检查点（高级技巧）

---

三、实战核心：batch_size的权衡艺术

batch_size 是我们在训练时最重要的超参数之一，它直接影响显存占用和训练效果。

• 为什么需要 Batch？

  1. 硬件效率：GPU 是为并行计算设计的，一次处理一个大批次的数据远比一次处理一个样本快得多。大的 batch_size 能更好地利用 GPU 的计算能力。
  2. 梯度稳定：一个批次的梯度是多个样本梯度的平均值，这比单个样本的梯度更能代表整体数据的趋势，使得训练过程更稳定，收敛更快。

• 如何选择 batch_size？

  1. 经验法则：对于大显存的 GPU，可以从 32 或 64 开始尝试，并以2的幂次（如 128, 256, 512）逐步增加。
  2. 监控工具：在终端中使用 nvidia-smi 命令实时监控 GPU 显存占用。
  3. 实践策略：不断增大 batch_size，直到出现 CUDA out of memory 错误。然后选择一个比该临界值稍小的值（如80%）作为你的最终 batch_size。

---

四、作业代码：用MLP处理图像数据并分析模型

我们将用一个简单的 MLP 模型来处理 MNIST 和 CIFAR-10 数据，并使用 torchsummary 来验证我们的参数计算。

import torch
import torch.nn as nn
from torchvision import datasets, transforms
from torchsummary import summary

# 检查是否有可用的GPU
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"Using device: {device}")

# --- 1. 针对 MNIST (灰度图) 的 MLP ---
class MLP_MNIST(nn.Module):
    def __init__(self):
        super(MLP_MNIST, self).__init__()
        self.flatten = nn.Flatten()  # 将 [1, 28, 28] 展平为 784
        self.layers = nn.Sequential(
            nn.Linear(28 * 28, 128),
            nn.ReLU(),
            nn.Linear(128, 10)
        )
        
    def forward(self, x):
        x = self.flatten(x)
        return self.layers(x)

# 实例化模型并移动到设备
model_mnist = MLP_MNIST().to(device)

print("\n--- MNIST MLP 模型结构分析 ---")
# 使用 torchsummary 查看模型详情，输入尺寸不包含 batch 维度
summary(model_mnist, input_size=(1, 28, 28))
# 参数计算验证:
# Layer 1: (784 * 128) + 128 = 100,352 + 128 = 100,480
# Layer 2: (128 * 10) + 10 = 1,280 + 10 = 1,290
# Total: 100,480 + 1,290 = 101,770 (与 summary 输出一致！)

# --- 2. 针对 CIFAR-10 (彩色图) 的 MLP ---
class MLP_CIFAR10(nn.Module):
    def __init__(self):
        super(MLP_CIFAR10, self).__init__()
        self.flatten = nn.Flatten() # 将 [3, 32, 32] 展平为 3072
        self.layers = nn.Sequential(
            nn.Linear(3 * 32 * 32, 128),
            nn.ReLU(),
            nn.Linear(128, 10)
        )
        
    def forward(self, x):
        x = self.flatten(x)
        return self.layers(x)

# 实例化模型并移动到设备
model_cifar10 = MLP_CIFAR10().to(device)

print("\n--- CIFAR-10 MLP 模型结构分析 ---")
summary(model_cifar10, input_size=(3, 32, 32))
# 参数计算验证:
# Layer 1: (3072 * 128) + 128 = 393,216 + 128 = 393,344
# Layer 2: (128 * 10) + 10 = 1,280 + 10 = 1,290
# Total: 393,344 + 1,290 = 394,634 (与 summary 输出一致！)

---

五、学习心得

今天的学习是一次从“算法理论”到“工程实践”的重要思维转变。

• 数据皆是张量：我深刻理解了不同类型的数据（结构化、灰度图、彩色图）在 PyTorch 中是如何被统一为“张量”这一数据结构的，关键在于理解其 shape 的含义。
• OOM 不再可怕：Out of Memory 错误不再是一个神秘的黑盒。现在我能清晰地分析出是哪“四大家族”耗尽了显存，并能通过调整 batch_size 这一最有效的手段来解决问题。
• 编程的严谨性：模型定义与 batch_size 无关，但在数据加载时必须指定；torchsummary 的输入尺寸不含 batch 维度；matplotlib 可视化需要转换维度顺序… 这些细节让我体会到深度学习工程的严谨性。

掌握了显存管理，就像掌握了开车时的油门和刹车，让我们在深度学习的道路上行得更稳、更远。

---

再次感谢 @浙大疏锦行 老师，将如此核心且复杂的工程问题讲解得如此透彻！