小Batch Size下Batch Normalisation性能问题及显存优化咨询

小Batch Size下Batch Normalisation性能恶化的合理性 #

这种情况完全合理。Batch Normalisation（BN）的核心是用当前训练batch的均值和方差对特征做归一化，同时维护全局的running mean和方差供测试阶段使用。当batch size小到16时，单batch的统计量（均值、方差）随机性极强，和全局真实统计量偏差极大。训练时BN层依赖这些不稳定的局部统计量更新参数，导致模型学习到的特征分布，和测试阶段用running mean归一化后的分布严重错位，最终模型完全无法泛化，表现出近乎随机分类的结果。尤其VGG16是深层模型，BN层的累积误差会被放大，这种现象会更显著。

突破GPU显存限制提升Batch Size的方法 #

针对GTX 1070的显存限制，以下方法可帮你将等效batch size提升到32：

• 混合精度训练：用PyTorch内置的torch.cuda.amp模块，自动将部分张量转为半精度（FP16）存储，大幅降低显存占用且几乎不损失精度。训练时只需添加少量代码：

  from torch.cuda.amp import GradScaler, autocast
  
  scaler = GradScaler()
  for epoch in range(epochs):
      for data in dataloader:
          inputs, labels = data
          optimizer.zero_grad()
          with autocast():
              outputs = model(inputs)
              loss = criterion(outputs, labels)
          scaler.scale(loss).backward()
          scaler.step(optimizer)
          scaler.update()
  

• 梯度累积：把多个小batch的梯度累积后再执行一次参数更新，等效于增大batch size。比如用batch size16，每累积2个batch的梯度再更新参数，就相当于等效batch size32。代码示例：

  accumulation_steps = 2
  for epoch in range(epochs):
      for i, data in enumerate(dataloader):
          inputs, labels = data
          outputs = model(inputs)
          loss = criterion(outputs, labels)
          loss = loss / accumulation_steps  # 均分损失避免梯度爆炸
          loss.backward()
          if (i + 1) % accumulation_steps == 0:
              optimizer.step()
              optimizer.zero_grad()
  

• 冻结更多底层参数：VGG16的底层卷积层学习的是通用视觉特征，迁移学习时可冻结前5个卷积块（共13层卷积），只训练最后1个卷积块、全连接层和新增的BN层，大幅减少需更新的参数数量，从而节省显存。

• 缩小输入图像尺寸：如果输入是标准224x224，可尝试缩小到192x192或160x160。每个样本的显存占用和图像尺寸平方成正比，缩小后能显著提升可容纳的batch size。注意训练和测试的预处理流程要一致，避免数据分布偏移。

• 激活值Checkpointing：使用torch.utils.checkpoint.checkpoint函数，通过重新计算部分中间层的激活值节省显存。对于VGG16这类激活值占显存较多的模型，能有效降低显存压力，但会略微增加训练时间。

内容的提问来源于stack exchange，提问作者RickarySanchez