解决显存不足问题的方案

在深度学习训练过程中，显存不足是一个常见的问题。以下是解决显存不足问题的一些常见方法，特别适用于使用PyTorch进行训练时。本文将提供逐步指导，帮助您找到合适的批次大小，并避免显存不足问题。

方法一：调整批次大小

批次大小（batch size）是指在每次迭代中处理的数据量。较大的批次大小会占用更多显存，而较小的批次大小则会减少显存使用。

步骤

1. 开始时使用较小的批次大小（例如2）。
2. 逐步增加批次大小（例如每次增加2或4），直到遇到内存错误。
3. 记录最大可行的批次大小，并在该值的基础上进行训练。

示例代码

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)

方法二：使用梯度累积

梯度累积允许您在多个小批次上累积梯度，以模拟较大的批次大小。这可以减少每个批次的显存使用。

修改训练函数以实现梯度累积

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)}')

方法三：清空缓存

在训练过程中定期清空CUDA缓存，可以释放不再使用的显存。

示例代码

在训练循环中添加清空缓存的代码：

torch.cuda.empty_cache()

在每个epoch结束后调用，以确保释放未使用的显存：

scheduler.step(epoch_loss)
torch.cuda.empty_cache()  # 清空缓存
print(f'Epoch {epoch + 1}/{num_epochs}, Loss: {epoch_loss / len(train_loader)}')

方法四：使用混合精度训练

混合精度训练可以通过使用torch.cuda.amp来实现。这可以显著减少显存使用，同时提高计算效率。

示例代码

在训练函数中使用混合精度训练：

from torch.cuda.amp import autocast, GradScaler

scaler = GradScaler()  # 初始化混合精度缩放器

# 在训练循环中使用autocast和scaler
with autocast():
    outputs = model(images)
    loss = criterion(outputs, labels)
    scaler.scale(loss).backward()

scaler.step(optimizer)
scaler.update()
optimizer.zero_grad()