pytorch打印模型信息——torchinfo

最新推荐文章于 2024-06-23 20:49:59 发布

wxyczhyza

最新推荐文章于 2024-06-23 20:49:59 发布

阅读量1.2k

点赞数

分类专栏： # pytorch 文章标签： pytorch 深度学习人工智能

本文链接：https://blog.csdn.net/wxyczhyza/article/details/128118418

版权

pytorch 专栏收录该内容

5 篇文章 1 订阅

订阅专栏

用torchinfo库的sumary函数可以打印模型信息，示例如下。注意输入张量中不需要batch维度。

import numpy as np

import torch
import torchvision
import torchinfo

model = torchvision.models.resnet50(pretrained=True)
torchinfo.summary(model, (3, 224, 224), batch_dim=0,
                  col_names=('input_size', 'output_size', 'num_params', 'kernel_size', 'mult_adds'), verbose=1
        )