How to assign a name for a pytorch layer?
Following a previous question, I want to plot weights, biases, activations and gradients to achieve a similar result to this.
Using
for name, param in model.named_parameters():
    summary_writer.add_histogram(f'{name}.grad', param.grad, step_index)

as was suggested in the previous question gives sub-optimal results, since layer names come out similar to '_decoder._decoder.4.weight', which is hard to follow, especially since the architecture is changing due to research. 4 in this run won't be the same in the next, and is really meaningless.
Thus, I wanted to give my own string names to each layer.

I found this Pytorch forum discussion, but no single best practice was agreed upon.
What is the recommended way to assign names to Pytorch layers?
Namely, layers defined in various ways:

Sequential:

self._seq = nn.Sequential(nn.Linear(1, 2), nn.Linear(3, 4),)


Dynamic:

self._dynamic = nn.ModuleList()
    for _ in range(self._n_features): 
        self._last_layer.append(nn.Conv1d(in_channels=5, out_channels=6, kernel_size=3, stride=1, padding=1,),)


Direct:

self._direct = nn.Linear(7, 8)


Other ways I didn't think about


I would like to be able to give a string name to each layer, defined in each of the above ways.

Sequential
Pass an instance of collections.OrderedDict. Code below gives conv1.weights, conv1.bias, conv2.weight, conv2.bias (notice lack of torch.nn.ReLU(), see end of this answer).
import collections

import torch

model = torch.nn.Sequential(
    collections.OrderedDict(
        [
            (&quot;conv1&quot;, torch.nn.Conv2d(1, 20, 5)),
            (&quot;relu1&quot;, torch.nn.ReLU()),
            (&quot;conv2&quot;, torch.nn.Conv2d(20, 64, 5)),
            (&quot;relu2&quot;, torch.nn.ReLU()),
        ]
    )
)

for name, param in model.named_parameters():
    print(name)

Dynamic
Use ModuleDict instead of ModuleList:
class MyModule(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.whatever = torch.nn.ModuleDict(
            {f&quot;my_name{i}&quot;: torch.nn.Conv2d(10, 10, 3) for i in range(5)}
        )

Will give us whatever.my_name{i}.weight (or bias) for each created module dynamically.
Direct
Just name it however you want and that's how it will be named
self.my_name_or_whatever = nn.Linear(7, 8)

You didn't think about

If you want to plot weights, biases and their gradients you can go along this route
You can't plot activations this way (or output from activations). Use PyTorch hooks instead (if you want per-layer gradients as they pass through network use this also)

For last task you can use third party library torchfunc (disclaimer: I'm the author) or go directly and write your own hooks.
