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

背景

我想要对一个大批量的数据分批输出模型失去所有的 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()居然是对原变量间接坐内存的变动，出其不意。