如何为 pytorch 层分配名称？

Question

在上一个问题之后，我想 plot 权重、偏差、激活和梯度来实现与此类似的结果。

使用

for name, param in model.named_parameters():
    summary_writer.add_histogram(f'{name}.grad', param.grad, step_index)

正如前一个问题中所建议的那样，给出了次优结果，因为层名称类似于'_decoder._decoder.4.weight' ，这很难理解，特别是因为架构正在因研究而发生变化。 这次跑的4下一次就不一样了，真的没有意义。

因此，我想给每一层我自己的字符串名称。

---

我发现了这个Pytorch 论坛讨论，但没有就单一的最佳实践达成一致。

为 Pytorch 层分配名称的推荐方法是什么？

即，以各种方式定义的层：

1. 顺序：

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

2. 动态的：

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,),)

3. 直接的：

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

4. 其他我没想到的方式

---

我希望能够为每个层提供一个字符串名称，以上述每种方式定义。

Answer 1

顺序的

传递collections.OrderedDict的实例。 下面的代码给出了conv1.weights ， conv1.bias ， conv2.weight ， conv2.bias （注意缺少torch.nn.ReLU() ，请参阅此答案的结尾）。

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)

动态的

使用ModuleDict而不是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)}
        )

将为每个创建的模块动态地给我们whatever.my_name{i}.weight （或bias ）。

直接的

随心所欲地命名它，这就是它的命名方式

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

你没想到

• 如果您想 plot 权重、偏差及其梯度，您可以沿着这条路线使用 go
• 您不能以这种方式激活plot （或激活 output）。 改用PyTorch 钩子（如果你想要每层梯度，因为它们通过网络也使用这个）

对于最后一个任务，您可以直接使用第三方库torchfunc （免责声明：我是作者）或 go 并编写自己的钩子。