MorphNet：快速且简单的资源受限深度网络结构学习

[目录]

新功能：FiGS：细粒度随机架构搜索

FiGS 是一种概率性的通道正则化方法，我们在《细粒度随机架构搜索》论文中首次提出（https://arxiv.org/pdf/2006.09581.pdf）。它优于我们之前的正则化方法，既可以作为剪枝算法使用，也可以作为一种完整的可微架构搜索方法。这是推荐的 MorphNet 使用方式，在以下文档中被称为 LogisticSigmoid 正则化器。

什么是 MorphNet？

MorphNet 是一种在训练过程中学习深度网络结构的方法。其核心思想是将网络结构学习问题进行连续松弛。简而言之，MorphNet 正则化器会逐渐降低滤波器的影响，当这些影响足够小的时候，对应的输出通道就会被标记为可以从网络中移除。

具体来说，通过添加针对特定资源消耗（如 FLOPs 或模型大小）的正则化项，可以诱导激活稀疏性。当正则化损失与训练损失相加，并通过随机梯度下降或其他类似优化器最小化时，学习问题就转化为一个受正则化项约束的网络结构优化问题。该方法最早在我们的 CVPR 2018 论文《MorphNet：快速且简单的资源受限深度网络结构学习》（https://arxiv.org/abs/1711.06798）中提出。关于该方法的概述以及针对不同设备延迟的正则化器，则在 GTC 2019 上进行了展示[幻灯片]，录像可在 YouTube 和 GTC 按需观看 中找到。我们新的概率性剪枝方法称为 FiGS，详细内容见《细粒度随机架构搜索》（https://arxiv.org/pdf/2006.09581.pdf）。

使用方法

假设你已经有一个用于图像分类的卷积神经网络，但希望缩小模型以满足某些约束条件（例如内存或延迟）。给定一个现有模型（“种子网络”）和一个目标指标，MorphNet 将通过调整每个卷积层的输出通道数来生成一个新的模型。

需要注意的是，MorphNet 不会改变网络的拓扑结构——生成的模型将与种子网络具有相同的层数和连接模式。

要使用 MorphNet，你需要：

1. 从 morphnet.network_regularizers 中选择一个正则化器。选择依据包括：

   • 你的目标成本（例如 FLOPs、延迟）
   • 你是否能够在模型中添加新层：
     • 在你希望剪枝的任何层之后添加我们的概率门控操作，并使用 LogisticSigmoid 正则化器。[推荐]
     • 如果无法添加新层，则根据你的网络架构选择正则化器类型：如果种子网络包含 BatchNorm，则使用 Gamma 正则化器；否则使用 GroupLasso 正则化器[已弃用]。

   注意：如果你使用 BatchNorm，必须启用缩放参数（“gamma 变量”），即在使用 tf.keras.layers.BatchNormalization 时设置 scale=True。

   注意：如果你使用 LogisticSigmoid，别忘了添加概率门控操作！示例如下。

2. 使用阈值以及模型的输出边界操作（可选地还包括输入边界操作）初始化正则化器。

   MorphNet 正则化器会从输出边界开始遍历你的计算图，并对遇到的一些操作应用正则化。当它遇到任何输入边界操作时，就不会再继续向下遍历（输入边界内的操作不会被正则化）。阈值决定了哪些输出通道可以被移除。

3. 将正则化项添加到你的损失函数中。

   和往常一样，正则化损失需要进行缩放。我们建议在以 1/初始成本 为中心的几个数量级范围内，沿对数尺度搜索缩放超参数（即正则化强度）。例如，如果种子网络的初始 FLOPs 数量为 1e9，则可以在 1e-9 左右探索正则化强度。

   注意：目前 MorphNet 尚未将正则化损失添加到 tf.GraphKeys.REGULARIZATION_LOSSES 集合中；这一选择可能会被修订。

   注意：不要混淆 get_regularization_term()（你应该添加到训练中的损失）和 get_cost()（如果应用提议的结构，网络的预计成本）。

4. 训练模型。

   注意：我们建议在此步骤中使用固定的学习率（不衰减），但这并非严格必要。

5. 使用 StructureExporter 保存提议的模型结构。

   导出的文件为 JSON 格式。需要注意的是，随着训练的进行，提议的模型结构会发生变化。虽然没有明确的停止时机指导，但通常建议等到通过摘要报告的正则化损失趋于稳定时再停止。

6. （可选）创建摘要操作，以便通过 TensorBoard 监控训练进度。

7. 使用 StructureExporter 的输出修改你的模型。

8. 从头开始重新训练模型，不再使用 MorphNet 正则化器。

   注意：所有超参数（如学习率调度）都应使用标准值。

9. （可选）均匀扩展网络，以按需调整准确率与成本之间的权衡。或者，这一步也可以在结构学习之前完成。

我们将第一轮训练称为“结构学习”，第二轮称为“重训”。

总结一下，MorphNet 的关键超参数包括：

• 正则化强度
• 存活阈值

需要注意的是，正则化器类型并不是超参数，因为它完全由你关注的指标（FLOPs、延迟）以及是否使用 BatchNorm 决定。

