pytorch 获取参数量代码片段

最新推荐文章于 2025-02-27 21:21:27 发布

原创最新推荐文章于 2025-02-27 21:21:27 发布 · 272 阅读

0 ·

CC 4.0 BY-SA版权

pytorch 专栏收录该内容

28 篇文章

订阅专栏

这段代码用于计算一个模型的总参数数量和训练参数数量。通过迭代模型的参数，`sum(p.numel() for p in model.parameters())`求和得到全部参数数，而`sum(p.numel() for p in model.parameters() if p.requires_grad)`则计算需要梯度更新的训练参数数。

# Find total parameters and trainable parameters
total_params = sum(p.numel() for p in model.parameters())
print(f'{total_params:,} total parameters.')
total_trainable_params = sum(
    p.numel() for p in model.parameters() if p.requires_grad)
print(f'{total_trainable_params:,} training parameters.')