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

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

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

经过仔细寻找,最后发现通过引入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)资源,也会带来问题,待续。

你可能感兴趣的:(pytorch,深度学习,机器学习)