pytorch和tensorflow一样展示summary

最新推荐文章于 2024-08-09 09:28:46 发布

其实也很简单

最新推荐文章于 2024-08-09 09:28:46 发布

阅读量381

点赞数

文章标签：深度学习神经网络算法

原文链接：https://zhuanlan.zhihu.com/p/105969106

版权

pytorch有一个库，叫torchsummary 。

转载：

import torch
from torchsummary import summary

from models import TSN

net = TSN(101, 8, 'RGB', base_model='resnet50', consensus_type='avg', dropout=0.5,
          partial_bn=False, pretrain='imagenet',
          is_shift=False, shift_div=8, shift_place='blockres',
          fc_lr5=True, temporal_pool=False, non_local=False)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = net.to(device)

summary(model, (1, 24, 28, 28))