pytorch利用thop计算网络参数量和浮点数

import torch
from torchsummary import summary
from thop import profile
# 网络输入仅为一个input
flops, params = profile(self.model, inputs=(input, ))

# 网络输入有多个组成部分
flops, params = profile(self.model, inputs=(batch_x, batch_x_mark, dec_inp, batch_y_mark))  # batch_x, batch_x_mark, dec_inp, batch_y_mark为网络的输入,根据具体网络的输入设置
print("%s | Params: %.2fM | FLOPs: %.2fM" % ("Informer", params / (1000 ** 2), flops / (1000 ** 3)))

你可能感兴趣的:(pytorch实现,大数据)