Day 40 图像数据与显存

一、图像数据的介绍

1. 灰度图像

import torch import torch.nn as nn import torch.optim as optim from torch.utils.data import DataLoader , Dataset # DataLoader 是 PyTorch 中用于加载数据的工具 from torchvision import datasets, transforms # torchvision 是一个用于计算机视觉的库，datasets 和 transforms 是其中的模块 import matplotlib.pyplot as plt # 设置随机种子，确保结果可复现 torch.manual_seed(42) # 1. 数据预处理，该写法非常类似于管道pipeline # transforms 模块提供了一系列常用的图像预处理操作 # 先归一化，再标准化 transform = transforms.Compose([ transforms.ToTensor(), # 转换为张量并归一化到[0,1] transforms.Normalize((0.1307,), (0.3081,)) # MNIST数据集的均值和标准差，这个值很出名，所以直接使用 ]) import matplotlib.pyplot as plt # 2. 加载MNIST数据集，如果没有会自动下载 train_dataset = datasets.MNIST( root='./data', train=True, download=True, transform=transform ) test_dataset = datasets.MNIST( root='./data', train=False, transform=transform ) # 随机选择一张图片，可以重复运行，每次都会随机选择 sample_idx = torch.randint(0, len(train_dataset), size=(1,)).item() # 随机选择一张图片的索引 # len(train_dataset) 表示训练集的图片数量；size=(1,)表示返回一个索引；torch.randint() 函数用于生成一个指定范围内的随机数,item() 方法将张量转换为 Python 数字 image, label = train_dataset[sample_idx] # 获取图片和标签 # 可视化原始图像（需要反归一化） def imshow(img): img = img * 0.3081 + 0.1307 # 反标准化 npimg = img.numpy() plt.imshow(npimg[0], cmap='gray') # 显示灰度图像 plt.show() print(f"Label: {label}") imshow(image)

图像数据与结构化数据的差异

结构化数据（如表格）的形状通常是 (样本数, 特征数)，例如 (1000, 5) 表示 1000 个样本，每个样本有 5 个特征。图像数据的形状更复杂，需要保留空间信息（高度、宽度、通道），因此不能直接用一维向量表示。其中颜色信息往往是最开始输入数据的通道的含义，因为每个颜色可以用红绿蓝三原色表示，因此一般输入数据的通道数是 3。

MNIST 数据集是手写数字的 灰度图像，每个像素点的取值范围为 0-255（黑白程度），因此 通道数为 1。图像尺寸统一为 28×28 像素。

2. 彩色图像

在 PyTorch 中，图像数据的形状通常遵循 (通道数, 高度, 宽度) 的格式（即 Channel First 格式），这与常见的 (高度, 宽度, 通道数)（Channel Last，如 NumPy 数组）不同。---注意顺序关系，

注意点：

1. 如果用matplotlib库来画图，需要转换下顺序，我们后续介绍

2. 模型输入通常需要 批次维度（Batch Size），形状变为 (批次大小, 通道数, 高度, 宽度)。例如，批量输入 10 张 MNIST 图像时，形状为 (10, 1, 28, 28)。

# 打印一张彩色图像，用cifar-10数据集 import torch import torchvision import torchvision.transforms as transforms import matplotlib.pyplot as plt import numpy as np # 设置随机种子确保结果可复现 torch.manual_seed(42) # 定义数据预处理步骤 transform = transforms.Compose([ transforms.ToTensor(), # 转换为张量并归一化到[0,1] transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)) # 标准化处理 ]) # 2. 加载CIFAR-10数据集（修改root为压缩包所在文件夹） trainset = datasets.CIFAR10( root=r"D:\PythonStudy", # 压缩包所在的文件夹路径（关键！不是压缩包本身） train=True, download=True, # 检测到文件夹内有压缩包时，仅解压不下载 transform=transform ) test_dataset = datasets.CIFAR10( root=r"D:\PythonStudy", # 同训练集的root路径 train=False, transform=transform ) # 创建数据加载器 trainloader = torch.utils.data.DataLoader( trainset, batch_size=4, shuffle=True ) # CIFAR-10的10个类别 classes = ('plane', 'car', 'bird', 'cat', 'deer', 'dog', 'frog', 'horse', 'ship', 'truck') # 随机选择一张图片 sample_idx = torch.randint(0, len(trainset), size=(1,)).item() image, label = trainset[sample_idx] # 打印图片形状 print(f"图像形状: {image.shape}") # 输出: torch.Size([3, 32, 32]) print(f"图像类别: {classes[label]}") # 定义图像显示函数（适用于CIFAR-10彩色图像） def imshow(img): img = img / 2 + 0.5 # 反标准化处理，将图像范围从[-1,1]转回[0,1] npimg = img.numpy() plt.imshow(np.transpose(npimg, (1, 2, 0))) # 调整维度顺序：(通道,高,宽) → (高,宽,通道) plt.axis('off') # 关闭坐标轴显示 plt.show() # 显示图像 imshow(image)

