python - 如何在图中获取所有张量？

Question

我想访问图的所有张量实例。例如，我可以检查张量是否分离，或者我可以检查大小。可以在 tensorflow 中完成.

我 不要想要图形的可视化。

Answer 1

您可以在运行时访问整个计算图。为此，您可以使用 Hook 。这些是插入到 nn.Module 上的功能s 用于推理和反向传播。
在推理时，您可以使用 register_forward_hook Hook 回调函数.同样对于反向传播，您可以使用 register_full_backward_hook .
注意:从 PyTorch 版本 1.8.0 register_backward_hook 开始已被弃用。
使用这两个函数，您基本上可以访问计算图上的任何张量。是否要打印所有张量、打印形状，甚至插入断点进行调查，这完全取决于您。
这是一个可能的实现:

def forward_hook(module, input, output):
    # ...

论据 input由 PyTorch 作为元组传递，并将包含传递给 Hook 模块的 forward 函数的所有参数。

def backward_hook(module, grad_input, grad_output):
    # ...

对于后钩， grad_input和 grad_output将是元组，并且根据模型的层具有不同的形状。
然后，您可以将这些回调 Hook 到任何现有的 nn.Module .例如，您可以遍历模型中的所有子模块:

for module in model.children():
    module.register_forward_hook(forward_hook)
    module.register_full_backward_hook(backward_hook)

要获取模块的名称，您可以将钩子(Hook)包裹起来，并在模型的 named_modules 上循环。 :

def forward_hook(name):
    def hook(module, x, y):
        print(f'{name}: {[tuple(i.shape) for i in x]} -> {list(y.shape)}')
    return hook

for name, module in model.named_children():
    module.register_forward_hook(forward_hook(name))

这可以在推理上打印以下内容:

fc1: [(1, 100)] -> (1, 10)
fc2: [(1, 10)] -> (1, 5)
fc3: [(1, 5)] -> (1, 1)

至于模型的参数，您可以通过调用 module.parameters 轻松访问两个钩子(Hook)中给定模块的参数。 .这将返回一个生成器。