深度学习显存管理新技巧：让你的代码飞速运行！
深度学习显存管理新技巧：让你的代码飞速运行！
在深度学习模型的训练和推理过程中，显存管理是一个非常重要的方面。特别是在使用大型模型时，显存管理不当可能会导致显存不足，从而导致程序崩溃或性能下降。本文介绍了一些显存管理的方法和技巧，如及时释放显存、使用上下文管理器、分批处理数据、动态调整显存分配以及混合精度训练等，帮助你在使用循环运行代码时更好地管理显存，提升代码运行效率。
什么是显存？
显存，全称图形处理器内存，是独立显卡上的存储器，主要用于存储和处理图形数据。在运行深度学习模型时，显存被用来存储和操作大量的模型参数和中间计算结果。然而，由于显卡内存容量有限，当模型过大或数据集过多时，很容易出现显存不足的问题，导致机器学习任务无法正常进行。
显存管理的重要性
- 提高机器学习速度：释放显存可以减少计算过程中的数据传输延迟，避免GPU资源的浪费，从而加快机器学习的速度。
- 减轻显卡压力：通过释放显存，可以使得显卡内存得到有效利用，降低显卡的压力，延长其使用寿命。
- 优化内存管理：合理地释放显存有助于优化内存管理，使得机器学习任务能够更好地适应不同的硬件环境，提高代码的鲁棒性。
- 避免内存溢出：在处理大型模型和数据集时，释放显存可以避免显存溢出的问题，保证机器学习任务的正常运行。
常见的显存优化技巧
1. 减小Batch Size
批次大小（batch size）是指在每次迭代中处理的数据量。较大的批次大小会占用更多显存，而较小的批次大小则会减少显存使用。
步骤：
- 开始时使用较小的批次大小（例如2）。
- 逐步增加批次大小（例如每次增加2或4），直到遇到内存错误。
- 记录最大可行的批次大小，并在该值的基础上进行训练。
示例代码：
initial_batch_size = 2
max_batch_size = initial_batch_size
# 尝试找到最大可行的批次大小
while True:
try:
train_loader = DataLoader(train_dataset, batch_size=max_batch_size, shuffle=True)
val_loader = DataLoader(val_dataset, batch_size=max_batch_size, shuffle=False)
# 进行一次测试训练
train_model(train_loader, model, criterion, optimizer, scheduler, num_epochs=1, accumulation_steps=4)
# 验证模型
validate_model(val_loader, model, criterion)
# 如果成功，增加批次大小
max_batch_size += 2
except RuntimeError as e:
if 'out of memory' in str(e):
print(f"Out of memory at batch size: {max_batch_size}")
torch.cuda.empty_cache()
break
else:
raise e
print(f"Max batch size found: {max_batch_size - 2}")
# 使用找到的最大批次大小进行训练
train_loader = DataLoader(train_dataset, batch_size=max_batch_size - 2, shuffle=True)
val_loader = DataLoader(val_dataset, batch_size=max_batch_size - 2, shuffle=False)
train_model(train_loader, model, criterion, optimizer, scheduler, accumulation_steps=4)
validate_model(val_loader, model, criterion)
2. 使用混合精度训练
一般默认情况下， 整个网络中采用的是32位的浮点数，如果切换到 16位的浮点数，其显存占用量将接近呈倍数递减。
修改训练函数以实现混合精度训练：
from torch.cuda.amp import autocast, GradScaler
def train_model(train_loader, model, criterion, optimizer, scheduler, num_epochs=100, accumulation_steps=4):
model.train()
scaler = GradScaler() # 初始化混合精度缩放器
for epoch in range(num_epochs):
epoch_loss = 0
optimizer.zero_grad() # 初始化优化器梯度
for i, (images, labels) in enumerate(tqdm(train_loader)):
images = images.cuda()
labels = labels.cuda()
batch_size, num_images, c, h, w = images.size()
images = images.view(-1, c, h, w) # 将批次的图像展平
with autocast(): # 使用混合精度
outputs = model(images)
outputs = outputs.view(batch_size, num_images, -1).mean(dim=1) # 平均每个病人的输出
loss = criterion(outputs, labels)
loss = loss / accumulation_steps # 将损失除以累积步数
scaler.scale(loss).backward() # 使用缩放器进行反向传播
if (i + 1) % accumulation_steps == 0: # 当累积步数达到时进行优化
scaler.step(optimizer)
scaler.update()
optimizer.zero_grad()
epoch_loss += loss.item() * accumulation_steps # 累加损失
scheduler.step(epoch_loss)
torch.cuda.empty_cache() # 清空缓存
print(f'Epoch {epoch + 1}/{num_epochs}, Loss: {epoch_loss / len(train_loader)}')
3. 梯度累积
梯度累积允许您在多个小批次上累积梯度，以模拟较大的批次大小。这可以减少每个批次的显存使用。
示例代码：
def train_model(train_loader, model, criterion, optimizer, scheduler, num_epochs=100, accumulation_steps=4):
model.train()
for epoch in range(num_epochs):
epoch_loss = 0
optimizer.zero_grad() # 初始化优化器梯度
for i, (images, labels) in enumerate(tqdm(train_loader)):
images = images.cuda()
labels = labels.cuda()
batch_size, num_images, c, h, w = images.size()
images = images.view(-1, c, h, w) # 将批次的图像展平
outputs = model(images)
outputs = outputs.view(batch_size, num_images, -1).mean(dim=1) # 平均每个病人的输出
loss = criterion(outputs, labels)
loss = loss / accumulation_steps # 将损失除以累积步数
loss.backward() # 反向传播
if (i + 1) % accumulation_steps == 0: # 当累积步数达到时进行优化
optimizer.step()
optimizer.zero_grad()
epoch_loss += loss.item() * accumulation_steps # 累加损失
scheduler.step(epoch_loss)
torch.cuda.empty_cache() # 清空缓存
print(f'Epoch {epoch + 1}/{num_epochs}, Loss: {epoch_loss / len(train_loader)}')
4. 模型结构优化
优化网络结构，或者改用深度可分离卷积代替常规卷积核，较小参数数量。
5. 显存释放机制
在TensorFlow中，可以使用tf.keras.backend.clear_session()来释放不再需要的显存，为后续任务腾出空间。
示例代码：
import tensorflow as tf
# 创建一个大型模型
model = tf.keras.models.Sequential([
tf.keras.layers.Dense(1024, activation='relu', input_shape=(784,)),
tf.keras.layers.Dense(10)
])
# 加载数据集
(x_train, y_train), (x_test, y_test) = tf.keras.datasets.mnist.load_data()
x_train = x_train.reshape((-1, 784)) / 255.0
x_test = x_test.reshape((-1, 784)) / 255.0
# 训练模型
for epoch in range(10):
model.compile(optimizer='adam', loss='sparse_categorical_crossentropy', metrics=['accuracy'])
history = model.fit(x_train, y_train, epochs=1, validation_data=(x_test, y_test))
# 在每个epoch结束后释放显存
del history
tf.keras.backend.clear_session() # 调用clear_session()释放显存
框架特定的显存管理
TensorFlow
TensorFlow提供了一个名为TF-GPU的工具，用于管理GPU内存。通过使用TF-GPU，可以有效地释放显存，提高计算效率。
在TensorFlow中，可以使用sess.run()方法来执行计算图（computation graph）中的操作。在sess.run()过程中，可以指定需要运行的节点（nodes），从而有效地释放未使用的显存。
PyTorch
PyTorch框架对用户的响应是实时互动的，采用了动态申请管理显存的方式，所以框架申请显存会实时调整（动态变化）。最大优点之一时：不会占用过量的显存，方便多人进行同时使用一张设备。但是在这种方式下不可避免会遇到一些问题，如，怎样设计申请频率，处理好API的时间消耗；如何处理管理机制带来的碎片问题。
PyTorch1.13版本显存的管理主要采用cudaMalloc方式实现，实现时需考虑的问题：
- cudaMalloc 动态申请多少显存？
- 申请了之后如何分配？
- 申请了之后各个stream之间如何协调？
- 什么时候调用cudaFree释放？
解决这些问题，在cudaCachingAllocator中采用池化的显存管理，下面按照显存使用的生命周期依次介绍。
首先是显存的申请/创建问题。在管理机制中，将显存的申请与使用过程进行了分离，即显存申请后会进行二次分配，其过程是：先通过cudaMalloc申请一个显存块segment，然后从segment分离出来的子显存块block，框架的上层应用使用的是分离后的block显存，上层应用不直接使用segment。
进一步，通过池化的方式将block按照块放入不同的显存池中，进行分类管理。
2.1 管理逻辑1: size触发创建
显存管理机制是根据申请size来决定从GPU创建多大的segment，以及是否要进行切分（split）。
总结
显存管理是深度学习开发中不可或缺的一环。通过合理运用上述技巧，可以有效提升模型训练和推理的效率。然而，最佳实践往往需要根据具体场景灵活调整。建议开发者在实际项目中不断尝试和优化，找到最适合自身需求的显存管理方案。