When we are using a deep learning model built with pytorch, we may want to know what parameters are contained in this model. In this tutorial, we will use an example to show you how to see these parameters name and shape.
For example:
torch.nn.Conv2d() is often used in many pytorch models. It contains two parameters: a kernel weight and bias.
Understand torch.nn.Conv2d() with Examples – PyTorch Tutorial
You may want to know their name and shape in torch.nn.Conv2d(). Here is the example code.
import torch
N = 40
C_in = 40
H_in = 32
W_in = 32
inputs = torch.rand([N, C_in, H_in, W_in])
padding = 3
kernel_size = 3
stride = 2
C_out = 100
x = torch.nn.Conv2d(C_in, C_out, kernel_size, stride=stride, padding=padding)
y = x(inputs)
for p, v in x.state_dict().items():
print(p, v.shape)
Run this code, we will see:
weight torch.Size([100, 40, 3, 3]) bias torch.Size([100])
It means the kernel weight is: 100*40*kernel_size. It is [c_out, c_in, kernal_h, kernel_w]
bias: is c_out = 100
We also can find: in order to get the parameter name and shape, we can use model.state_dict().items(). It is easy to use.
Moreover, if you want to get the size of all parameters in a pytorch model. You can read: