计算神经网络每一层需要训练的参数和每一层的FLOPs

FLOPs是什么？

FLOPs：注意s小写，是floating point operations的缩写（s表复数），意指浮点运算数，理解为计算量。可以用来衡量算法/模型的复杂度。

开始实现

使用的是pytorch来实现神经网络，先导入这几个包。为啥那，我再想想。。。

import torch
import torch.nn.functional as F
import torch.nn as nn
from thop import profile

定义神经网络，这里常用了继承torch.nn.Module的方法，至于其他的方法，我只是一个小白。。。要什么其他方法。

class Net(torch.nn.Module):
    def __init__(self, n_features, n_hidden, n_output):
        super(Net, self).__init__()
        self.hidden = torch.nn.Linear(n_features, n_hidden)
        self.predict = torch.nn.Linear(n_hidden, n_output)

    def forward(self, x):
        x = F.relu(self.hidden(x))
        x = self.predict(x)
        return x

定义了一个获得神经网络每层参数的函数getParameter，

在getParameter里面进一步写出我想要的东西，返回的是一个字典

def getParameter(net):  # return a dictionary that includes every layer name and count of parameter

    nameList = []
    myReturnList = {}
    for k in net.state_dict():
        nameList.append(k)

    for index in range(0, len(nameList), 2):
        parameter = net.state_dict()[nameList[index]].numel() + net.state_dict()[nameList[index + 1]].numel()
        myReturnList[nameList[index].split('.')[0]] = parameter

    return myReturnList

定义了一个获得每层FLOPs的函数getFlops

在getFlops里面进一步写出我想要的东西，返回的依然是一个字典。

def getFlops(net):
    myParameterlist = getParameter(net)
    myInput = torch.unsqueeze(torch.linspace(-1, 1, 100), dim=1)
    flops, params = profile(net, inputs=(myInput,))
    average = flops/params

    myFlopslist = {}

    for name in myParameterlist:
        myFlopslist[name] = myParameterlist[name] * average

    return myFlopslist

这个getFlops方法里啊，主要的思路是，获得了整个神经网络的FLOPs和需要训练的参数，然后，求出一个需要训练的参数所需要的FLOPs，又通过是一个定义的函数getParameter知道每一层的参数个数，就求出每一层的参数了。

最后输出

net = Net(1, 10, 1)
print(net)
print("Total number of paramerters in networks is {}  ".format(sum(temp.numel() for temp in net.parameters())))
myReturnList = getParameter(net)
print(myReturnList)
myflops = getFlops(net)
print(myflops)

这一行print("Total number of paramerters in networks is {} ".format(sum(temp.numel() for temp in net.parameters())))是为了看一下神经网络里面所有的参数