使用thop计算放入GPU中的模型的参数量和FLOPs

安装

pip install thop

使用

from thop import profile


input = torch.randn(1, 3, 1024).to(device)  # 放进和model相同的GPU中
flops, params = profile(model, inputs=(input,), verbose=True)
print("%s ------- params: %.2fMB ------- flops: %.2fG" % (model, params / (1000 ** 2), flops / (1000 ** 3)))  # 这里除以1000的平方，是为了化成M的单位

注意：

1. 这里的input需要和实际data的尺寸相同，前面的 1 是batchsize，一般都设为 1 来计算模型的参数量和FLOPs。
2. 这里使用thop需要注释掉模型的分布式训练命令，保证模型和数据在同一个GPU上。

# model = nn.DataParallel(model)