Pytorch：为什么`tensor`变量占用的内存这么小？

Question

在Pytorch 1.0.0中，我发现一个

tensor

变量占用的内存非常小。我想知道它是如何存储这么多数据的。 这是代码。

a = np.random.randn(1, 1, 128, 256)
b = torch.tensor(a, device=torch.device('cpu'))

a_size = sys.getsizeof(a)
b_size = sys.getsizeof(b)

a_size

是 262288。

b_size

是 72。

Answer 1

答案分为两部分。从

sys.getsizeof

的文档来看，首先

所有内置对象都将返回正确的结果，但这不一定适用于第三方扩展，因为它是特定于实现的。

因此，对于张量

__sizeof__

可能未定义或定义与您期望的不同 - 这个函数不是您可以依赖的。其次

仅考虑直接归因于该对象的内存消耗，而不考虑其引用的对象的内存消耗。

这意味着如果

torch.Tensor

对象仅保存对实际内存的引用，则不会在

sys.getsizeof

中显示。确实是这样，如果你检查底层存储的大小，你会看到预期的数字

import torch, sys
b = torch.randn(1, 1, 128, 256, dtype=torch.float64)
sys.getsizeof(b)
>> 72
sys.getsizeof(b.storage())
>> 262208

注意：我将

dtype

显式设置为

float64

，因为这是

dtype

中的默认

numpy

，而

torch

默认使用

float32

。

Answer 2

如果你想在cuda中获取张量或网络的大小，你可以使用此代码来计算它的大小：

import torch

device = 'cuda:0'

# before
torch._C._cuda_clearCublasWorkspaces()
memory_before = torch.cuda.memory_allocated(device)

# your tensor or network
data5 = torch.randn((10000,100),device=device)

# after
memory_after = torch.cuda.memory_allocated(device)
latent_size = memory_after - memory_before

latent_size
# 4000256

从中获取想法：https://github.com/pytorch/pytorch/blob/ee28b865ee9c87cce4db0011987baf8d125cc857/torch/distributed/pipeline/sync/_balance/profile.py#L102