How to obtain feature map generated from first convolution layer of YOLOv5?

Question

I would like to obtain feature map generated from first convolution layer of YOLOv5 (PyTorch). How can I get the feature map?

Answer 1

It is trivial to modify the network structure to save an intermediate output. Here I save the output from the first conv and pooling layer as x1, and return this value as well in the forward call. Now if you try to use this approach in an off-the-shelf manner (i.e. if you didn't write your own training script and loss function) this will likely cause an error during training as the loss function call probably does not expect an additional output from forward. Thus, the preferred solution for drop-in modifications to existing nn libraries is to use the register_hook approach linked in the comments.

Here's a trivial example network from pytorch.

import torch.nn as nn
import torch.nn.functional as F


class Net(nn.Module):
  def __init__(self):
    super().__init__()
    self.conv1 = nn.Conv2d(3, 6, 5)
    self.pool = nn.MaxPool2d(2, 2)
    self.conv2 = nn.Conv2d(6, 16, 5)
    self.fc1 = nn.Linear(16 * 5 * 5, 120)
    self.fc2 = nn.Linear(120, 84)
    self.fc3 = nn.Linear(84, 10)

  def forward(self, x):
    x1 = self.pool(F.relu(self.conv1(x)))
    x = self.pool(F.relu(self.conv2(x1)))
    x = torch.flatten(x, 1) # flatten all dimensions except batch
    x = F.relu(self.fc1(x))
    x = F.relu(self.fc2(x))
    x = self.fc3(x)
    return x,x1


net = Net()