Question

在PyTorch中，我发现对多个张量进行stack或cat进行处理将使所有这些数组的大小之和增加内存使用量。一个例子如下：

import torch as tc
import torch.autograd as tag
import sys
import psutil
import os
import resource

def get_ru_maxrss():
    """ Return max RSS usage (in kilobytes) """
    size = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss
    if sys.platform == 'darwin':
        # on Mac OS X ru_maxrss is in bytes, on Linux it is in KB
        size //= 1024
    return size / 1024

def cpuStats():
        print(sys.version)
        print(psutil.cpu_percent())
        print(psutil.virtual_memory())  # physical memory usage
        pid = os.getpid()
        py = psutil.Process(pid)
        memoryUse = py.memory_info()[0] / 2. ** 30  # memory use in GB...I think
        print('memory GB:', memoryUse)

m0 = get_ru_maxrss()

x1 = tc.ones([8192, 8192], requires_grad=True)
print(x1.dtype)
print(get_ru_maxrss() - m0)
print('=======')

y = x1 * 1.1
print(y.dtype)
print(get_ru_maxrss() - m0)
print('=======')

for i in range(10):
    y = tc.cat([y, x1])
    print(y.dtype)
    print(get_ru_maxrss() - m0)
    print('=======')

loss = tc.mean(y)

print(get_ru_maxrss() - m0)
print('=======')

loss.backward()

print(get_ru_maxrss() - m0)
print('=======')

我们可以看到，在for循环中每次出现cat都会使内存使用增加512 MB，这是y和x1的总和（每个256 MB）。目前这对我来说不是主要问题，但我对此感到很好奇。如果我正确理解的话，堆栈或串联的向量-雅可比积只是对梯度向量执行相反的操作，将其分成多个数组，其形状与堆栈或cat函数的原始输入相匹配。此过程不需要在前向遍历中计算的堆叠张量的中间值，这类似于线性运算，但是后者不会招致额外的内存使用（例如，如果我更改了cat for循环到诸如y = y * 1.1之类的线性运算，不会增加内存消耗）。所以我想知道增加的内存使用量是否实际上是分配的一个空白空间，以在向后传递中保留``拆分的''梯度数组，该空间必须是连续的，因此拆分操作不必大步跨入内存？因此，线性运算不需要此额外的内存，因为可以直接覆盖输入梯度的内存空间，而不必担心内存连续性。是吗？

堆叠/级联张量会增加内存使用量

0 个答案: