内存使用情况

import os
import psutildef 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

整理自别的博客，侵删。