如何为 pytorch 层指定名称？

Posted 2023-03-27

【中文标题】如何为 pytorch 层指定名称？

【英文标题】：How to assign a name for a pytorch layer?

【发布时间】：2021-05-15 01:51:54

【问题描述】：

在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 层分配名称的推荐方法是什么？

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

顺序：

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

动态：

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

直接：

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

其他我没想到的方式

---

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

【问题讨论】：

对于nn.Sequential，您可以提供nn.Modules 的OrderedDict，从而为它们命名

【参考方案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_namei": torch.nn.Conv2d(10, 10, 3) for i in range(5)
        )

将为每个创建的模块动态地提供whatever.my_namei.weight（或bias）。

直接

随便取个名字就行了

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

你没想到

如果您想绘制权重、偏差及其梯度，您可以沿着这条路线走

您不能以这种方式绘制激活

（或激活的输出）。请改用PyTorch hooks（如果您希望每层梯度在通过网络时也可以使用）

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

【讨论】：

谢谢！这就是我想要的。您能否解释一下torchfunc 与此有何关系？我能用 torchfunc 做什么？例如，我将如何用它绘制激活和梯度？

@Gulzar 它会自动注册记录器（通过钩子），这些记录器获取来自指定层的数据（您可以稍后绘制），请参阅documentation。 flashtorch 也可能会有所帮助，尽管我对此没有任何经验。

好的，现在尝试实现这一点，我明白你关于激活的观点。你能解释一下为什么我不能像这样绘制激活吗？另请注意，我使用的是 Pytorch Lightning，可能钩子已经实现。我需要哪一个？ on_after_backward 可能吗？

另外，我发现保留这些统计数据会减慢 X100 的训练速度。在这种情况下，权衡日志记录和速度的最佳做法是什么？

@Gulzar 请避免多问，发布一个关于激活的新问题并使其更具体，谢谢。