Утечка памяти в Python коде при манипулировании весами в модели CNN - PullRequest
Новая
       73

Утечка памяти в Python коде при манипулировании весами в модели CNN

0 голосов
/ 18 марта 2020

Я новичок в Python, и я пытаюсь реализовать сокращение каналов в Resnet50, что в основном означает, что моя программа должна уменьшить размеры тензоров веса для каждого слоя в сети. Вход представляет собой 4-мерный тензор, а выходной должен быть тензор, размер которого уменьшается, при этом удаляются все каналы, индексы которых отсутствуют в списке «idx». 

# this code is in a function that takes *blocks* and *amount* as arguments
# *blocks* is a nested dictionary containing the layers in each bottleneck in resnet50
# the output is *reduced_blocks* which is also a nested-dictionary
reduced_blocks={}

for block in blocks:
    bottleneck={}
    for layer in blocks[block]:
        #store dimentions of the layer to be pruned in this case, its conv3 of each bottleneck in resnet50 
        (c_,h_,w_,d_)=blocks[block][layer].conv3.weight.shape
        #store the content of the weight tensor of layer conv3
        layer3 = blocks[block][layer].conv3.weight
        # a new 4-dimensional tensor that has the reduced dimensions by *amount* 
        w = torch.zeros(c_,int(h_*(1-amount)),w_,d_)

        w_plane=torch.zeros(int(h_*(1-amount)),w_,d_)
        for i in range(c_):
            m=0
            for j in idx:
                # *idx* contains only the indexes of channels that i want to keep
                # this is where the problem appears to be, if I comment this line, the code runs without memory saturation
                w_plane[m]=layer3[i][j]
                m+=1
            # I assign the pruned plane to the *w* tensor
            w[i]=w_plane
            dic = {'conv3':w }
            bottleneck.update({layer:dic})
        reduced_blocks[block] = bottleneck.copy()
        bottleneck.clear()
 # the function then returns *bottleneck_pruned*

Всякий раз, когда я запускаю код в лаборатории Jupyter использование памяти начинает расти, пока не осталось ничего. Машина либо зависает, либо Jupyter перезапускает ядро. Я не уверен, почему я получаю эту ошибку, может потому, что у меня слишком много вложенных циклов for? Потому что, если я применяю сокращение только к одному слою, программа работает нормально. Я попытался определить, где именно проблема, кажется. Я прокомментировал программу построчно, пока не обнаружил, что виновником было назначение в последнем l oop. Если кто-то поймет, почему у меня возникла эта проблема, ваша помощь будет принята с благодарностью.

...