如何使用fvcore处理含多参数forward方法的模型？

你的包装函数问题出在闭包绑定参数的方式无法被fvcore正确解析，FlopCountAnalysis需要明确追踪所有模型输入张量的计算流程，而你原来的代码只把input_tensor作为输入传给包装后的函数，其他参数提前绑定在闭包里，导致fvcore无法识别这些参数是模型的输入部分，进而报错。

修正后的实现代码 #

from fvcore.nn import FlopCountAnalysis, parameter_count_table

def modelCount(model, *inputs, **kwargs):
    def _wrapped_forward(*args):
        # 直接将所有位置参数和关键字参数传递给原模型的forward
        return model(*args, **kwargs)
    
    # 将所有输入张量传入FlopCountAnalysis，而非单个张量
    flops = FlopCountAnalysis(_wrapped_forward, inputs).total()
    params = parameter_count_table(model)
    return flops, params

使用方式示例 #

假设你的模型forward定义如下：

import torch.nn as nn
import torch

class MyModel(nn.Module):
    def forward(self, x, y, attention_mask=None, dropout_rate=0.1):
        # 模型计算逻辑示例
        feat = nn.Linear(1024, 512)(x + y)
        if attention_mask is not None:
            feat = feat * attention_mask.unsqueeze(-1)
        feat = nn.Dropout(dropout_rate)(feat)
        return feat

调用统计函数时，直接按模型forward的参数传递即可：

# 准备输入张量
x = torch.randn(1, 1024)
y = torch.randn(1, 1024)
attention_mask = torch.ones(1, 1024)

# 统计Flops和参数
total_flops, param_table = modelCount(MyModel(), x, y, attention_mask=attention_mask, dropout_rate=0.1)
print(f"Total Flops: {total_flops}")
print(param_table)

关键说明 #

• 所有张量类型的位置参数直接作为modelCount的位置参数传入，顺序和模型forward的位置参数一致。
• 非张量类型的配置参数（如dropout_rate）或可选张量参数（如attention_mask）以关键字参数形式传入即可。
• 修正后的包装函数不再提前绑定参数，而是让FlopCountAnalysis能追踪到所有输入张量的计算路径，避免解析错误。

内容的提问来源于stack exchange，提问作者H.M