PyTorch 张量切片和内存使用

DBa*_*ker 2 python pytorch

import torch
T = torch.FloatTensor(range(0,10 ** 6)) # 1M

#case 1:
torch.save(T, 'junk.pt')
# results in a 4 MB file size

#case 2:
torch.save(T[-20:], 'junk2.pt')
# results in a 4 MB file size

#case 3:
torch.save(torch.FloatTensor(T[-20:]), 'junk3.pt')
# results in a 4 MB file size

#case 4:
torch.save(torch.FloatTensor(T[-20:].tolist()), 'junk4.pt')
# results in a 405 Bytes file size

Run Code Online (Sandbox Code Playgroud)

我的问题是:

  1. 在第 3 种情况下,生成的文件大小似乎令人惊讶,因为我们正在创建一个新的张量。为什么这个新张量不仅仅是切片?

  2. 情况 4 是仅保存张量的一部分(切片)的最佳方法吗?

  3. 更一般地说,如果我想通过删除它的前半部分值来“修剪”一个非常大的一维张量以节省内存,我是否必须像情况 4 那样继续,或者是否有更直接且计算量更少的方法不涉及创建 python 列表的昂贵方式。

Mic*_*ngo 5

(i) 在情况 3 中,由于我们正在创建一个新的张量,因此生成的文件大小似乎令人惊讶。为什么这个新张量不仅仅是切片?

切片创建张量的视图,该视图共享底层数据,但包含有关用于可见数据的内存偏移量的信息。这避免了必须频繁复制数据,从而使许多操作更加高效。有关受影响操作的列表,请参阅PyTorch - 张量视图

您正在处理少数几种情况之一,其中基础数据很重要。为了保存张量,它需要保存底层数据,否则偏移量将不再有效。

torch.FloatTensor如果没有必要,不会创建张量的副本。您可以验证它们的底层数据是否仍然相同(它们具有完全相同的内存位置):

torch.FloatTensor(T[-20:]).storage().data_ptr() == T.storage().data_ptr()
# => True
Run Code Online (Sandbox Code Playgroud)

(ii) 情况 4 是仅保存张量的一部分(切片)的最佳方法吗?

(iii) 更一般地说,如果我想通过删除其值的前半部分来“修剪”一个非常大的一维张量以节省内存,我是否必须像情况 4 那样继续,还是有更直接的方法?以及不涉及创建 python 列表的计算成本较低的方法。

你很可能不会复制切片的数据,但至少你可以避免从它创建一个 Python 列表并从列表中创建一个新的张量,torch.Tensor.clone而是使用:

torch.save(T[-20:].clone(), 'junk5.pt')
Run Code Online (Sandbox Code Playgroud)