二、图像相关的神经网络定义

1. 黑白图像模型的定义

# 先归一化，再标准化 transform = transforms.Compose([ transforms.ToTensor(), # 转换为张量并归一化到[0,1] transforms.Normalize((0.1307,), (0.3081,)) # MNIST数据集的均值和标准差，这个值很出名，所以直接使用 ]) import matplotlib.pyplot as plt # 2. 加载MNIST数据集，如果没有会自动下载 train_dataset = datasets.MNIST( root='./data', train=True, download=True, transform=transform ) test_dataset = datasets.MNIST( root='./data', train=False, transform=transform ) # 定义两层MLP神经网络 class MLP(nn.Module): def __init__(self): super(MLP, self).__init__() self.flatten = nn.Flatten() # 将28x28的图像展平为784维向量 self.layer1 = nn.Linear(784, 128) # 第一层：784个输入，128个神经元 self.relu = nn.ReLU() # 激活函数 self.layer2 = nn.Linear(128, 10) # 第二层：128个输入，10个输出（对应10个数字类别） def forward(self, x): x = self.flatten(x) # 展平图像 x = self.layer1(x) # 第一层线性变换 x = self.relu(x) # 应用ReLU激活函数 x = self.layer2(x) # 第二层线性变换，输出logits return x # 初始化模型 model = MLP() device = torch.device("cuda" if torch.cuda.is_available() else "cpu") model = model.to(device) # 将模型移至GPU（如果可用） from torchsummary import summary # 导入torchsummary库 print("\n模型结构信息：") summary(model, input_size=(1, 28, 28)) # 输入尺寸为MNIST图像尺寸

模型结构信息： ---------------------------------------------------------------- Layer (type) Output Shape Param # ================================================================ Flatten-1 [-1, 784] 0 Linear-2 [-1, 128] 100,480 ReLU-3 [-1, 128] 0 Linear-4 [-1, 10] 1,290 ================================================================ Total params: 101,770 Trainable params: 101,770 Non-trainable params: 0 ---------------------------------------------------------------- Input size (MB): 0.00 Forward/backward pass size (MB): 0.01 Params size (MB): 0.39 Estimated Total Size (MB): 0.40 ----------------------------------------------------------------

我们关注和之前结构化 MLP 的差异

1. 输入需要展平操作MLP 的输入层要求输入是一维向量，但 MNIST 图像是二维结构（28×28 像素），形状为 [1,28,28]（通道 × 高 × 宽）。nn.Flatten () 展平操作将二维图像 “拉成” 一维向量（784=28×28 个元素），使其符合全连接层的输入格式。其中不定义这个 flatten 方法，直接在前向传播的过程中用 x = x.view (-1, 28 * 28) 将图像展平为一维向量也可以实现
2. 输入数据的尺寸包含了通道数 input_size=(1,28,28)
3. 参数的计算

• 第一层 layer1（全连接层）

权重参数：输入维度 × 输出维度 = 784 × 128 = 100,352

偏置参数：输出维度 = 128合计：100,352 + 128 = 100,480

• 第二层 layer2（全连接层）

权重参数：输入维度 × 输出维度 = 128 × 10 = 1,280

偏置参数：输出维度 = 10 合计：1,280 + 10 = 1,290

• 总参数：100,480（layer1） + 1,290（layer2） = 101,770

2. 彩色图像模型的定义

