9

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.

Gulzar
  • 23,452
  • 27
  • 113
  • 201
  • For `nn.Sequential` you can provide an _OrderedDict_ of `nn.Modules` thus naming them – Shai Feb 11 '21 at 10:51

1 Answers1

13

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.

Szymon Maszke
  • 22,747
  • 4
  • 43
  • 83
  • Thanks! this is what I wanted. Can you please explain how `torchfunc` relates to this? What will I be able to di with torchfunc? For example, How would I plot activations and gradients with it? – Gulzar Feb 12 '21 at 13:46
  • @Gulzar it automatically registers recorders (via hooks) which get data coming out of specified layers (which you can later plot), see [documentation](https://szymonmaszke.github.io/torchfunc/packages/torchfunc.hooks/torchfunc.hooks.recorders.html). Also [`flashtorch`](https://github.com/MisaOgura/flashtorch) might help, though I don't have any experience with it. – Szymon Maszke Feb 12 '21 at 15:41
  • Ok now trying to implement this, I see your point about activations. Can you explain please why I can't plot activations like this? Also please note I am using Pytorch Lightning, maybe the hook is already implemented. Which one do I need? `on_after_backward` maybe? – Gulzar Feb 13 '21 at 07:58
  • Also, I see keeping these statistics slows training X100. What is the best practice of trading off logging with speed in this case? – Gulzar Feb 13 '21 at 08:03
  • @Gulzar please avoid multiquestion, post a new one about activations and make it more concrete, thanks. – Szymon Maszke Feb 13 '21 at 19:25