pytorch的Gpu内存使用率计算

Question

我有一个在服务器上运行的API服务，该服务器通过pytorch和GPU使用机器学习。如何计算一次调用API的GPU内存使用量？

Answer 1

您可以使用

torch.cuda.memory_summary

Answer 2

找到了这个解决方案！

# setting device on GPU if available, else CPU
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print('Using device:', device)
print()

#Additional Info when using cuda
if device.type == 'cuda':
    print(torch.cuda.get_device_name(0))
    print('Memory Usage:')
    print('Allocated:', round(torch.cuda.memory_allocated(0)/1024**3,1), 'GB')
    print('Cached:   ', round(torch.cuda.memory_cached(0)/1024**3,1), 'GB')