正则化器类型

正则化器类位于 network_regularizers/ 目录下。它们的命名基于所使用的算法以及试图最小化的目标成本。例如，LogisticSigmoidFlopsRegularizer 使用逻辑斯蒂- sigmoid 概率方法来正则化 FLOPs 成本，而 GammaModelSizeRegularizer 则利用批归一化中的 gamma 参数来正则化模型大小成本。

正则化器算法

• [新] LogisticSigmoid 适用于控制任何类型的模型，但需要在你的模型中添加简单的“门控层”。
• GroupLasso 专为不含批归一化的模型设计。
• Gamma 专为含有批归一化的模型设计；它要求批归一化的缩放参数必须启用。

正则化器目标成本

• FLOPs 以推理网络的 FLOP 计数为目标。
• 模型大小 以网络的权重数量为目标。
• 延迟 根据特定硬件特性，优化网络的估计推理延迟。

示例

添加 FLOPs 正则化器

下面的示例展示了如何使用 MorphNet 来减少模型中的 FLOPs 数量。在本例中，正则化器将从 logits 开始遍历计算图，并且不会跨越任何在图中位于 inputs 或 labels 之前的算子；这允许指定 MorphNet 需要优化的子图。

from morph_net.network_regularizers import flop_regularizer
from morph_net.tools import structure_exporter

def build_model(inputs, labels, is_training, ...):
  gated_relu = activation_gating.gated_relu_activation()

  net = tf.layers.conv2d(inputs, kernel=[5, 5], num_outputs=256)
  net = gated_relu(net, is_training=is_training)

  ...
  ...

  net = tf.layers.conv2d(net, kernel=[3, 3], num_outputs=1024)
  net = gated_relu(net, is_training=is_training)

  logits = tf.reduce_mean(net, [1, 2])
  logits = tf.layers.dense(logits, units=1024)
  return logits

inputs, labels = preprocessor()
logits = build_model(inputs, labels, is_training=True, ...)

network_regularizer = flop_regularizer.LogisticSigmoidFlopsRegularizer(
    output_boundary=[logits.op],
    input_boundary=[inputs.op, labels.op],
    alive_threshold=0.1  # 值范围为 [0, 1]。此默认值在大多数情况下效果良好。
)
regularization_strength = 1e-10
regularizer_loss = (network_regularizer.get_regularization_term() * regularization_strength)

model_loss = tf.nn.sparse_softmax_cross_entropy_with_logits(labels, logits)

optimizer = tf.train.MomentumOptimizer(learning_rate=0.01, momentum=0.9)

train_op = optimizer.minimize(model_loss + regularizer_loss)

您应该通过 TensorBoard 监控结构学习训练的进展。特别是，您应考虑添加一个摘要，用于计算当前的 MorphNet 正则化损失以及采用当前提议结构时的成本。

tf.summary.scalar('RegularizationLoss', regularizer_loss)
tf.summary.scalar(network_regularizer.cost_name, network_regularizer.get_cost())

regularization_strength 的值越大，最终收敛到的有效 FLOP 计数就越小。如果 regularization_strength 足够大，FLOP 计数将降为零。相反，如果它足够小，FLOP 计数将保持初始值，网络结构也不会发生变化。regularization_strength 参数是您用来控制希望在价格性能曲线上所处位置的调节旋钮。alive_threshold 参数则用于确定激活是否处于“存活”状态。

提取 MorphNet 学习到的架构

在训练过程中，您应保存一个包含网络学习到的结构的 JSON 文件，即由 MorphNet 保留（而非移除）的某一层中激活的数量。

exporter = structure_exporter.StructureExporter(
    network_regularizer.op_regularizer_manager)

with tf.Session() as sess:
  tf.global_variables_initializer().run()
  for step in range(max_steps):
    _, structure_exporter_tensors = sess.run([train_op, exporter.tensors])
    if (step % 1000 == 0):
      exporter.populate_tensor_values(structure_exporter_tensors)
      exporter.create_file_and_save_alive_counts(train_dir, step)

其他

联系方式：morphnet@google.com

维护者

• Elad Eban，GitHub: eladeban
• Andrew Poon，GitHub: ayp-google
• Yair Movshovitz-Attias，GitHub: yairmov
• Max Moroz，GitHub: pkch

贡献者

• Ariel Gordon，GitHub: gariel-google。

MorphNet 快速上手指南

MorphNet 是一种在训练过程中学习深度网络结构的方法。它通过添加正则化项，根据资源约束（如 FLOPs、模型大小或延迟）自动调整卷积层的输出通道数量，从而实现模型压缩和架构搜索。

环境准备

• 操作系统: Linux, macOS, Windows (需配置好 TensorFlow 环境)
• Python 版本: 推荐 Python 3.6 - 3.8 (取决于 TensorFlow 版本兼容性)
• 核心依赖:
  • TensorFlow 1.x (官方示例基于 TF1，若使用 TF2 需开启兼容模式或自行适配)
  • NumPy
  • Matplotlib (用于可视化)
