计算网络模型的参数量和计算量

from thop import profile
import torch

device = torch.device('cuda')
model = IFNet().to(device)
input = torch.randn(1, 6, 256, 256).to(device)  # 根据你的输入分辨率设置
flops, params = profile(model, inputs=(input,), verbose=False)

print(f"FLOPs: {flops / 1e9:.2f} GFLOPs")
print(f"参数量: {params / 1e6:.2f} M")