如何在 PyTorch 中检查我的模型在 GPU 上的天气

如果你正在使用 PyTorch 训练深度学习模型并且使用 GPU 来加速训练过程，那么你可能想要检查模型是否确实在 GPU 上运行，以及 GPU 的使用情况。本文将介绍如何在 PyTorch 中检查模型是否在 GPU 上运行，并显示 GPU 的使用情况。

检查模型是否在 GPU 上运行

在 PyTorch 中，默认情况下模型和数据张量（tensor）是在 CPU 上运行的。要将它们移到 GPU 上，你需要将它们转换为 torch.cuda module 中的张量。如果你已经这样做了，那么你可以使用 torch.Tensor.device 属性来检查张量所在的设备。

import torch

# 检查 GPU 是否可用
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

# 定义模型并将其移动到 GPU 上
model = MyModel()
model.to(device)

# 定义输入数据并将其移动到 GPU 上
input_data = torch.randn(batch_size, input_size).to(device)

# 运行模型
output = model(input_data)

# 检查输出的张量所在的设备
output_dev = output.device
print(f"output device: {output_dev}")  # output device: cuda:0

如上所示，我们首先检查 GPU 是否可用，然后将我们的模型和输入数据转换为 GPU 上的张量，最后运行模型并检查输出张量所在的设备。在这个例子中，输出张量的设备是 cuda:0，说明它在 GPU 上运行。

显示 GPU 的使用情况

要检查 GPU 的使用情况，我们可以使用 torch.cuda.memory_allocated() 来获取当前 GPU 上已分配的内存量（以字节为单位），使用 torch.cuda.max_memory_allocated() 来获取 GPU 在某个时刻分配的最大内存量，使用 torch.cuda.memory_cached() 来获取 GPU 上当前缓存的内存量（即可以立即重新分配的内存），以及使用 torch.cuda.max_memory_cached() 来获取 GPU 在某个时刻缓存的最大内存量。

我们可以在模型训练过程中打印这些值，以便查看模型在 GPU 上的内存使用情况。以下是一个例子：

import torch

# 检查 GPU 是否可用
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

# 定义模型并将其移动到 GPU 上
model = MyModel()
model.to(device)

# 定义损失函数和优化器
criterion = torch.nn.MSELoss()
optimizer = torch.optim.SGD(model.parameters(), lr=learning_rate)

# 模型训练
for epoch in range(num_epochs):

    # 对训练数据进行迭代
    for i, (inputs, labels) in enumerate(train_loader):

        # 将输入和标签数据移动到 GPU 上
        inputs = inputs.to(device)
        labels = labels.to(device)

        # 前向传播
        outputs = model(inputs)
        loss = criterion(outputs, labels)

        # 反向传播和优化
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        # 显示 GPU 的使用情况
        memory_allocated = torch.cuda.memory_allocated(device=device)
        max_memory_allocated = torch.cuda.max_memory_allocated(device=device)
        memory_cached = torch.cuda.memory_cached(device=device)
        max_memory_cached = torch.cuda.max_memory_cached(device=device)

        print(f"Epoch [{epoch+1}/{num_epochs}], Batch [{i+1}/{total_batches}], "
              f"Loss: {loss.item():.4f}, "
              f"Memory Allocated: {memory_allocated/1024**2:.2f} MB, "
              f"Max Memory Allocated: {max_memory_allocated/1024**2:.2f} MB, "
              f"Memory Cached: {memory_cached/1024**2:.2f} MB, "
              f"Max Memory Cached: {max_memory_cached/1024**2:.2f} MB")

如上所示，我们在每个训练批次结束时打印出当前的损失和 GPU 的使用情况。在这个例子中，我们完成了一个 epoch 之后打印 GPU 的使用情况，但是你可以将它们放在任何你认为合适的位置。

总结

在本教程中，我们介绍了如何在 PyTorch 中检查模型是否在 GPU 上运行，并显示 GPU 的使用情况。这些技巧对于调试和优化深度学习模型非常有用，因为它们可以帮助你了解模型的内存使用情况，避免内存泄露和优化 GPU 计算。