如何为 pytorch 层指定名称？

Question

在a previous question 之后，我想绘制权重、偏差、激活和梯度以达到与this 相似的结果。

使用

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

正如the previous question 中所建议的那样，给出了次优结果，因为层名称与'_decoder._decoder.4.weight' 相似，这很难理解，尤其是由于研究正在改变架构。 4这一次跑，下一次就不一样了，真的没意义。

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

---

我发现this Pytorch 论坛讨论，但没有达成一致的最佳实践。

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

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

1. 顺序：

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

1. 动态：

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. 直接：

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)

你没想到

• 如果您想绘制权重、偏差及其梯度，您可以沿着这条路线走
• 您不能以这种方式绘制激活（或激活的输出）。请改用PyTorch hooks（如果您希望每层梯度在通过网络时也可以使用）

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