class MLP(nn.Module): def __init__(self, input_size=3072, hidden_size=128, num_classes=10): super(MLP, self).__init__() # 展平层：将3×32×32的彩色图像转为一维向量 # 输入尺寸计算：3通道 × 32高 × 32宽 = 3072 self.flatten = nn.Flatten() # 全连接层 self.fc1 = nn.Linear(input_size, hidden_size) # 第一层 self.relu = nn.ReLU() self.fc2 = nn.Linear(hidden_size, num_classes) # 输出层 def forward(self, x): x = self.flatten(x) # 展平：[batch, 3, 32, 32] → [batch, 3072] x = self.fc1(x) # 线性变换：[batch, 3072] → [batch, 128] x = self.relu(x) # 激活函数 x = self.fc2(x) # 输出层：[batch, 128] → [batch, 10] return x # 初始化模型 model = MLP() device = torch.device("cuda" if torch.cuda.is_available() else "cpu") model = model.to(device) # 将模型移至GPU（如果可用） from torchsummary import summary # 导入torchsummary库 print("\n模型结构信息：") summary(model, input_size=(3, 32, 32)) # CIFAR-10 彩色图像（3×32×32）

模型结构信息： ---------------------------------------------------------------- Layer (type) Output Shape Param # ================================================================ Flatten-1 [-1, 3072] 0 Linear-2 [-1, 128] 393,344 ReLU-3 [-1, 128] 0 Linear-4 [-1, 10] 1,290 ================================================================ Total params: 394,634 Trainable params: 394,634 Non-trainable params: 0 ---------------------------------------------------------------- Input size (MB): 0.01 Forward/backward pass size (MB): 0.03 Params size (MB): 1.51 Estimated Total Size (MB): 1.54 ----------------------------------------------------------------

• 第一层 layer1（全连接层）

权重参数：输入维度 × 输出维度 = 3072 × 128 = 393,216

偏置参数：输出维度 = 128合计：393,216 + 128 = 393,344

• 第二层 layer2（全连接层）

权重参数：输入维度 × 输出维度 = 128 × 10 = 1,280

偏置参数：输出维度 = 10合计：1,280 + 10 = 1,290

• 总参数：393,344（layer1） + 1,290（layer2） = 394,634

3. 模型定义与batchsize的关系

实际定义中，输入图像还存在batchsize这一维度

在 PyTorch 中，模型定义和输入尺寸的指定不依赖于 batch_size，无论设置多大的 batch_size，模型结构和输入尺寸的写法都是不变的。

class MLP(nn.Module): def __init__(self): super().__init__() self.flatten = nn.Flatten() # nn.Flatten()会将每个样本的图像展平为 784 维向量，但保留 batch 维度。 self.layer1 = nn.Linear(784, 128) self.relu = nn.ReLU() self.layer2 = nn.Linear(128, 10) def forward(self, x): x = self.flatten(x) # 输入：[batch_size, 1, 28, 28] → [batch_size, 784] x = self.layer1(x) # [batch_size, 784] → [batch_size, 128] x = self.relu(x) x = self.layer2(x) # [batch_size, 128] → [batch_size, 10] return x

三、显存占用的主要组成成分

在面对数据集过大的情况下，由于无法一次性将数据全部加入到显存中，所以采取了分批次加载这种方式。即一次只加载一部分数据，保证在显存的范围内。

那么显存设置多少合适呢？如果设置的太小，那么每个batchsize的训练不足以发挥显卡的能力，浪费计算资源；如果设置的太大，会出现OOT（out of memory）

显存一般被以下内容占用：

1. 模型参数与梯度：模型的权重（Parameters）和对应的梯度（Gradients）会占用显存，尤其是深度神经网络（如 Transformer、ResNet 等），一个 1 亿参数的模型（如 BERT-base），单精度（float32）参数占用约 400MB（1e8×4Byte），加上梯度则翻倍至 800MB（每个权重参数都有其对应的梯度）。

2. 部分优化器（如 Adam）会为每个参数存储动量（Momentum）和平方梯度（Square Gradient），进一步增加显存占用（通常为参数大小的 2-3 倍）

3. 其他开销。

OOM处理方法：

1.减小Batch Size

2.清理GPU缓存，释放内存

3.推理阶段关闭梯度计算（节省显存）

4.使用混合精度训练

5.梯度积累

6.控制图像输入尺寸

7.使用轻量级模型结构

以手写数据集为例

手写数据集（MNIST）和当前 MLP 模型，显存占用的计算可以简化为以下几个部分。

1. 模型参数与梯度（FP32 精度）

参数数量：101,770 个参数

• 1 字节（Byte）= 8 位（bit），是计算机存储的最小寻址单位。
• 位（bit）是二进制数的最小单位（0 或 1），例如 0b1010 表示 4 位二进制数。
• 1KB=1024 字节；1MB=1024KB=1,048,576 字节

