计算模型大小


import tensorflow as tf

# 构建一个简单的模型
model = tf.keras.Sequential([
    tf.keras.layers.Dense(64, activation='relu', input_shape=(10,)),
    tf.keras.layers.Dense(32, activation='relu'),
    tf.keras.layers.Dense(1)
])

# 计算模型参数数量
total_params = 0
for variable in tf.trainable_variables():
    shape = variable.get_shape()
    variable_params = 1
    for dim in shape:
        variable_params *= dim
    total_params += variable_params

# 计算模型大小（以兆字节为单位）
model_size_mb = total_params * 4 / (1024 * 1024)  # 假设每个参数占用 4 字节（float32）

print(f"模型参数数量: {total_params}")
print(f"模型大小 (MB): {model_size_mb:.4f}")