Python pytorch网络结构的推导

对于我的用例，我需要能够获取pytorch模块并解释模块中的层序列，以便能够以某种文件格式在层之间创建“连接”。现在让我们假设我有一个简单的模块，如下所示

class mymodel(nn.Module):
    def __init__(self, input_channels):
        super(mymodel, self).__init__()
        self.fc = nn.Linear(input_channels, input_channels)
    def forward(self, x):
        out = self.fc(x)
        out += x
        return out


if __name__ == "__main__":
    net = mymodel(5)

    for mod in net.modules():
        print(mod) 

这里的输出结果是：

mymodel(
  (fc): Linear(in_features=5, out_features=5, bias=True)
)
Linear(in_features=5, out_features=5, bias=True)

如您所见，由于加号等于操作或加号操作不是forward函数中的NN模块，因此未捕获有关加号等于操作或加号操作的信息。我的目标是能够从pytorch模块对象创建一个图形连接，用json表示如下内容：

layers {
"fc": {
"inputTensor" : "t0",
"outputTensor": "t1"
}
"addOp" : {
"inputTensor" : "t1",
"outputTensor" : "t2"
}
}

输入张量名称是任意的，但它捕获了图形的本质和层之间的连接

---

我的问题是，有没有办法从pytorch对象中提取信息？我本来打算使用.modules（），但后来意识到手写操作不是以这种方式作为模块捕获的。我想如果所有的东西都是nn.module，那么.modules（）可能会给我网络层的安排。这里需要帮助。我想知道张量之间的联系，以创建如上所述的格式

您要查找的信息不存储在

nn.Module

中，而是存储在输出张量的

grad\u fn

属性中：

model=mymodel（通道）
pred=型号（火炬号（（1个通道））
所有信息都在输出张量的计算图中

提取这些信息并不是件小事。您可能想看看从

grad\u fn

信息中绘制一个漂亮图形的包