pytorch和tensorflow一样展示summary

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))