PyTorch深度学习代码获取内存、显存使用情况

内存使用情况

import os
import psutil

def get_current_memory_gb():
    # 获取当前进程内存占用。
    pid = os.getpid()
    p = psutil.Process(pid)
    info = p.memory_full_info()
    return info.rss / 1024. / 1024. / 1024.

显存情况

import torch
import pynvml

# Total memory of a GPU card (MB)
mem_card = torch.cuda.get_device_properties(self.device).total_memory/1024/1024

# Get free memory of a GPU card (MB)
# gpu_index = device.index
def get_memory_free_MiB(gpu_index):
    pynvml.nvmlInit()
    handle = pynvml.nvmlDeviceGetHandleByIndex(int(gpu_index))
    mem_info = pynvml.nvmlDeviceGetMemoryInfo(handle)
    return mem_info.free // 1024 ** 2

整理自别的博客，侵删。