Advertisement

Pytorch高效展示网络结构与参数方法

阅读量:

首先,需要完成torchsummary的安装步骤。

复制代码
    pip isntall torchsummary
    
    
      
    

以下为示例代码,其中(3, 100, 100)分别对应单张图像的通道数量、高度以及宽度

复制代码
    from torchsummary import summary
    import torchvision
    import torch
    
    def model_info(model):
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    backbone = model.to(device)
    summary(backbone, (3, 100, 100))
    
    model = torchvision.models.vgg16(pretrained=False)
    model_info(model)
    
    
      
      
      
      
      
      
      
      
      
      
      
    

最终输出结果包含13个卷积层(Conv2d)与3个全连接层

全部评论 (0)

还没有任何评论哟~