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())))
是为了看一下神经网络里面所有的参数