pytorch 训练过程内存泄露/显存泄露debug记录:dataloader和dataset导致的泄露

这篇具有很好参考价值的文章主要介绍了pytorch 训练过程内存泄露/显存泄露debug记录:dataloader和dataset导致的泄露。希望对大家有所帮助。如果存在错误或未考虑完全的地方,请大家不吝赐教,您也可以点击"举报违法"按钮提交疑问。

背景

微调 mask-rcnn 代码,用的是 torchvision.models.detection.maskrcnn_resnet50_fpn 代码,根据该代码的注释,输入应该是:
images, targets=None
(List[Tensor], Optional[List[Dict[str, Tensor]]]) -> Tuple[Dict[str, Tensor], List[Dict[str, Tensor]]]

所以我写的 dataset 是这样的:

def _load_ann(self):
	transformed_anns = {'boxes': boxes : List[List],
						'labels': categories: List[int],
						'masks': masks: List[str]}
	self.anns[filename] = transformed_anns

def __getitem__(self, item) -> (Tensor, Optional[Dict[str, Tensor]]):
	img_name = self.img_list[item]
	img = cv2.imread(os.path.join(self.root_path, self.split, img_name))
	if self.split == 'train':
		return img, self.anns[img_name]

大概思路是:先把所有的标注信息读入内存,然后按照 img_name 把标注信息(也就是 target )取出来

这里有个令人纠结的地方:__get_item__ 返回的到底是什么格式的数据?一开始我是直接把 boxes labels masks 都直接写成 tensor 返回,但是一次性把所有的 masks 都读到内存,太大了。再加上看了这个 pytorch内存泄露-dataloader - 知乎 ,这里建议 dataset 的 __get_item__ 返回的都是 python 的基础数据类型,所以我就改成了上面的样子。其实返回什么类型的都行,只要在 dataloader 的 collate_fn 方法里面都转成可以送入模型的数据形式就行了。

因为 dataset 是上面的写法,所以对应的 collate_fn 写法是:

def collate_fn(datas: List[Tuple[Tensor, Dict]]):
	imgs = []
	targets = []
	for data in datas:
		img, target = data
		imgs.append(transforms.ToTensor()(img))

		target['boxes'] = torch.tensor(target['boxes'], dtype=torch.float)
		target['labels'] = torch.tensor(target['labels'], dtype=torch.int64)

		masks = target['masks']
		masks = [cv2.imread(mask, 0) for mask in masks]
		masks = np.stack(masks, axis=0)
		masks = masks / 255
		masks = masks.astype(np.uint8)
		target['masks'] = torch.from_numpy(masks)

		targets.append(target)
	return imgs, targets

错误排查及解决方法

把所有的数据送入 model 的代码都注释掉,只保留如下代码:

for e in range(epoch):
	for i, (imgs, targets) in enumerate(train_dataloader):
		imgs = [img.to(device) for img in imgs]
		targets = [_to_device(target, device) for target in targets]

watch -n 1 nvidia-smi 监控显存占用,发现一直在涨。毫无疑问肯定是 dataloader 导致的显存泄露 😭

然后就是排查,到底是谁?是谁想害朕??

排查方法是:分别注释掉 imgs / boxes / labels / masks ,观察注释掉谁的时候不会显存泄露。

发现,是 masks 导致的内存泄露。

但是这很怪啊,明明 masks 和 imgs 是一样的数据类型,为什么前者会显存泄露,但是后者不会?于是我把 masks 单独拿出来,像 imgs 一样放在 list 里面,不会内存泄露。但是一旦把 imgs 嵌套放在 targets 这个 dict 里面,就会显存泄露 orz

于是,既然是 masks 没有释放,所以我加一句:

for e in range(epoch):
	for i, (imgs, targets) in enumerate(train_dataloader):
		imgs = [img.to(device) for img in imgs]
		targets = [_to_device(target, device) for target in targets]
		# ... 传入模型的计算
		for target in targets:
			del target['masks']

