PyTorch中的dataloader环节内存爆炸问题

最新推荐文章于 2025-05-13 01:55:56 发布

Yore_999

最新推荐文章于 2025-05-13 01:55:56 发布

阅读量2.5k

点赞数

文章标签： pytorch 深度学习机器学习

本文链接：https://blog.csdn.net/qq_35284646/article/details/130495367

版权

最近遇到在数据加载完成后切片导致数据占据内存过大，最后爆掉无法正常运行的情况，

经过仔细排查发现是由于切片后数据量不断增加，最终超过上限所致。

经过仔细寻找，最后发现通过引入pin_memory这一参数可以解决这一问题。

ataLoader(dataset, batch_size=1, shuffle=False, sampler=None,
           batch_sampler=None, num_workers=0, collate_fn=None,
           pin_memory=False, drop_last=False, timeout=0,
           worker_init_fn=None, *, prefetch_factor=2,
           persistent_workers=False)

可以看到，默认情况下pin_memory是false,这里需要手动设置为True。

这里的原理是：当我们设置pin_memory=True时，PyTorch会预先将数据从CPU内存复制到固定（钉住）的内存，然后再将其移动到GPU。这可以减少在训练过程中将数据从CPU复制到GPU所需的时间，从而提高性能。

这可以解决我们内存占用过多的情况。

但是，这种方法会使用额外的共享内存（shm）资源，也会带来问题，待续。