• 硬件要求: 支持 CUDA 的 NVIDIA GPU (可选，但推荐用于加速训练)

注意: 由于 MorphNet 主要基于 TensorFlow 1.x 开发，请确保您的环境中已正确安装 tensorflow-gpu (TF1 版本)。国内用户可使用清华源加速安装：

pip install tensorflow-gpu==1.15.0 -i https://pypi.tuna.tsinghua.edu.cn/simple

安装步骤

目前 MorphNet 未作为标准包发布在 PyPI 上，建议直接从 GitHub 源码安装。

1. 克隆仓库

   git clone https://github.com/google-research/morphnet.git
   cd morphnet
   

2. 安装依赖 将当前目录添加到 Python 路径，或安装为可编辑模式：

   pip install -e .
   

   (如果目录下没有 setup.py，请直接将 morphnet 文件夹复制到您的项目目录中，并确保能导入 morph_net 模块)

基本使用

以下是使用 MorphNet 压缩模型 FLOPs 的最简流程。核心思想是：结构学习训练 -> 导出结构 -> 重构并重新训练。

第一步：构建带正则化的模型

在定义模型时，需要在激活函数后插入概率门控层（Gating Layer），并选择 LogisticSigmoid 正则化器（推荐方式）。

import tensorflow as tf
from morph_net.network_regularizers import flop_regularizer
from morph_net.tools import structure_exporter
from morph_net.framework import activation_gating

def build_model(inputs, labels, is_training):
  # 1. 定义门控激活函数 (Probabilistic Gating)
  gated_relu = activation_gating.gated_relu_activation()

  # 构建网络示例
  net = tf.layers.conv2d(inputs, filters=256, kernel_size=[5, 5])
  net = gated_relu(net, is_training=is_training)

  # ... 中间层 ...

  net = tf.layers.conv2d(net, filters=1024, kernel_size=[3, 3])
  net = gated_relu(net, is_training=is_training)

  logits = tf.reduce_mean(net, [1, 2])
  logits = tf.layers.dense(logits, units=10)
  return logits

# 准备数据占位符
inputs = tf.placeholder(tf.float32, shape=[None, 224, 224, 3])
labels = tf.placeholder(tf.int32, shape=[None])

logits = build_model(inputs, labels, is_training=True)

# 2. 初始化正则化器
# output_boundary: 从哪个节点开始向后遍历
# input_boundary: 遇到哪些节点停止遍历 (通常不包含输入和标签)
network_regularizer = flop_regularizer.LogisticSigmoidFlopsRegularizer(
    output_boundary=[logits.op],
    input_boundary=[inputs.op, labels.op],
    alive_threshold=0.1  # 阈值，决定通道是否被保留
)

# 3. 组合损失函数
regularization_strength = 1e-10  # 正则化强度，需根据初始 FLOPs 调整 (约 1/initial_flops)
regularizer_loss = network_regularizer.get_regularization_term() * regularization_strength

base_loss = tf.nn.sparse_softmax_cross_entropy_with_logits(labels=labels, logits=logits)
total_loss = base_loss + regularizer_loss

# 4. 定义优化器 (建议使用固定学习率)
optimizer = tf.train.MomentumOptimizer(learning_rate=0.01, momentum=0.9)
train_op = optimizer.minimize(total_loss)

# 5. (可选) 添加 TensorBoard 监控
tf.summary.scalar('RegularizationLoss', regularizer_loss)
tf.summary.scalar('EstimatedFLOPs', network_regularizer.get_cost())

第二步：训练并导出结构

在训练循环中，定期保存当前的网络结构配置到 JSON 文件。

exporter = structure_exporter.StructureExporter(
    network_regularizer.op_regularizer_manager)

with tf.Session() as sess:
  sess.run(tf.global_variables_initializer())
  
  for step in range(max_steps):
    # 执行训练
    _, summary_str, exporter_tensors = sess.run(
        [train_op, tf.summary.merge_all(), exporter.tensors]
    )
    
    # 每 1000 步保存一次结构
    if step % 1000 == 0:
      exporter.populate_tensor_values(exporter_tensors)
      # 生成包含各层存活通道数的 JSON 文件
      exporter.create_file_and_save_alive_counts('./train_log', step)
      
    # 写入 TensorBoard
    # writer.add_summary(summary_str, step)

第三步：重构与重训练

1. 解析 JSON: 读取上一步生成的 JSON 文件，获取每一层需要保留的通道数量。
2. 修改模型: 根据 JSON 中的通道数，手动修改原始模型代码中 conv2d 等层的 filters 参数，移除所有的门控层（gating layers）和 MorphNet 正则化项。
3. 从头训练: 使用标准的超参数（如学习率衰减策略）对裁剪后的新模型进行从头训练（Retraining），以恢复精度。

提示: regularization_strength 是控制压缩率的关键旋钮。值越大，最终模型的 FLOPs 越小；值过小则结构不会发生变化。建议在对数尺度上搜索该参数（例如围绕 1/初始 FLOPs 上下浮动几个数量级）。