Pytorch basic learning-torchsummary calculation parameter quantity

Here is a well-used tool - torchsummary, which can view the structure of each layer of the network, its size, and the amount of parameters, and finally calculate the size and amount of parameters of the entire network.

1. Installation

pip install torchsummary

2. Test

Enter in python modeimport torchsummary

3. Use

import torchsummary

net = xxx # 自己的网络
print(torchsummary.summary(net, (3, 256, 256), device="cuda"))

• torchsummary.summary(model, input_size, batch_size=-1, device="cuda")
• model: own network model
• input_size: input size, the shape is C, H, W, respectively the number of channels, height, and width
• batch_size: batch_size, the default is -1, the batch_size displayed when displaying the output shape of each layer of the model
• device: select "cuda" or "cpu"

4. Results display