基础介绍
在计算机视觉深度学习网络中，在训练阶段数据输入通常是一个批次，即不是一次输入单张图片，而是一次性输入多张图片，而神经网络的结构内部一次只能处理一张图片，这时候很自然就会考虑为什么要这样的输入？神经网络是如何处理多个数据的，下面从硬件架构的角度去分析处理。
GPU
硬件架构
GPU的硬件架构设计是批处理能够高效运行的关键原因之一。GPU现阶段一般采用SIMT架构，它的特点如下：
SIMT（Single Instruction, Multiple Threads）架构是一种并行计算模型，主要用于图形处理单元（GPU）和其他高性能计算平台。它与传统的 SIMD（Single Instruction, Multiple Data）架构有所不同，尽管两者都允许同时执行多个操作。
在SIMT架构中，多个线程同时执行相同的指令，这类似于SIMD。不同之处在于，SIMT允许每个线程拥有自己的寄存器和状态，因此它们可以在执行过程中有不同的执行路径。这种灵活性使得SIMT能够处理更多种类的计算任务，适合于并行计算和图形处理中的高度数据并行任务
SIMT可以将计算任务分为多个线程，每个线程独立运行，这使得程序可以有效的利用并行硬件资源，提升计算性能。注意有的cpu也支持了SIMD架构，如intel的SSE和avx。所以很多使用使用cpu进行训练时就是调用的cpu的SSE或avx。
GPU的硬件架构特征：
CPU: 少量强大的核心，适合串行处理
[Core 1] [Core 2] [Core 3] [Core 4]GPU: 大量简单的核心，适合并行处理
[Core][Core][Core][Core][Core][Core]...(数千个核心)
[Core][Core][Core][Core][Core][Core]...
[Core][Core][Core][Core][Core][Core]...
有的GPU可以多大上千个核心，而CPU核心的数量则远远达不到这个数量。
GPU的特点
- 高吞吐量
单精度浮点运算性能（FLOPS）示例：
高端CPU： ~1-2 TFLOPS
高端GPU： ~30-40 TFLOPS
- 高内存带宽
内存带宽示例：
CPU： ~50-100 GB/s
GPU： ~700-1000 GB/s
流处理器
# GPU处理批次数据的示意
batch_images = torch.randn(32, 3, 224, 224) # 32张图片
# GPU会自动将计算分配到多个流处理器
# 假设一个简单的操作
def gpu_parallel_operation(batch):
# 在GPU中，这个操作会自动分配到多个处理器并行执行
return batch * 2 # 每个元素的操作都可以并行
GPU与批处理的关系
并行计算能力
class ConvNet(nn.Module):
def __init__(self):
super().__init__()
self.conv = nn.Conv2d(3, 64, 3)
def forward(self, x):
# batch_size=32时
# GPU可以同时处理32张图片的卷积运算
# 每个CUDA核心可以负责部分计算
return self.conv(x)
# 使用GPU
model = ConvNet().cuda()
batch = torch.randn(32, 3, 224, 224).cuda()
output = model(batch) # 并行处理32张图片
内存层级
GPU内存层级：
Global Memory (显存) -> L2 Cache -> L1 Cache/Shared Memory -> Registers
数据流动示例：
batch_data (Global Memory)
→ 加载到Shared Memory
→ 分配给各个CUDA核心的Registers
→ 并行计算
→ 结果写回Global Memory
如何选择合适批处理大小
可以从一个角度出发，即显存的利用率。通过设置一个合适的显存利用率来判断当前GPU的利用是否充分。一个简单的方法如下：
# 批大小需要平衡以下因素：
# 1. GPU显存大小
# 2. 计算效率
# 3. 训练效果
# 示例：显存管理
def get_optimal_batch_size(model, input_size, max_memory=0.8):
"""估算最优批大小"""
try:
batch_size = 1
while True:
# 尝试运行一个批次
x = torch.randn(batch_size, *input_size).cuda()
_ = model(x)
# 检查显存使用
memory_used = torch.cuda.memory_allocated() / torch.cuda.max_memory_allocated()
if memory_used > max_memory:
return batch_size - 1
batch_size *= 2
except RuntimeError: # 显存溢出
return batch_size // 2
利用gpu的流水线
# 使用DataLoader的pin_memory和num_workers优化数据加载
train_loader = DataLoader(
dataset,
batch_size=32,
pin_memory=True, # 将数据固定在内存中，加速GPU传输
num_workers=4 # 多进程数据加载
)
计算和数据传输重叠
# 使用CUDA流实现计算和数据传输重叠
stream1 = torch.cuda.Stream()
stream2 = torch.cuda.Stream()
# 在两个流上并行处理
with torch.cuda.stream(stream1):
output1 = model(batch1)
with torch.cuda.stream(stream2):
output2 = model(batch2)
模型参数体积大于GPU显存应该如何处理？
首先检查显存的状态：显存总量，显存当前使用量。方法如下：
import torch
import psutil
import GPUtil
def memory_check():
# 检查系统内存
system_memory = psutil.virtual_memory()
print(f"系统内存总量: {system_memory.total / 1e9:.2f}GB")
print(f"系统内存使用: {system_memory.used / 1e9:.2f}GB")
# 检查GPU显存
if torch.cuda.is_available():
gpu = GPUtil.getGPUs()[0]
print(f"GPU显存总量: {gpu.memoryTotal/1024:.2f}GB")
print(f"GPU显存使用: {gpu.memoryUsed/1024:.2f}GB")
查明了GPU的当前状态后，如果模型参数的大小超过了可使用的显存，基本有如下几种思路去解决：
- 梯度检查点：这个属于不懂的地方
- 混合精度训练：这个也属于不懂的
- 模型并行策略，即将模型数据按块分别加载到不同的GPU上，每个GPU完成模型推理的某一部分
- 流水线并行：这个也属于不懂的
- 模型压缩技术：量化和知识蒸馏
//模型压缩技术
# 1. 量化
def quantization_example():
import torch.quantization
# 动态量化
model_int8 = torch.quantization.quantize_dynamic(
model, # 原FP32模型
{torch.nn.Linear}, # 量化层类型
dtype=torch.qint8 # 量化为8位整数
)
# 显存节省约75%
# 2. 知识蒸馏
class SmallModel(torch.nn.Module):
def __init__(self):
super().__init__()
# 更小的架构
self.features = torch.nn.Sequential(
torch.nn.Conv2d(3, 64, 3),
torch.nn.ReLU()
)
def distillation_training(teacher, student, data):
criterion = torch.nn.KLDivLoss()
temp = 3.0 # 温度参数
# 教师模型输出
with torch.no_grad():
teacher_outputs = teacher(data)
# 学生模型输出
student_outputs = student(data)
# 蒸馏损失
loss = criterion(
torch.log_softmax(student_outputs/temp, dim=1),
torch.softmax(teacher_outputs/temp, dim=1)
)
//模型分割和流水线并行
# 1. 模型分割
class SplitModel(torch.nn.Module):
def __init__(self):
super().__init__()
# 前半部分放在第一个GPU
self.features1 = torch.nn.Sequential(
torch.nn.Conv2d(3, 64, 3),
torch.nn.ReLU()
).to('cuda:0')
# 后半部分放在第二个GPU
self.features2 = torch.nn.Sequential(
torch.nn.Conv2d(64, 128, 3),
torch.nn.ReLU()
).to('cuda:1')
def forward(self, x):
x = self.features1(x.to('cuda:0'))
x = self.features2(x.to('cuda:1'))
return x
# 2. 梯度累积
def gradient_accumulation(model, dataloader, num_accumulation_steps=4):
optimizer.zero_grad()
for i, (data, target) in enumerate(dataloader):
output = model(data)
# 损失除以累积步数
loss = criterion(output, target) / num_accumulation_steps
loss.backward()
if (i + 1) % num_accumulation_steps == 0:
optimizer.step()
optimizer.zero_grad()
//显存优化技术
# 1. 显存优化器
def memory_efficient_training():
# 使用Adam优化器的显存效率版本
from torch.optim import AdamW
optimizer = AdamW(model.parameters(), lr=1e-4, weight_decay=0.01)
# 使用梯度检查点
from torch.utils.checkpoint import checkpoint
def forward_with_checkpoint(self, x):
return checkpoint(self.block, x)
# 2. 混合精度训练
def mixed_precision_example():
scaler = torch.cuda.amp.GradScaler()
for data, target in dataloader:
with torch.cuda.amp.autocast():
output = model(data)
loss = criterion(output, target)
scaler.scale(loss).backward()
scaler.step(optimizer)
scaler.update()
GPU缓存的作用
- 加速数据传输（Pin_memory）
- 减少内存碎片
- 优化内存分配效率
- 提供快速的内存重用
如何理解GPU快速处理矩阵计算
假设我们要计算两个矩阵相乘：C = A × B
# 假设矩阵大小
A: (M × K)
B: (K × N)
C: (M × N)
GPU内部如何对这个计算任务进行分配呢？请看下面的介绍
# 每个线程负责计算C矩阵中的一个元素
# 例如计算C[2,3]的线程：
Thread(2,3) 负责计算：
C[2,3] = sum(A[2,:] * B[:,3])
每个线程负责计算矩阵中的一个元素，等所有线程执行完成后，矩阵的计算结果也出来了。每个线程将计算的结果放到共享内存的指定位置。从这个角度看可以矩阵运算实现了并行。
总结
GPU非常适合于矩阵运算、卷积运算、元素级操作（比如每个元素乘2）。现在的GPU有的显存能够达到80GB甚至更高。
CPU
CPU也可以处理批次数据，这里不重点介绍。主要有以下手段：
- 向量化操作（SIMD）:Numpy和Pytorch在cpu上的运算会自动使用SIMD指令（如SSE,AVX）
- 多线程处理
- 多进程处理