可视化LSTM模型参数

 请预先安装torchinfo:

pip install torchinfo

 然后导入库：

from torchinfo import summary

使用resnet进行测试:

import torchvision

model = torchvision.models.resnet152()
summary(model, (1, 3, 224, 224), depth=3)

 出现如下界面表示成功：

 就可以查看每一层的输出的shape和每一层的参数数量。

---

查看LSTM参数

from torchinfo import summary

batch, seq_len, insize = 64, 10, 56
input = torch.zeros((batch, seq_len, insize))
model = CLSTM(input_size=insize, channels=[64, 128, 256], hidden_size=150, out_size=1, num_layers=2)
state = model.init_state(batch)
y = model(input, state)
print(model, model(input, state)[0].shape)
summary(model, input.shape)

请注意CLSTM是我自己的模型，可以替换为你自己的LSTM模型，summary传参的时候，第一个参数为模型，第二个参数为input的shape，后面还有一些参数，例如input_data、batch_dim等，可以在源码中查看

那么对于LSTM来说，请将模型的输入的state默认参数设置为None，不然传参的时候会报错，提示传入的参数不够，需要传入state参数。

 注意state=none

def forward(self, x: torch.Tensor, state=None):
    # input shape:[batch_size,length,in_size]
    # input shape of cnn:[batch_size,in_size,length]
    out = self.block1(x.permute(0, 2, 1))
    out = self.block2(out)
    # shape change:[batch_size,in_size,length]->[batch_size,length,in_size]
    out = self.block3(out).permute(0, 2, 1)
    out, state = self.rnn(out, state)

    out = self.dense(out)
    return out, state