Skip to content

contribs.torchsummary.summary

Function ยท Source

params_info = mdnc.contribs.torchsummary.summary(
    model, input_size, batch_size=-1, device='cuda:0', dtypes=None
)

Iterate the whole pytorch model and summarize the infomation as a Keras-style text report. The output would be store in a str.

Arguments

Requries

Argument Type Description
model nn.Module The pyTorch network module instance. It is to be analyzed.
input_size (seq/int, ) A sequence (list/tuple) or a sequence of sequnces, indicating the size of the each model input variable.
batch_size int The batch size used for testing and displaying the results.
device str or
torch.device		 Should be set according to the deployed device of the argument model.
dtypes (torch.dtype, ) A sequence of torch data type for each input variable. If set None, would use float type for all variables.

Returns

Argument Description
params_info A tuple of two values. The first value is the total parameter numbers. The second value is the trainable parameter numbers.

Examples

Example
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
import torch
import torch.nn as nn
import torch.nn.functional as F
import mdnc

class TestTupleOutModule(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc1a = nn.Linear(300, 50)
        self.fc1b = nn.Linear(50, 10)

        self.fc2a = nn.Linear(300, 50)
        self.fc2b = nn.Linear(50, 10)

    def forward(self, x1, x2):
        x1 = F.relu(self.fc1a(x1))
        x1 = self.fc1b(x1)
        x2 = x2.type(torch.FloatTensor)
        x2 = F.relu(self.fc2a(x2))
        x2 = self.fc2b(x2)
        # set x2 to FloatTensor
        x = torch.cat((x1, x2), 0)
        return F.log_softmax(x, dim=1), F.log_softmax(x1, dim=1), F.log_softmax(x2, dim=1)

input1 = (1, 300)
input2 = (1, 300)
dtypes = (torch.FloatTensor, torch.LongTensor)
total_params, trainable_params = mdnc.contribs.torchsummary.summary(
    TestTupleOutModule(), (input1, input2), device='cpu', dtypes=dtypes)

----------------------------------------------------------------
        Layer (type)               Output Shape         Param #
================================================================
            Linear-1                [-1, 1, 50]          15,050
            Linear-2                [-1, 1, 10]             510
            Linear-3                [-1, 1, 50]          15,050
            Linear-4                [-1, 1, 10]             510
TestTupleOutModule-5                [-1, 1, 10]               0
                                    [-1, 1, 10]                
                                    [-1, 1, 10]                
================================================================
Total params: 31,120
Trainable params: 31,120
Non-trainable params: 0
----------------------------------------------------------------
Input size (MB): 0.00
Forward/backward pass size (MB): 0.00
Params size (MB): 0.12
Estimated Total Size (MB): 0.12
----------------------------------------------------------------
Last update: March 14, 2021

Comments