有时候用pytorch跑模型很容易出现内存溢出的情况，因此要具体查看内存情况： 简单版：
#直接调用返回当前可用内存/显存 import torch import pynvml import psutil def get_avaliable_memory(device): if device==torch.device('cuda:0'): pynvml.nvmlInit() handle = pynvml.nvmlDeviceGetHandleByIndex(0) # 0表示第一块显卡 meminfo = pynvml.nvmlDeviceGetMemoryInfo(handle) ava_mem=round(meminfo.free/1024**2) print('current available video memory is' +' : '+ str(round(meminfo.free/1024**2)) +' MIB') elif device==torch.device('cpu'): mem = psutil.virtual_memory() print('current available memory is' +' : '+ str(round(mem.used/1024**2)) +' MIB') ava_mem=round(mem.used/1024**2) return ava_mem复杂版，有一个github开源项目： https://github.com/Oldpan/Pytorch-Memory-Utils 主要代码如下：
import gc import datetime import pynvml import torch import numpy as np class MemTracker(object): """ Class used to track pytorch memory usage Arguments: frame: a frame to detect current py-file runtime detail(bool, default True): whether the function shows the detail gpu memory usage path(str): where to save log file verbose(bool, default False): whether show the trivial exception device(int): GPU number, default is 0 """ def __init__(self, frame, detail=True, path='', verbose=False, device=0): self.frame = frame self.print_detail = detail self.last_tensor_sizes = set() self.gpu_profile_fn = path + f'{datetime.datetime.now():%d-%b-%y-%H:%M:%S}-gpu_mem_track.txt' self.verbose = verbose self.begin = True self.device = device self.func_name = frame.f_code.co_name self.filename = frame.f_globals["__file__"] if (self.filename.endswith(".pyc") or self.filename.endswith(".pyo")): self.filename = self.filename[:-1] self.module_name = self.frame.f_globals["__name__"] self.curr_line = self.frame.f_lineno def get_tensors(self): for obj in gc.get_objects(): try: if torch.is_tensor(obj) or (hasattr(obj, 'data') and torch.is_tensor(obj.data)): tensor = obj else: continue if tensor.is_cuda: yield tensor except Exception as e: if self.verbose: print('A trivial exception occured: {}'.format(e)) def track(self): """ Track the GPU memory usage """ pynvml.nvmlInit() handle = pynvml.nvmlDeviceGetHandleByIndex(self.device) meminfo = pynvml.nvmlDeviceGetMemoryInfo(handle) self.curr_line = self.frame.f_lineno where_str = self.module_name + ' ' + self.func_name + ':' + ' line ' + str(self.curr_line) with open(self.gpu_profile_fn, 'a+') as f: if self.begin: f.write(f"GPU Memory Track | {datetime.datetime.now():%d-%b-%y-%H:%M:%S} |" f" Total Used Memory:{meminfo.used/1000**2:<7.1f}Mb\n\n") self.begin = False if self.print_detail is True: ts_list = [tensor.size() for tensor in self.get_tensors()] new_tensor_sizes = {(type(x), tuple(x.size()), ts_list.count(x.size()), np.prod(np.array(x.size()))*4/1000**2) for x in self.get_tensors()} for t, s, n, m in new_tensor_sizes - self.last_tensor_sizes: f.write(f'+ | {str(n)} * Size:{str(s):<20} | Memory: {str(m*n)[:6]} M | {str(t):<20}\n') for t, s, n, m in self.last_tensor_sizes - new_tensor_sizes: f.write(f'- | {str(n)} * Size:{str(s):<20} | Memory: {str(m*n)[:6]} M | {str(t):<20} \n') self.last_tensor_sizes = new_tensor_sizes f.write(f"\nAt {where_str:<50}" f"Total Used Memory:{meminfo.used/1000**2:<7.1f}Mb\n\n") pynvml.nvmlShutdown()使用方法： 定义好之后，有序的放到各个模块之间，就可以生成一个txt文件，显示显存的使用情况
mport torch import inspect from torchvision import models from gpu_mem_track import MemTracker device = torch.device('cuda:0') frame = inspect.currentframe() # define a frame to track gpu_tracker = MemTracker(frame) # define a GPU tracker gpu_tracker.track() # run function between the code line where uses GPU cnn = models.vgg19(pretrained=True).features.to(device).eval() gpu_tracker.track() # run function between the code line where uses GPU dummy_tensor_1 = torch.randn(30, 3, 512, 512).float().to(device) # 30*3*512*512*4/1000/1000 = 94.37M dummy_tensor_2 = torch.randn(40, 3, 512, 512).float().to(device) # 40*3*512*512*4/1000/1000 = 125.82M dummy_tensor_3 = torch.randn(60, 3, 512, 512).float().to(device) # 60*3*512*512*4/1000/1000 = 188.74M gpu_tracker.track()