Representation projection to 2D plot

Jingles: <p>Let’s say I have a model, <a href="https://github.com/pytorch/vision/blob/master/torchvision/models/alexnet.py" rel="noopener nofollow ugc">AlexNet</a>:</p>
<pre data-code-wrap="python"><code class="lang-python">class AlexNet(nn.Module):

    def __init__(self, num_classes: int = 1000) -&gt; None:
        super(AlexNet, self).__init__()
        self.features = nn.Sequential(
            nn.Conv2d(3, 64, kernel_size=11, stride=4, padding=2),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(kernel_size=3, stride=2),
            nn.Conv2d(64, 192, kernel_size=5, padding=2),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(kernel_size=3, stride=2),
            nn.Conv2d(192, 384, kernel_size=3, padding=1),
            nn.ReLU(inplace=True),
            nn.Conv2d(384, 256, kernel_size=3, padding=1),
            nn.ReLU(inplace=True),
            nn.Conv2d(256, 256, kernel_size=3, padding=1),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(kernel_size=3, stride=2),
        )
        self.avgpool = nn.AdaptiveAvgPool2d((6, 6))
        self.classifier = nn.Sequential(
            nn.Dropout(),
            nn.Linear(256 * 6 * 6, 4096),
            nn.ReLU(inplace=True),
            nn.Dropout(),
            nn.Linear(4096, 4096),
            nn.ReLU(inplace=True),
            nn.Linear(4096, num_classes),
        )

    def forward(self, x: torch.Tensor) -&gt; torch.Tensor:
        x = self.features(x)
        ### I want to extract `x` and plot the representation in a 2D space
        x = self.avgpool(x)
        x = torch.flatten(x, 1)
        x = self.classifier(x)
        return x
</code></pre>
<p>Is there any way to visualize the weights learn from <code>self.features</code> layer and plot into a heatmap that look like this (taken from <a href="https://umap-learn.readthedocs.io/en/latest/basic_usage.html#digits-data" rel="noopener nofollow ugc">umap</a>):</p>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/2/4/241172916c47497e7f3e0ef377fc62a406b397b4.jpeg" data-download-href="https://discuss.pytorch.org/uploads/default/241172916c47497e7f3e0ef377fc62a406b397b4" title="image"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/2/4/241172916c47497e7f3e0ef377fc62a406b397b4_2_618x500.jpeg" alt="image" data-base62-sha1="594zlsMx9dgJ0B5SufVw0dVTQry" width="618" height="500" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/2/4/241172916c47497e7f3e0ef377fc62a406b397b4_2_618x500.jpeg, https://discuss.pytorch.org/uploads/default/optimized/3X/2/4/241172916c47497e7f3e0ef377fc62a406b397b4_2_927x750.jpeg 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/2/4/241172916c47497e7f3e0ef377fc62a406b397b4_2_1236x1000.jpeg 2x" data-dominant-color="F9F8F7"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">image</span><span class="informations">1356×1096 72.4 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>
<p>Where each color is the class labels of the data points from the <code>test</code> dataloader. Preferably with Python packages and plot with matplotlib and without tensorboard.</p>

ptrblck: <p>Yes, you could use forward hooks to store the activation from a specific layer (I assume you meant activation instead of weight) as described <a href="https://discuss.pytorch.org/t/how-can-l-load-my-best-model-as-a-feature-extractor-evaluator/17254/6">here</a> and use the linked <code>umap</code> library to reduce the dimensionality. Once this is done, you could use <code>matplotlib</code> to create a scatter plot (take a look at the examples at the bottom of the <a href="https://matplotlib.org/3.3.3/api/_as_gen/matplotlib.pyplot.scatter.html">docs</a>).</p>