import torch.nn as nn
import torch
import torch.nn.functional as F
class FP_Conv2d(nn.Module):
    def __init__(self, input_channels, output_channels,
            kernel_size=-1, stride=-1, padding=-1, dropout=0, groups=1, channel_shuffle=0, shuffle_groups=1, last=0, first=0):
        super(FP_Conv2d, self).
 
__init__()
        self.dropout_ratio = dropout
        self.last = last
        self.first_flag = first
        if dropout!=0:
            self.dropout = nn.Dropout(dropout)
        self.conv = nn.Conv2d(input_channels, output_channels,
                kernel_size=kernel_size, stride=stride, padding=padding, groups=groups)
        self.bn 
 
= nn.BatchNorm2d(output_channels)
        self.relu = nn.ReLU(inplace=True)
    def forward(self, x):
        if self.first_flag:
            x = self.relu(x)
        if self.dropout_ratio!=0:
            x = self.dropout(x)
        x = self.conv(x)
        x = self.bn(x)
        x = self.relu(x)
        
 
return x

class Net(nn.Module):
    def __init__(self, cfg = None):
        super(Net, self).__init__()
        if cfg is None:
            cfg = [192, 160, 96, 192, 192, 192, 192, 192]
        self.tnn_bin = nn.Sequential(
                nn.Conv2d(3, cfg[0], kernel_size=5, stride=1, padding=2),#預設0填充，3為輸出通道數，cfg[0]為輸出通道數
                nn.BatchNorm2d(cfg[0]),
                FP_Conv2d(cfg[0], cfg[1], kernel_size=1, stride=1, padding=0, first=1),
                FP_Conv2d(cfg[1], cfg[2], kernel_size=1, stride=1, padding=0),
                nn.MaxPool2d(kernel_size=3, stride=2, padding=1),
                FP_Conv2d(cfg[2], cfg[3], kernel_size=5, stride=1, padding=2),
                FP_Conv2d(cfg[3], cfg[4], kernel_size=1, stride=1, padding=0),
                FP_Conv2d(cfg[4], cfg[5], kernel_size=1, stride=1, padding=0),
                nn.AvgPool2d(kernel_size=3, stride=2, padding=1),
                FP_Conv2d(cfg[5], cfg[6], kernel_size=3, stride=1, padding=1),
                FP_Conv2d(cfg[6], cfg[7], kernel_size=1, stride=1, padding=0),
                nn.Conv2d(cfg[7],  10, kernel_size=1, stride=1, padding=0),
                nn.BatchNorm2d(10),
                nn.ReLU(inplace=True),
                nn.AvgPool2d(kernel_size=8, stride=1, padding=0),#平均值
                )
    
    def forward(self, x):
        x = self.tnn_bin(x)
        #x = self.dorefa(x)
        x = x.view(x.size(0), 10)
        return x

需要安裝thop:pip install thop

無法安裝的話：pip install --upgrade git+https://github.com/Lyken17/pytorch-OpCounter.git

#pytorch 計算模型的Para和GFLOPs
from torchvision.models import resnet18
import torch
from thop import profile
from models import nin
model = nin.Net()
#checkpoint = torch.load('models_save/nin.pth',map_location='cpu')
checkpoint = torch.load('models_save/nin_preprune.pth',map_location='cpu')
model.load_state_dict(checkpoint['state_dict'])
input = torch.randn(1, 3, 32, 32) #模型輸入的形狀,batch_size=1
flops, params = profile(model, inputs=(input, ))
print("GFLOPs :{:.2f}, Params : {:.2f}".format(flops/1e9,params/1e6)) #flops單位G，para單位M