常数数据类型的字节占用

• MNIST 数据集的原始图像像素值为 0-255 的整数（uint8 类型，占 1 字节），表示灰度值（0 = 黑，255 = 白）。
• 但 PyTorch 的 transforms.ToTensor () 会将其归一化到 [0,1] 范围，并转换为 float32 类型（浮点型更适合神经网络计算）。

计算示例：单张 MNIST 图像的显存占用

1. 原始像素值（uint8，未转换时）
   • 尺寸：28×28 像素
   • 单像素占用：1 字节（uint8）
   • 总占用：28×28×1 = 784 字节 ≈ 0.766 KB
2. 转换为 float32 张量后
   • 尺寸：1×28×28（通道 × 高 × 宽）
   • 单像素占用：4 字节（float32）
   • 总占用：1×28×28×4 = 3136 字节 ≈ 3.06 KB

单精度（float32）参数占用：101,770 × 4 Byte ≈ 403 KB

梯度是损失函数对模型参数的导数（dLoss/dWeight），用于指示参数更新的方向和幅度。梯度是损失函数对模型参数的导数（dLoss/dWeight），用于指示参数更新的方向和幅度。因此在默认情况下，梯度的数据类型和数目与参数相同。

梯度占用（反向传播时）：与参数相同，合计约 806 KB

2. 优化器状态

SGD

• SGD 优化器不存储额外动量，因此无额外显存占用。
• SGD 随机梯度下降，最基础的优化器，直接沿梯度反方向更新参数。
• 参数更新公式：w = w - learning_rate * gradient

Adam

• Adam 优化器：自适应学习率优化器，结合了动量（Momentum）和梯度平方的指数移动平均。
• 每个参数存储动量（m）和平方梯度（v），占用约 101,770 × 8 Byte = 806 KB
• 动量（m）：每个参数对应一个动量值，数据类型与参数相同（float32），占用 403 KB。
• 梯度平方（v）：每个参数对应一个梯度平方值，数据类型与参数相同（float32），占用 403 KB。

3. 数据批量（batch_size）的显存占用

• 单张图像尺寸：1×28×28（通道 × 高 × 宽），归一化转换为张量后为 float32 类型
  • 单张图像显存占用：1×28×28×4 Byte = 3,136 Byte ≈ 3 KB
• 批量数据占用：batch_size × 单张图像占用
  • 例如：batch_size=64 时，数据占用为 64×3 KB = 192 KB
  • batch_size=1024 时，数据占用为 1024×3 KB = 3 MB

4. 前向 / 反向传播中间变量

• 对于两层 MLP，中间变量（如 Layer1 的输出）占用较小：
  • batch_size×128 维向量：batch_size×128×4 Byte = batch_size×512 Byte
  • 例如 batch_size=1024 时，中间变量约 512 KB

以 SGD 为例，此时其他参数占用固定，batchsize 会影响显存占用

在 PyTorch 中，在使用 DataLoader 加载数据时，如果不指定 batch_size 参数，默认值是 1，即每次迭代返回一个样本。这与一次性使用全部数据进行训练是完全不同的概念。如果想要一次性使用全部数据进行训练，需要手动将 batch_size 设置为数据集的大小，但对于大型数据集，这样做通常会导致内存不足，因为一次性将所有数据加载到内存中可能会超出硬件的内存限制。

大规模数据时，通常从 16 开始测试，然后逐渐增加，确保代码运行正常且不报错，直到出现 内存不足（OOM）报错或训练效果下降，此时选择略小于该值的 batch_size。

训练时候搭配 nvidia-smi 监控显存占用，合适的 batch_size = 硬件显存允许的最大值 × 0.8（预留安全空间），并通过训练效果验证调整。

补充说明：batchsize 对于训练的影响在深度学习中，使用较大的 batch_size（批量大小）相比单样本训练（batch_size=1）有以下核心优势

• 并行计算能力最大化，减小训练时间；且大幅减少更新次数
• 梯度方向更准确，单样本训练的梯度仅基于单个数据点，可能包含大量噪声（尤其是数据分布不均或存在异常值时）。大 batch_size 的梯度是多个样本的平均值，能抵消单个样本的随机性，梯度方向更接近真实分布的 “全局最优方向”。会让训练过程更稳定，波动更小

勇闯python的第40天@浙大疏锦行