How to assign a name for a pytorch layer?

user_913098 0
Viewed 4584

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:

  1. Sequential:
self._seq = nn.Sequential(nn.Linear(1, 2), nn.Linear(3, 4),)
  1. 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,),)
  1. Direct:
self._direct = nn.Linear(7, 8)
  1. 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.

pytorch deep-learning machine-learning neural-network python
1 Answers

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(
        [
            ("conv1", torch.nn.Conv2d(1, 20, 5)),
            ("relu1", torch.nn.ReLU()),
            ("conv2", torch.nn.Conv2d(20, 64, 5)),
            ("relu2", 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"my_name{i}": 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.

user_108864200
Related