morph-net - 在训练过程中优化深度网络结构的方法

MorphNet 项目介绍

项目概述

MorphNet 是一种在训练过程中学习深度网络结构的方法，其核心理念是将网络结构学习问题进行连续松弛。简单来说，MorphNet 通过正则化器压缩网络中过滤器的影响，当它们变得足够小时，相应的输出通道会被标记为可从网络中移除。

MorphNet 和传统的神经网络优化不同，通过加入一些正则化器，专门针对消耗特定资源（如浮点运算次数FLOPs或模型大小）的部分进行约束优化。在训练损失中加入正则化损失，并通过随机梯度下降或相似的优化器最小化它们的和，形成了一个约束优化问题。这个方法最早被介绍在 CVPR 2018 的论文中。

新进展：FiGS（精细化随机架构搜索）

FiGS 是一种基于概率的通道正则化方法，是当前推荐使用 MorphNet 的方法。它可以作为剪枝算法使用，也可以用作完整的可微分架构搜索的方法。在 MorphNet 的文档中通常被称为 LogisticSigmoid 正则化器。

如何使用 MorphNet

设想你有一个可用于图像分类的卷积神经网络，并希望缩小模型以满足某些约束（如内存或延迟）。在已存在模型（“种子网络”）和目标标准下，MorphNet 会通过调整每个卷积层的输出通道数提出一个新的模型。

注意，MorphNet不会改变网络的拓扑结构——所提议的模型将与种子网络拥有相同的层数和连接模式。

使用步骤如下：

选择正则化器
- 从 morphnet.network_regularizers 中选择正则化器。选择基于你的目标成本（如 FLOPs 或延迟）以及你是否能够向模型添加新层：
  - 如果能够，推荐使用基于 LogisticSigmoid 的正则化器。
  - 如果不能，根据网络架构选择正则化器类型。
初始化正则化器
- 使用阈值以及模型的输出和可选输入边界操作初始化正则化器。正则化器将从输出边界开始分析图，并在遇到输入边界操作时停止。阈值决定哪些输出通道可以被移除。
加入正则化损失
- 将正则化项加入到训练损失中。建议在 1/(初始成本) 周围的对数刻度上搜索正则化强度。
训练模型
- 建议使用固定的学习率进行该步骤的训练。
保存模型结构
- 使用 StructureExporter 保存模型结构，导出的文件为 JSON 格式。
监控训练进程
- 可选择通过 TensorBoard 创建 summary 操作以监控训练进展。
调整和重训练模型
- 修改你的模型结构后，从头开始重新训练。
从新调整网络
- 可选择在结构学习步骤之前或之后均匀扩展网络，以调整精度与成本的权衡。

正则化器类型

LogisticSigmoid: 适用于任何模型类型，但需要在模型中添加简单的“门控层”。
GroupLasso: 适用于无批归一化的模型。
Gamma: 针对具有批归一化的模型，需要启用批归一化的缩放。

目标成本

浮点运算次数 (Flops): 针对推理网络的 FLOPs 计数优化。
模型大小: 针对网络权重数量优化。
延迟: 基于具体硬件特性优化网络的估算推理延迟。

例子：添加 FLOPs 正则化器

以下示例展示了如何使用 MorphNet 减少模型中的 FLOPs 数量：

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)

保存和提取架构

在训练过程中，保存包含网络学习到的结构的 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)

项目维护者和贡献者

维护者:
- Elad Eban
- Andrew Poon
- Yair Movshovitz-Attias
- Max Moroz
贡献者:
- Ariel Gordon

总体来说，MorphNet 是一项用于优化神经网络结构，特别是在资源受限场景下的实用工具，通过其灵活的正则化器和结构学习方法，可以帮助用户有效地管理和优化模型资源使用。