计算nn.Linear的计算量FLOPs

import torch
import torch.nn as nn
m = nn.Linear(20, 30)
input = torch.randn(128, 3,20)
output = m(input)
print(output.size())
flops = (torch.prod(torch.LongTensor(list(output.size()))) \
         * input[0].size(1)).item()
print((list(output.size())))
print(input[0].size(1))
print(flops)

官网对nn.Linear的介绍：

Applies a linear transformation to the incoming data:
$$y=x{A}^T+b$$
其函数签名为
torch.nn.Linear(in_features, out_features, bias=True)
其参数为：

• in_features - size of each input sample

• out_features -size of each output sample

• bias If set to false, the layer will not learn an additive bias.
  其输入输出的维度为：
• Input: $(N,\ast ,H_{in})$, where *∗ means any number of additional dimensions and $H_{in}$=in_features
• Output: $(N,\ast ,H_{out})$, where all but the last dimension are the same shape as the input and $H_{out}$=out_features

故代码运行得到的计算量FLOPS为 128×3×20×30 = 230400，即计算量FLOPs为batch_size×特征维度值×输入特征数×输出特征数。