pytorch查看模型参数总结

1:DNN_printer
其中(3, 32, 32)是输入的大小,其他方法中的参数同理

from DNN_printer import DNN_printer

batch_size = 512
def train(epoch):
    print('\nEpoch: %d' % epoch)
    net.train()
    train_loss = 0
    correct = 0
    total = 0
    // put the code here and you can get the result
    DNN_printer(net, (3, 32, 32),batch_size)

结果

2:parameters

def cnn_paras_count(net):
    """cnn参数量统计, 使用方式cnn_paras_count(net)"""
    # Find total parameters and trainable parameters
    total_params = sum(p.numel() for p in net.parameters())
    print(f'{total_params:,} total parameters.')
    total_trainable_params = sum(p.numel() for p in net.parameters() if p.requires_grad)
    print(f'{total_trainable_params:,} training parameters.')
    return total_params, total_trainable_params

cnn_paras_count(net)

直接输出参数量,然后自己计算
需要注意的是,一般模型中参数是以float32保存的,也就是一个参数由4个bytes表示,那么就可以将参数量转化为存储大小。
例如:

44426个参数*4 / 1024 ≈ 174KB

3:get_model_complexity_info()

from ptflops import get_model_complexity_info
from torchvision import models

net = models.mobilenet_v2()
ops, params = get_model_complexity_info(net, (3, 224, 224), as_strings=True, 
										print_per_layer_stat=True, verbose=True)

pytorch查看模型参数总结_第1张图片
4:torchstat

from torchstat import stat
import torchvision.models as models
model = models.resnet152()
stat(model, (3, 224, 224))

输出
pytorch查看模型参数总结_第2张图片

参考链接
第一个方法
第二个方法
第三个方法
第四个方法

你可能感兴趣的:(pytorch,深度学习)