如何统计TensorFlow Hub加载的目标检测模型的参数数量

方案可行性说明 #

将hub.KerasLayer嵌入tf.keras.Sequential构建的模型中，完全可以实现调用summary()查看参数的需求，你之前第二种方法返回空权重列表，是因为KerasLayer未完成初始化：没有指定输入形状、也没有传入实际输入的情况下，层的权重不会被加载。

具体实现代码 #

import tensorflow as tf
import tensorflow_hub as hub

# 构建Sequential模型
model = tf.keras.Sequential([
    # 需指定input_shape匹配模型要求的输入尺寸，ssd_mobilenet_v2要求输入为[batch, 320, 320, 3]的uint8张量
    hub.KerasLayer("https://tfhub.dev/tensorflow/ssd_mobilenet_v2/2", input_shape=(320, 320, 3), dtype=tf.uint8)
])

# 直接调用summary即可查看参数统计
model.summary()

# 如需单独统计权重数量/总参数量，可调用以下代码
print("权重张量个数：", len(model.weights))
total_params = sum(tf.size(w).numpy() for w in model.weights)
print("总参数量：", total_params)

原有方法的优化补充 #

你第一种方法获取的是权重张量的个数，不是总参数量，每个张量内部包含多个参数值，可直接通过以下代码计算总参数量，结果和Kerassummary()输出一致：

hub_model = hub.load("https://tfhub.dev/tensorflow/ssd_mobilenet_v2/2")
variables = hub_model.signatures['serving_default'].variables
total_params = sum(tf.size(var).numpy() for var in variables)
print("总参数量：", total_params)

内容的提问来源于stack exchange，提问作者Davide Brattelli