import torch
import torchvision

model = torchvision.models.detection.fasterrcnn_resnet50_fpn(pretrained=True)

for name, p in model.named_parameters():
    print(name)
    print(p.requires_grad)
    print(...)

#或者

for p in model.parameters():
    print(p)
    print(...)

class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()

        cfg = [64, 64, 'M', 128, 128, 'M', 256, 256, 256, 'M', 512, 512, 512, 'M', 512, 512, 512]
        self.features = self._vgg_layers(cfg)

    def _vgg_layers(self, cfg):
        layers = []
        in_channels = 3
        for x in cfg:
            if x == 'M':
                layers += [nn.MaxPool2d(kernel_size=2, stride=2)]
            else:
                layers += [nn.Conv2d(in_channels, x ,kernel_size=3, padding=1),
                        nn.BatchNorm2d(x),
                        nn.ReLU(inplace=True)
                        ]
                in_channels = x
            
        return nn.Sequential(*layers)

    def forward(self, data):
        out_map = self.features(data)
        return out_map
    
Model = Net()

for name, p in model.named_parameters():
    print(name)
    print(p.requires_grad)
    print(...)

#或者

for p in model.parameters():
    print(p)
    print(...)

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)

def cnn_paras_count(net):
    """cnn參數(shù)量統(tǒng)計, 使用方式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)

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)

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