把pytorch变量转移到cpu中释放gpu内存

把pytorch变量转移到cpu中释放gpu内存,第1张

背景

我想要对一个大批量的数据分批输入模型得到所有的embedding,但是过程进行到中间报 cuda out of memory。经查勘是因为早些的embedding没有及时从cuda中释放内存导致后面的model无法进行forward。

解决

只需要pytorch的tensor.cpu()就可以释放了。
举个例子:

embeddings=[]
for batch in data_loader:
	embedding_batch=model(input_data)
	embeddings.append(embedding_batch.cpu()) # 不带cpu不能释放cuda内存,导致oom

若不加cpu(),会进行20多batch后cuda out of memory。

后记

其实接触torch tensor已经四年有余,往往对其cpu()是因为torch与numpy数据不符合,往往和numpy()一起连用,本以为是对原tensor返回一个cpu device上的tensor,不改变原变量,没想到cpu()竟然是对原变量直接坐内存的变化,出乎意料。

欢迎分享,转载请注明来源:内存溢出

原文地址: http://outofmemory.cn/langs/943835.html

(0)
打赏 微信扫一扫 微信扫一扫 支付宝扫一扫 支付宝扫一扫
上一篇 2022-05-18
下一篇 2022-05-18

发表评论

登录后才能评论

评论列表(0条)

保存