在多GPU设置（Pytorch）上反向传播期间的记忆中cuda cuda

Question

这里是我的培训功能的简化版本

def train_model(model, data_loader, num_epochs, sparse_matrix_path, save_path):
    device_net = torch.device('cuda:0')
    device_matrix = torch.device('cuda:1')
    model = model.to(device_net)

    # Load model and sparse matrix
    model.train()
    
    sparse_matrix = torch.load(sparse_matrix_path, weights_only=True).to_sparse().to(device_matrix)
    sparse_matrix.requires_grad = False

    optimizer = torch.optim.Adam(model.parameters(), lr=1e-4)
   

    for epoch in range(num_epochs):
        for step, input_data in enumerate(data_loader):
            input_data = input_data.to(device_net).unsqueeze(1)
            optimizer.zero_grad()

            # Model forward pass
            output, _ = model(input_data)
            

            # Sparse matrix operation
            batch_size = output.size(0)
            flat_output = output.view(batch_size, -1).to(device_matrix)
            sparse_output = torch.sparse.mm(sparse_matrix, flat_output.T)

            # Compute loss
            transformed_output = sparse_output.T.view(batch_size, 1, 700, 384)
            loss = F.mse_loss(transformed_output.to(device_net), input_data)
            loss.backward()  # Error occurs here
            
            optimizer.step()

键观察：

在device_matrix（GPU 1）上加载了稀疏矩阵（6 GB）。

当在CPU上运行时，lose.backward（）将RAM使用增加〜9 GB，总CPU为40GB.

• 在GPU中，内存使用量在损失期间大大增加。后返回（），即使批量大小为1.

• 问题：

• 为什么loss.backward（）导致GPU的大幅度记忆增加，其中有一个矩阵？

  如何优化此设置以在GPU上运行，而不会用完存储器？ 是否有特定的策略或pytorch功能可以减少向后通过的GPU内存使用量？

1. 为什么loss.backward（）导致GPU的大幅度记忆增加，其中有一个矩阵？

当您致电

Answer 1

时，计算网络的所有梯度。每个张量的参数有一个梯度。基本上，这将记忆使用量乘以两个。

如何优化此设置以在GPU上运行，而不会用完存储器？是否有特定的策略或pytorch功能可以减少向后通过的GPU内存使用量？

您可以尝试使用

混合精确训练

，降低输入的尺寸，使用SGD代替Adam，使用更少的FLOPS的网络（例如，很早的输入就下一个示例示例）或使用较大的GPU.

。