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