但是没用,还是泄露。然后查了 pytorch 怎么释放 tensor,发现要主动调用 torch.cuda.empty_cache() 才会释放,所以我又加了一句:

for e in range(epoch):
	for i, (imgs, targets) in enumerate(train_dataloader):
		imgs = [img.to(device) for img in imgs]
		targets = [_to_device(target, device) for target in targets]
		# ... 传入模型的计算
		for target in targets:
			del target['masks']
		torch.cuda.empty_cache()

这回没有显存泄露了。

但是出现了新的问题,在 epoch=2 的时候报错 targets 没有 masks 这个 key;然后我 debug 发现,由 dataloader 取到的数据 label 和 boxes 在 collate_fn 之前就已经是 tensor 状态了,再往前倒,发现 dataset.anns 里面的数据居然被改了!这实在是太荒谬了。

所以我把 __get_item 改成:

def __getitem__(self, item) -> (Tensor, Optional[Dict[str, Tensor]]):
	img_name = self.img_list[item]
	img = cv2.imread(os.path.join(self.root_path, self.split, img_name))
	if self.split == 'train':
		return img, deepcopy(self.anns[img_name])

这样就没问题了文章来源地址https://www.toymoban.com/news/detail-630932.html

总结

  1. 查找内存泄露/显存泄露的位置:
    • 把数据送入模型的代码全部注释掉,观察显存是否上涨;上涨说明内存泄露出现在 dataloader(出现在非 dataloader 地方的最常见的显存泄露原因是,loss 打印/统计的时候没有写 loss.item()
    • 把不同的 data 组成部分注释掉,观察具体是哪个 data 导致的内存泄露
  2. pytorch 释放内存的方法:把 tensor 读到 gpu 就会有显存占用,一般可以自动释放,但是显存泄露的时候就没法释放。找到没有及时释放的代码位置,然后首先 del tensor 标记删除,随后需要调用 torch.cuda.empty_cache() 才能真正释放
  3. dataset 的 __get_item__ 方法注意,如果要返回内部维护的 list 类型的数据的话,不要直接返回该数据切片,而是返回 deepcopy() 防止内部维护的数据被外部修改

到了这里,关于pytorch 训练过程内存泄露/显存泄露debug记录:dataloader和dataset导致的泄露的文章就介绍完了。如果您还想了解更多内容,请在右上角搜索TOY模板网以前的文章或继续浏览下面的相关文章,希望大家以后多多支持TOY模板网!

本文来自互联网用户投稿,该文观点仅代表作者本人,不代表本站立场。本站仅提供信息存储空间服务,不拥有所有权,不承担相关法律责任。如若转载,请注明出处: 如若内容造成侵权/违法违规/事实不符,请点击违法举报进行投诉反馈,一经查实,立即删除!

领支付宝红包 赞助服务器费用

相关文章

  • 记录线上排查内存泄露问题

    记录一次云上排查内存泄露的问题,最近监控告警云上有空指针异常报出,于是找到运维查日志定位到具体是哪一行代码抛出的空指针异常, 发现是在解析cookie的一个方法内,调用HttpServletRequest.getServerName()获取不到抛出的NPE,这个获取服务名获取不到,平时都没有出现过的

    2024年02月03日
    浏览(39)
  • 记录一次发现Webpack源码泄露(js.map泄露)过程

    Vue使用webpack(静态资源打包器)的时候,如果未进行正确配置,会产生一个js.map文件,而这个js.map可以通过工具来反编译还原Vue源代码,产生代码泄露。 使用npm -v查看版本 npm -v 使用如下命令安装 npm install --global reverse-sourcemap 在目录扫描中,发现某个压缩包中存在js.map文件 解

    2024年03月20日
    浏览(56)
  • 一行代码解决PyTorch训练模型时突然出现的For debugging consider passing CUDA_LAUNCH_BLOCKING=1报错

    一、问题描述         今天在调试模型的代码,然后代码一直运行得好好地,就突然出现了一下的错误: RuntimeError: CUDA error: invalid device ordinal CUDA kernel errors might be asynchronously reported at some other API call,so the stacktrace below might be incorrect. For debugging consider passing CUDA_LAUNCH_BLOCK

    2024年02月02日
    浏览(46)
  • 人工智能(pytorch)搭建模型9-pytorch搭建一个ELMo模型,实现训练过程

    大家好,我是微学AI,今天给大家介绍一下人工智能(pytorch)搭建模型9-pytorch搭建一个ELMo模型,实现训练过程,本文将介绍如何使用PyTorch搭建ELMo模型,包括ELMo模型的原理、数据样例、模型训练、损失值和准确率的打印以及预测。文章将提供完整的代码实现。 ELMo模型简介 数据

    2024年02月07日
    浏览(67)
  • Pytorch运行过程中解决出现内存不足的问题

    利用Transformer模型进行O3浓度的反演 一开始模型是在CPU上面跑的,为了加快速度,我改成了在GPU上跑 方法如下: 1、验证pytorch是否存在GPU版本 在Pycharm命令行输入 我的输出为True,说明pytorch是存在GPU版本的 2、将模型从CPU版本转换到GPU版本 声明使用GPU(指定具体的卡) PS:

    2023年04月16日
    浏览(45)
  • Pytorch训练深度强化学习时CPU内存占用一直在快速增加

    最近在用MATD3算法解决多机器人任务,但是在训练过程中,CPU内存一直在增加(注意,不是GPU显存)。我很头疼,以为是算法代码出了问题,导致了内存泄漏,折腾了1天也没解决。后来用memory_profiler对代码分析,才发现是这个函数占用的内存一直在增加:  def store_transition(

    2024年02月07日
    浏览(51)
  • PyTorch深度学习实战(1)——神经网络与模型训练过程详解

    人工神经网络 ( Artificial Neural Network , ANN ) 是一种监督学习算法,其灵感来自人类大脑的运作方式。类似于人脑中神经元连接和激活的方式,神经网络接受输入,通过某些函数在网络中进行传递,导致某些后续神经元被激活,从而产生输出。函数越复杂,网络对于输入的数据拟

    2024年02月06日
    浏览(51)
  • 【深度学习】BasicSR训练过程记录,如何使用BasicSR训练GAN

    BasicSR支持两种灵活的使用场景,以满足用户不同的需求: 本地克隆仓库使用: 用户可以直接克隆BasicSR的本地仓库,查看完整的代码并进行修改,例如在BasicSR中训练SRGAN或StyleGAN2。安装方式包括先执行 git clone ,然后运行 python setup.py develop/install 。详细信息请参考安装指南。

    2024年01月23日
    浏览(47)
  • SwinIR实战:详细记录SwinIR的训练过程

    论文地址:https://arxiv.org/pdf/2108.10257.pdf 预训练模型下载:https://github.com/JingyunLiang/SwinIR/releases 训练代码下载:https://github.com/cszn/KAIR 测试代码:https://github.com/JingyunLiang/SwinIR 论文翻译:https://blog.csdn.net/hhhhhhhhhhwwwwwwwwww/article/details/124434886 测试:https://wanghao.blog.csdn.net/article/d

    2024年01月16日
    浏览(44)
  • nerf训练自己的数据,过程记录十分详细

           之前跑很多项目没有记录,后来再弄就不行了。这次特别记录一下,在梳理流程的同时希望给大家带来小小的帮助!我自己是在cuda11.2,windows环境下成功的,过程十分详细,有需要的朋友耐心看完。有问题可以评论区交流         首先,本文nerf是基于pytorch训练的,

    2024年02月03日
    浏览(47)

觉得文章有用就打赏一下文章作者

支付宝扫一扫打赏

博客赞助

微信扫一扫打赏

请作者喝杯咖啡吧~博客赞助

支付宝扫一扫领取红包,优惠每天领

二维码1

领取红包

二维码2

领红包