Extracting feature maps from ResNet
TLDR - what is considered best practice when extracting feature maps from ResNet?

I'm trying to feed the entire CIFAR10 dataset through ResNet18, to extract a new dataset that consists of some non-output activation of every sample in CIFAR10. I have implemented a code that generates this dataset, but the running time takes too long (exceeds Google Colab free RAM access, which is quite some RAM). The code I've implemented is based on a blog post called Intermediate Activations — the forward hook.
activation = {}

def get_activation(name):
    &quot;&quot;&quot;
    when given as input to register_forward_hook, this function is implicitly called when model.forward() is performed
    and saves the output of layer 'name' in the dictionary described above.
    :param name:
    :return:
    &quot;&quot;&quot;

    def hook(model, input, output):
        activation[name] = output.detach()

    return hook

the get_activation helper function is used inside the activation_maps function which takes the feature map provided from the 4th layer, 2nd BasicBlock,  conv1 layer (batch-size, 3,224,224) -&gt; (batch-size,512,7,7) of ResNet18
(PS - this layer was arbitrarily chosen - is there a known layer from which the activations are better?)
ResNet18 = torch.hub.load('pytorch/vision:v0.10.0', 'resnet18', pretrained=True)

def activation_maps(name='conv1'):
    &quot;&quot;&quot;
    This function takes a batch and returns some non - last activation alongside the true labels
    :return: train_activations_and_true_labels: array of tuples (Activation,True_labels) as train data
               
    &quot;&quot;&quot;
    non_output_activation_map = ResNet18.layer4[1].register_forward_hook(get_activation(name))
    # now we create a list of activations and true labels  for every sample.
    # This means that if we looped over (X,y) in a dataloader, we can now loop (activation,y) which is
    # an element in the arrays below, like a regular dataloader.
    train_activations_and_true_labels = []
    for i, (X_train, y_train) in enumerate(train_dataloader):
        out = ResNet18(X_train)
        train_activations_and_true_labels.append((activation[name], y_train))
        print(f&quot;Training data [{i}/{len(train_dataloader)}]&quot;, end='\r')
    non_output_activation_map.remove()  # detaching hooks
    return train_activations_and_true_labels

Now, this code runs - but exceeds the memory capacity of my PyCharm/Google-Colab Am I missing something? what is the best approach when extracting feature maps?

What batch size are you using, and how much RAM do you have available? Resnet is a somewhat large model, and the layer you're extracting is quite large as well so storing all that in memory might be causing issues.
Try reducing your batch size, or storing intermediary results to disk and clearing them from memory.
You might also consider turning off the gradient computation when calling the ResNet18 model, this would save a good bit of memory. Putting the @torch.no_grad() decorator on activation_maps(name='conv1') might work.
