Python打卡第40天

@浙大疏锦行

图像数据、MLP图像模型、Batch Size与显存计算

来源：复试班CV学习笔记，PyTorch实现MNIST/CIFAR10图像MLP，详解图像维度、模型定义、显存估算、OOM原理

一、图像数据介绍

1.1 灰度图像（MNIST手写数字）

图像数据和表格结构化数据最大区别：结构化数据形状一般为(样本数, 特征数)；图像需要保留空间信息（高度、宽度、通道），形状为(通道，高，宽)。

MNIST是28×28灰度图：

通道数=1，只有黑白信息；像素原始取值0~255
PyTorch采用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',
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()
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)

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',
    train=True,
    download=True,
    transform=transform
)

trainloader = torch.utils.data.DataLoader(trainset, batch_size=4, shuffle=True)
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_size=10，MNIST输入shape：[10,1,28,28]

二、图像MLP模型定义

MLP全连接网络只能接收一维向量 ，图像是二维空间结构，所以必须先展平（Flatten）。两种写法：nn.Flatten() 或者 x.view(-1, 784)

2.1 MNIST灰度图像MLP

复制代码

class 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×128=100352，偏置128，合计100480
layer2：权重128×10=1280，偏置10，合计1290
总参数量：101770

2.2 CIFAR10彩色图像MLP

3通道32×32图像，展平后维度：3×32×32=3072

复制代码

class MLP(nn.Module):
    def __init__(self, input_size=3072, hidden_size=128, num_classes=10):
        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))

参数： fc1：3072×128+128=393344；fc2：128×10+10=1290；总参数394634

2.3 模型与batch_size的关系（重点！）

✅ 模型定义本身和batch_size完全无关

nn.Module写网络时，不用指定batch大小；
nn.Flatten会自动保留batch维度，只对单样本做展平；
torchsummary的input_size只写单样本尺寸，不带batch；
batch_size是在DataLoader中设置，仅控制每次迭代读取多少样本。

from torch.utils.data import DataLoader
train_loader = DataLoader(train_dataset, batch_size=64, shuffle=True)
test_loader = DataLoader(test_dataset, batch_size=1000, shuffle=False)

三、显存占用组成 & 显存估算

训练时显存由4大部分构成：

模型参数Parameters：权重，FP32每个参数占4Byte
梯度Gradients：反向传播生成，数量和参数一致，FP32占4Byte
优化器状态：Adam额外保存动量m、方差v；SGD无额外存储
中间激活值（前向反向中间张量）+ 当前batch数据 ：这部分随batch_size变大而增加

基础单位换算

float32：4 Byte；uint8：1 Byte
1KB=1024 Byte；1MB=1024KB

MNIST模型参数101770：

参数：101770 ×4 ≈403KB
梯度：同参数≈403KB，合计806KB
Adam优化器额外存储m、v，再加806KB；SGD不加额外显存

单张MNIST图像转float32：1×28×28×4 = 3136 Byte ≈3KB batch_size越大，batch数据+中间激活占用显存越高。

batch_size	数据占用	中间变量	总显存占用（SGD，近似）
64	192 KB	32 KB	~1 MB
256	768 KB	128 KB	~1.7 MB
1024	3 MB	512 KB	~4.5 MB
4096	12 MB	2 MB	~15 MB

OOM（Out Of Memory）：显存不够，直接报错。 batch调参策略：从小batch开始试，逐步增大；设置为显卡显存上限×0.8预留安全空间；训练时nvidia-smi监控显存。

四、Batch Size原理回顾

batch_size：一次前向+反向传播送入网络的样本数量。一个batch内：计算全部样本loss，求平均loss，反向得到一组平均梯度，更新一次权重。

❗不是每个样本单独训练得到一套参数，最后参数取平均；是梯度求平均。

batch_size大：梯度更稳定，GPU并行效率高；同epoch参数更新次数变少，显存占用高，更容易过拟合。
batch_size小：梯度噪声大，震荡；噪声相当于弱正则，泛化能力更好；显存开销低。
batch_size=1：SGD随机梯度下降，梯度抖动剧烈。

五、总结

图像数据是[C,H,W]（PyTorch），送入网络自动带上batch维度[B,C,H,W]，绘图要转成[H,W,C]。
MLP处理图像必须展平，会丢失空间位置信息，后续CNN就是为了解决这个问题。
模型结构定义和batch size无关，batch只在DataLoader配置。
显存分为固定部分（参数、梯度、优化器状态）和动态部分（batch数据、中间激活），动态部分随batch_size增大而上升。
batch本质：用一批样本的平均梯度近似全集梯度，在梯度稳定性、显存、泛化能力之间做权衡。

Python打卡第40天

图像数据、MLP图像模型、Batch Size与显存计算

一、图像数据介绍

1.1 灰度图像（MNIST手写数字）

设置随机种子，保证实验可复现

数据预处理管道

加载数据集，自动下载

随机选取一张图片查看

反归一化，可视化灰度图

1.2 彩色图像（CIFAR10）

二、图像MLP模型定义

2.1 MNIST灰度图像MLP

2.2 CIFAR10彩色图像MLP

2.3 模型与batch_size的关系（重点！）

三、显存占用组成 & 显存估算

基础单位换算

四、Batch Size原理回顾

五、总结