使用thop库计算flops和params

from thop import profile
from thop import clever_format

input_data = torch.randn(1, 3, 256, 256).cuda()

flops, params = profile(model, inputs=(input_data,))
flops, params = clever_format([flops, params], "%.3f")
print(f"FLOPs: {flops}, Params: {params}")