torch 测试模型显存占用

要测试PyTorch模型的显存占用，可以使用以下方法。

首先，需要导入PyTorch和相关类库。然后，创建一个模型并加载到设备上，例如使用以下代码加载一个预训练的ResNet模型：

import torch
import torchvision.models as models

# 创建模型并将其加载到设备上
model = models.resnet50()
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model.to(device)

接下来，可以生成一些测试数据并传递给模型进行前向传播。为了测试显存占用，可以使用`torch.cuda.max_memory_allocated()`和`torch.cuda.max_memory_cached()`方法来获取模型在GPU上分配的最大内存使用情况。

# 生成测试数据
inputs = torch.randn(1, 3, 224, 224).to(device)

# 前向传播
outputs = model(inputs)

# 打印最大的分配内存和缓存内存
print(f"Max allocated memory: {torch.cuda.max_memory_allocated() / (1024**3):.2f} GB")
print(f"Max cached memory: {torch.cuda.max_memory_cached() / (1024**3):.2f} GB")

运行以上代码后，将会输出模型在GPU上的最大内存使用情况。

注意，为了准确地测试显存占用，需要通过`torch.cuda.empty_cache()`方法清理缓存，并在每个循环迭代或训练批次之后调用它。

# 清理缓存
torch.cuda.empty_cache()

这是一种简单的方法来测试PyTorch模型的显存占用。