如何在Google Colab中提升A100 GPU显存利用率以加速PyTorch训练？

问题描述 #

基于PyTorch开发的全连接神经网络（FCNN）项目，在Google Colab Pro使用A100（40GB显存）训练时，速度与本地CPU几乎无差异，显存仅占用0.6GB。已更换GPU实例、添加device = torch.device("cuda:0")并迁移模型至GPU，但问题仍存在。

已尝试操作 #

• 更换Colab中的GPU实例
• 添加CUDA设备指定代码并将网络迁移至GPU

原代码 #

import torch
import torchvision
import torchvision.transforms as transforms
import torch.nn as nn
import torch.optim as optim


device = torch.device("cuda:0")
# Define transform
transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.5,), (0.5,))
])

# Load FashionMNIST dataset
trainset = torchvision.datasets.FashionMNIST('./data', download=True, train=True, transform=transform)
testset = torchvision.datasets.FashionMNIST('./data', download=True, train=False, transform=transform)

# Create data loaders
trainloader = torch.utils.data.DataLoader(trainset, batch_size=1, shuffle=True, num_workers=2)
testloader = torch.utils.data.DataLoader(testset, batch_size=1  , shuffle=False, num_workers=2)

# Define constant for classes
classes = ('T-shirt/top', 'Trouser', 'Pullover', 'Dress', 'Coat',
           'Sandal', 'Shirt', 'Sneaker', 'Bag', 'Ankle Boot')




# Define the fully connected neural network
class FCNN(nn.Module):
    def __init__(self, num_layers=1):
        super(FCNN, self).__init__()
        self.num_layers = num_layers
        self.fc_layers = nn.ModuleList()
        if self.num_layers == 1:
            self.fc_layers.append(nn.Linear(28 * 28, 1024))
        elif self.num_layers == 2:
            self.fc_layers.append(nn.Linear(28 * 28, 1024))
            self.fc_layers.append(nn.Linear(1024, 1024))
        self.output_layer = nn.Linear(1024, 10)

    def forward(self, x):
        x = x.view(-1, 28 * 28)
        for layer in self.fc_layers:
            x = nn.functional.relu(layer(x))
        x = self.output_layer(x)
        return x

# Modify the train function to move inputs and labels to the GPU
def train(net, criterion, optimizer, epochs=15):
    for epoch in range(epochs):
        running_loss = 0.0
        for i, data in enumerate(trainloader, 0):
            inputs, labels = data[0].to(device), data[1].to(device)
            optimizer.zero_grad()

            outputs = net(inputs)
            loss = criterion(outputs, labels)
            loss.backward()
            optimizer.step()

            running_loss += loss.item()
            if i % 2000 == 1999:
                print('[%d, %5d] loss: %.2f' %
                      (epoch + 1, i + 1, running_loss / 2000))
                running_loss = 0.0

# Define function to test accuracy
def test(net):
    correct = 0
    total = 0
    with torch.no_grad():
        for data in testloader:
            images, labels = data
            outputs = net(images)
            _, predicted = torch.max(outputs.data, 1)
            total += labels.size(0)
            correct += (predicted == labels).sum().item()

    print('Accuracy: %d %%' % (
            100 * correct / total))

# Main function
if __name__ == "__main__":
    # Define the network
    net1 = FCNN(num_layers=1)
    net2 = FCNN(num_layers=2)
    net2.to(device)

    # Define loss function and optimizer
    criterion = nn.CrossEntropyLoss()
    optimizer1 = optim.SGD(net1.parameters(), lr=0.001, momentum=0.0)
    optimizer2 = optim.SGD(net2.parameters(), lr=0.001, momentum=0.0)

    # Train and test network with 1 FC layer
    #print("Training network with 1 layer...")
    #train(net1, criterion, optimizer1)
    #test(net1)

    # Train and test network with 2 FC layers
    print("Training network with 2 layers...")
    train(net2, criterion, optimizer2)
    test(net2)

问题根源分析 #

1. 极小的batch_size：当前设置batch_size=1，GPU无法发挥并行计算优势，大部分时间消耗在数据传输而非计算上，导致显存利用率极低。
2. 测试环节未使用GPU：test函数中未将输入数据和模型移至GPU，全程在CPU运行，影响整体效率。
3. 优化器参数不合理：lr=0.001过小且momentum=0.0，导致训练收敛极慢，进一步放大了GPU未被充分利用的问题。
4. 数据加载效率不足：num_workers=2在Colab环境中可适当调高，提升数据预处理和加载速度，避免GPU等待数据。

解决方案 #

1. 调整batch_size #

将DataLoader的batch_size增大至64或128，充分利用GPU的并行计算能力：

trainloader = torch.utils.data.DataLoader(trainset, batch_size=128, shuffle=True, num_workers=4)
testloader = torch.utils.data.DataLoader(testset, batch_size=128, shuffle=False, num_workers=4)

2. 修正test函数 #

将测试数据和模型移至GPU：

def test(net):
    correct = 0
    total = 0
    net.to(device)  # 确保模型在GPU上
    with torch.no_grad():
        for data in testloader:
            images, labels = data[0].to(device), data[1].to(device)
            outputs = net(images)
            _, predicted = torch.max(outputs.data, 1)
            total += labels.size(0)
            correct += (predicted == labels).sum().item()

    print('Accuracy: %d %%' % (100 * correct / total))

3. 优化器参数调整 #

增大学习率并添加动量，加速收敛：

optimizer2 = optim.SGD(net2.parameters(), lr=0.01, momentum=0.9)

4. 迁移损失函数至GPU #

将损失函数也移至GPU，避免张量设备不匹配：

criterion = nn.CrossEntropyLoss().to(device)

修改后完整代码片段（关键部分） #

# 调整数据加载器
trainloader = torch.utils.data.DataLoader(trainset, batch_size=128, shuffle=True, num_workers=4)
testloader = torch.utils.data.DataLoader(testset, batch_size=128, shuffle=False, num_workers=4)

# 修正测试函数
def test(net):
    correct = 0
    total = 0
    net.to(device)
    with torch.no_grad():
        for data in testloader:
            images, labels = data[0].to(device), data[1].to(device)
            outputs = net(images)
            _, predicted = torch.max(outputs.data, 1)
            total += labels.size(0)
            correct += (predicted == labels).sum().item()

    print('Accuracy: %d %%' % (100 * correct / total))

# 主函数中调整优化器和损失函数
if __name__ == "__main__":
    net2 = FCNN(num_layers=2)
    net2.to(device)

    criterion = nn.CrossEntropyLoss().to(device)
    optimizer2 = optim.SGD(net2.parameters(), lr=0.01, momentum=0.9)

    print("Training network with 2 layers...")
    train(net2, criterion, optimizer2)
    test(net2)

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