Register_forward_hook with global variable

1

Hello,
I have memory issues with forward_hooks.

So basically what I’m doing is the following:

loading pretrained model

vgg16 = models.vgg16(pretrained=True)
…

register forward hook

vgg16.features[3].register_forward_hook(get_layer3)
…

the function which is called in layer3 with my hook

def get_layer3(self, input, output):
global output_layer3
output_layer3 = output.data # saving the output.data from this layer

This is my network (simplified)

class Net(nn.Module):

def __init__(self, vgg):
	super(Net, self).__init__()
            self.features = vgg.features
            [...]
def forward(self, x):
	x = self.features(x)
            # do something with output_layer3
            [...]

So each time my model goes through layer3 of my vgg16-network it gives me the output of this layer and saves it in a global variable so that I can do stuff with it.

The problem is now, that this output_layer3 is saved in the GPU. And in the next epoch the new output_layer3 is been written to the GPU (without deleting the former output_layer3 or overwriting it) as well so that after several epochs the memory is full and the program crashes.

How can I program it in a way that it overwrites the previous output_layer3 with the next epoch?

I tried to define the variable inside mit network and set the forward_hook there. But I can’t define the function which is called inside the class. It tells me that it doesn’t know the function.

Do you have an idea what I can do here?

2

I found the problem. It was in upsampling of the output from the layers:
def get_layer8(self, input, output):
global output_layer8
upsampling_layer8 = nn.Upsample(scale_factor=2, mode=‘nearest’)
output_layer8 = upsampling_layer8(output.data)
output_layer8 = output_layer8.data

When I don’t upsample the tensor it works fine. I just have to manually get the correct pixels from the tensor which gets smaller after each layer.