监控显存使用

PyTorch提供了两个函数用于显存查询:

memory_allocated

torch.cuda.memory_allocated(device=None)

查询指定GPU中使用的显存大小(字节)

max_memory_allocated

torch.cuda.max_memory_allocated(device=None)

返回给定设备的张量占用的最大GPU内存(字节)