Укладка / конкатенация тензоров увеличивает использование памяти - PullRequest
Новая
       2

Укладка / конкатенация тензоров увеличивает использование памяти

1 голос
/ 27 мая 2020

В PyTorch я обнаружил, что несколько тензоров stack -ing или cat -ing увеличивают использование памяти на сумму размеров всех этих массивов. Пример выглядит следующим образом: 

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('=======')

И мы можем видеть, что каждое появление cat в for l oop увеличивает использование памяти на 512 МБ, что является суммой y и x1 (по 256 МБ). Сейчас это не является для меня серьезной проблемой, но мне это просто интересно. Если я правильно понимаю, векторно-якобиевское произведение стека или конкатенации - это просто выполнение обратной операции с вектором градиента, разделение его обратно на несколько массивов, формы которых соответствуют исходным входам для функции стека или кошки. Этому процессу не требуется промежуточное значение сложенного тензора, вычисленное в прямом проходе, что похоже на линейную операцию, но последняя не требует дополнительного использования памяти (например, если я изменяю cat в for l oop в линейную операцию, такую ​​как y = y * 1.1, это не увеличит потребление памяти). Поэтому мне интересно, является ли увеличенное использование памяти, по сути, просто пустым пространством, выделенным для хранения `` разделенных '' массивов градиентов в обратном проходе, который должен быть непрерывным, поэтому операция разделения не должна проходить в памяти ? Отсюда следует, что линейная операция не требует этой дополнительной памяти, потому что пространство памяти входного градиента можно напрямую перезаписать, не беспокоясь о смежности памяти. Это правда?

...