pytorch obtains the parameter amount of the model and the size of the model

reference

• Finding model size
• Parameter and buffer in Pytorch model
• What pytorch means by buffers?
• The difference between Module, Parameter and Buffer in Pytorch
• torch.Tensor.element_size
• torch.Tensor.nelement

buffers and parameters

In the model, there are two types of buffer and parameter. The parameter is the parameter of the model that we generally think of. It has a gradient and can be updated by training. But the buffer has no gradient and cannot be updated by training.
We can pass torch.nn.Module.buffers()and torch.nn.Module.named_buffers()return the buffer in the model. The second function returns its own defined name and buffer at the same time.
At the same time, we can pass torch.nn.Module.parameters()and torch.nn.Module.named_parameters()return parameters in the model. The second function returns both the name and parameter defined by itself.
Both functions have a bool parameter recurse, which defaults to true. If yes true, all sublayer parameters will be recursively searched. Otherwise only sublayers of the first layer are searched.

torch.Tensor.nelement和torch.Tensor.element_size

The parameters and buffers we get are both Tensor type parameters, and for Tensor, the first function can return the number of elements in this Tensor, such as how many are in the matrix. The second function can return the byte size of the data type corresponding to this Tensor. For example, float32 is 4 bytes.

get the size of the model

def getModelSize(model):
    param_size = 0
    param_sum = 0
    for param in model.parameters():
        param_size += param.nelement() * param.element_size()
        param_sum += param.nelement()
    buffer_size = 0
    buffer_sum = 0
    for buffer in model.buffers():
        buffer_size += buffer.nelement() * buffer.element_size()
        buffer_sum += buffer.nelement()
    all_size = (param_size + buffer_size) / 1024 / 1024
    print('模型总大小为：{:.3f}MB'.format(all_size))
    return (param_size, param_sum, buffer_size, buffer_sum, all_size)

The function is also easy to understand. By model.parameters()returning an iterator that can iterate over all parameters, you can get all of them through the for loop parameter. The buffer is also similar.
What is returned param_sizeis the parameter byte MB size of all parameters, buffer_sizewhich is the parameter byte MB size of all buffers, all_sizewhich is the MB size of the model.