浙大疏锦行图像数据、MLP图像模型、Batch Size与显存计算来源复试班CV学习笔记PyTorch实现MNIST/CIFAR10图像MLP详解图像维度、模型定义、显存估算、OOM原理一、图像数据介绍1.1 灰度图像MNIST手写数字图像数据和表格结构化数据最大区别 结构化数据形状一般为(样本数, 特征数)图像需要保留空间信息高度、宽度、通道形状为(通道高宽)。MNIST是28×28灰度图通道数1只有黑白信息像素原始取值0~255PyTorch采用Channel First[C, H, W]也就是[1, 28, 28]import torch import torch.nn as nn import torch.optim as optim from torch.utils.data import DataLoader , Dataset from torchvision import datasets, transforms import matplotlib.pyplot as plt # 设置随机种子保证实验可复现 torch.manual_seed(42) # 数据预处理管道 transform transforms.Compose([ transforms.ToTensor(), # 转为张量像素归一化至[0,1] transforms.Normalize((0.1307,), (0.3081,)) # MNIST专用均值方差 ]) # 加载数据集自动下载 train_dataset datasets.MNIST( root./data, trainTrue, downloadTrue, transformtransform ) test_dataset datasets.MNIST( root./data, trainFalse, transformtransform ) # 随机选取一张图片查看 sample_idx torch.randint(0, len(train_dataset), size(1,)).item() image, label train_dataset[sample_idx] # 反归一化可视化灰度图 def imshow(img): img img * 0.3081 0.1307 npimg img.numpy() plt.imshow(npimg[0], cmapgray) plt.show() print(fLabel: {label}) imshow(image)1.2 彩色图像CIFAR10彩色RGB图像通道数3形状[3,32,32]。重点区分PyTorch张量Channel First(C, H, W)Numpy/Matplotlib绘图Channel Last(H, W, C)绘图时必须用np.transpose(npimg, (1,2,0))调换维度顺序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(), transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)) ]) trainset torchvision.datasets.CIFAR10( root./data, trainTrue, downloadTrue, transformtransform ) trainloader torch.utils.data.DataLoader(trainset, batch_size4, shuffleTrue) 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}) print(f图像类别: {classes[label]}) def imshow(img): img img / 2 0.5 npimg img.numpy() plt.imshow(np.transpose(npimg, (1, 2, 0))) plt.axis(off) plt.show() imshow(image)图像送入网络训练时会额外增加batch维度[batch, C, H, W]。 例如batch_size10MNIST输入shape[10,1,28,28]二、图像MLP模型定义MLP全连接网络只能接收一维向量图像是二维空间结构所以必须先展平Flatten。 两种写法nn.Flatten()或者x.view(-1, 784)2.1 MNIST灰度图像MLPclass MLP(nn.Module): def __init__(self): super(MLP, self).__init__() self.flatten nn.Flatten() # [B,1,28,28] - [B,784] self.layer1 nn.Linear(784, 128) self.relu nn.ReLU() self.layer2 nn.Linear(128, 10) def forward(self, x): x self.flatten(x) x self.layer1(x) x self.relu(x) x self.layer2(x) return x model MLP() device torch.device(cuda if torch.cuda.is_available() else cpu) model model.to(device) from torchsummary import summary print(\n模型结构信息) summary(model, input_size(1, 28, 28))参数计算layer1权重784×128100352偏置128合计100480layer2权重128×101280偏置10合计1290总参数量1017702.2 CIFAR10彩色图像MLP3通道32×32图像展平后维度3×32×323072class MLP(nn.Module): def __init__(self, input_size3072, hidden_size128, num_classes10): super(MLP, self).__init__() 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) # [B,3,32,32] - [B,3072] x self.fc1(x) x self.relu(x) x self.fc2(x) return x model MLP() device torch.device(cuda if torch.cuda.is_available() else cpu) model model.to(device) from torchsummary import summary print(\n模型结构信息) summary(model, input_size(3, 32, 32))参数 fc13072×128128393344fc2128×10101290总参数3946342.3 模型与batch_size的关系重点✅模型定义本身和batch_size完全无关nn.Module写网络时不用指定batch大小nn.Flatten会自动保留batch维度只对单样本做展平torchsummary的input_size只写单样本尺寸不带batchbatch_size是在DataLoader中设置仅控制每次迭代读取多少样本。from torch.utils.data import DataLoader train_loader DataLoader(train_dataset, batch_size64, shuffleTrue) test_loader DataLoader(test_dataset, batch_size1000, shuffleFalse)三、显存占用组成 显存估算训练时显存由4大部分构成模型参数Parameters权重FP32每个参数占4Byte梯度Gradients反向传播生成数量和参数一致FP32占4Byte优化器状态Adam额外保存动量m、方差vSGD无额外存储中间激活值前向反向中间张量 当前batch数据这部分随batch_size变大而增加基础单位换算float324 Byteuint81 Byte1KB1024 Byte1MB1024KBMNIST模型参数101770参数101770 ×4 ≈403KB梯度同参数≈403KB合计806KBAdam优化器额外存储m、v再加806KBSGD不加额外显存单张MNIST图像转float321×28×28×4 3136 Byte ≈3KB batch_size越大batch数据中间激活占用显存越高。batch_size数据占用中间变量总显存占用SGD近似64192 KB32 KB~1 MB256768 KB128 KB~1.7 MB10243 MB512 KB~4.5 MB409612 MB2 MB~15 MBOOMOut Of Memory显存不够直接报错。 batch调参策略从小batch开始试逐步增大设置为显卡显存上限×0.8预留安全空间训练时nvidia-smi监控显存。四、Batch Size原理回顾batch_size一次前向反向传播送入网络的样本数量。 一个batch内计算全部样本loss求平均loss反向得到一组平均梯度更新一次权重。❗不是每个样本单独训练得到一套参数最后参数取平均是梯度求平均。batch_size大梯度更稳定GPU并行效率高同epoch参数更新次数变少显存占用高更容易过拟合。batch_size小梯度噪声大震荡噪声相当于弱正则泛化能力更好显存开销低。batch_size1SGD随机梯度下降梯度抖动剧烈。五、总结图像数据是[C,H,W]PyTorch送入网络自动带上batch维度[B,C,H,W]绘图要转成[H,W,C]。MLP处理图像必须展平会丢失空间位置信息后续CNN就是为了解决这个问题。模型结构定义和batch size无关batch只在DataLoader配置。显存分为固定部分参数、梯度、优化器状态和动态部分batch数据、中间激活动态部分随batch_size增大而上升。batch本质用一批样本的平均梯度近似全集梯度在梯度稳定性、显存、泛化能力之间做权衡。