Pytorch 计算量FLOPs和参数量Parameter的统计

在PyTorch中，计算模型的参数（Parameter）和浮点运算数（FLOP）可以帮助我们了解模型的复杂度和资源消耗。以下是如何在PyTorch中计算Parameter和FLOP的步骤： ### 计算模型参数（Parameter） 要计算模型的参数数量，可以使用`torch.nn.Module`提供的`parameters()`方法。这个方法会返回模型中所有的参数，然后我们可以使用`numel()`方法计算每个参数的元素数量，并求和得到总参数数量。 ```python import torch.nn as nn def count_parameters(model): return sum(p.numel() for p in model.parameters() if p.requires_grad) # 假设我们有一个模型 model = nn.Sequential( nn.Linear(10, 50), nn.ReLU(), nn.Linear(50, 1) ) # 计算参数数量 num_params = count_parameters(model) print(f'The model has {num_params} trainable parameters') ``` ### 计算浮点运算数（FLOP） 计算FLOP稍微复杂一些，因为需要根据模型的架构手动计算每层的计算量。以下是一个简单的例子，展示了如何计算一个包含线性层和ReLU激活函数的模型的FLOP。 ```python import torch import torch.nn as nn def count_flops(model, input_size): flops = 0 input_shape = (1, *input_size) device = next(model.parameters()).device x = torch.randn(input_shape).to(device) hooks = [] def hook_fn(module, input, output): if isinstance(module, nn.Linear): out_features, in_features = module.weight.shape flops_per_sample = in_features * out_features flops += flops_per_sample * x.size(0) elif isinstance(module, nn.Conv2d): # 简单的卷积层FLOP计算 out_h = output.size(2) out_w = output.size(3) flops_per_sample = module.in_channels * module.out_channels * module.kernel_size[0] * module.kernel_size[1] * out_h * out_w flops += flops_per_sample * x.size(0) for m in model.modules(): if isinstance(m, (nn.Linear, nn.Conv2d)): hooks.append(m.register_forward_hook(hook_fn)) with torch.no_grad(): model(x) for hook in hooks: hook.remove() return flops # 假设我们有一个模型 model = nn.Sequential( nn.Linear(10, 50), nn.ReLU(), nn.Linear(50, 1) ) # 计算FLOP flops = count_flops(model, input_size=(10,)) print(f'The model requires {flops} FLOPs') ``` ### 总结 - **参数计算**：使用`model.parameters()`和`numel()`方法。 - **FLOP计算**：需要手动计算每层的计算量，并使用`register_forward_hook`方法注册钩子函数。