PAYATHON 2026

Keras 如何不分配权重内存地查看网络结构？

数据小陈

结论

tf.config.experimental.set_memory_growth() 不适用于 CPU。示例中的 ValueError 是因为传入了 CPU 设备，而该设置只支持 GPU 等设备。即使改用 GPU 内存增长设置，它也只会改变 GPU 显存的分配方式，无法让已经创建并构建的 Keras 模型免于占用权重内存。

只查看模型结构时，可以单独保存层、连接关系、输入输出形状和参数数量等元数据，再用 Graphviz 等工具绘图。这样的描述对象不是 Keras Model，不能直接传给 tf.keras.utils.model_to_dot()。

为什么创建模型会占用内存

原问题中的 model = create_vgg_like_model() 会创建一个 Keras 模型。模型的层在构建时通常会创建权重变量，这些变量本身就需要内存，无论模型之后是否用于训练或推理。调用 model_to_dot(model) 时，工具会根据已有模型结构生成图，不会自动把完整模型转换成不含权重的描述对象。

TensorFlow 的 GPU 内存增长设置不支持 CPU。错误信息 Cannot set memory growth on non-GPU and non-Pluggable devices 表示传入的设备类型不受支持。将 CPU 设为可见设备也不会延迟分配 CPU 权重。

只画结构时使用轻量元数据

可以为每层记录绘图所需的信息，并维护层之间的连接关系。例如：

layers = [
    {"name": "input", "shape": "224×224×3", "params": 0},
    {"name": "conv1", "shape": "224×224×64", "params": 1792},
    {"name": "pool1", "shape": "112×112×64", "params": 0},
]
edges = [("input", "conv1"), ("conv1", "pool1")]

这里的数字只用于演示数据格式，不代表 VGG16 的实际层配置或参数量。可以直接根据这些元数据生成 Graphviz 图，无需创建 Keras 权重：

from graphviz import Digraph

layers = [
    {"name": "input", "shape": "224×224×3", "params": 0},
    {"name": "conv1", "shape": "224×224×64", "params": 1792},
    {"name": "pool1", "shape": "112×112×64", "params": 0},
]
edges = [("input", "conv1"), ("conv1", "pool1")]

dot = Digraph("model")
for layer in layers:
    label = (
        f"{layer['name']}\n"
        f"output: {layer['shape']}\n"
        f"params: {layer['params']}"
    )
    dot.node(layer["name"], label)

for source, target in edges:
    dot.edge(source, target)

dot.render("model_architecture", format="png", cleanup=True)

计算参数数量时，需要根据层类型和配置生成元数据。例如，带 bias 的 Conv2D 参数量为 kernel_height × kernel_width × input_channels × filters + filters；带 bias 的 Dense 参数量为 input_features × units + units。实际计算还要考虑是否启用 bias、分组卷积等配置，不能只根据层名称推断。

需要使用 Keras 绘图工具时

如果要使用 model_to_dot() 或 tf.keras.utils.plot_model()，通常需要传入实际的 Keras Model。构建模型时会创建相应的变量，这些工具没有通用的“只构建图、不创建权重”开关。可以检查模型定义中是否意外包含很大的全连接层，或只创建需要展示的子模型。能否因此减少内存占用，取决于具体层和模型结构。

在 GPU 上运行时，可以为 GPU 设备配置内存增长，以控制显存的分配方式；这不会减少模型权重所需的总内存，也不适用于 CPU。只为绘制架构时，通常更适合维护元数据并独立绘图。

备注：内容仅供参考。