作者：Arjun Gopalan，Google Research 软件工程师

编辑：Robert Crowe，Google Research TensorFlow 开发者倡导者

简介

神经结构化学习 (NSL) 是 TensorFlow 中的一个框架，可用于训练具有结构化信号的神经网络。它以两种方式处理结构化输入：(i) 作为显式图，或 (ii) 作为隐式图，其中邻居在模型训练期间动态生成。具有显式图的 NSL 通常用于 神经图学习，而具有隐式图的 NSL 通常用于 对抗性学习。这两种技术都作为 NSL 框架中的一种正则化形式实现。因此，它们只会影响训练工作流程，因此模型服务工作流程保持不变。在本帖的其余部分，我们将主要关注如何在 TFX 中使用 NSL 框架实现图正则化。

使用 NSL 构建图正则化模型的高级工作流程涉及以下步骤

1. 如果不存在图，则构建一个图。
2. 使用图和输入示例特征来增强训练数据。
3. 使用增强的训练数据将图正则化应用于给定模型。

这些步骤不会直接映射到现有的 TFX 管道组件。但是，TFX 支持 自定义组件，允许用户在他们的 TFX 管道中实现自定义处理。有关 TFX 中自定义组件的介绍，请参阅 这篇博文。因此，要在 TFX 中创建包含上述步骤的图正则化模型，我们将使用额外的自定义 TFX 组件。

为了说明使用 NSL 的示例 TFX 管道，让我们考虑对 IMDB 数据集 进行情感分类的任务。一个基于 colab 的教程演示了使用 NSL 完成此任务（使用原生 TensorFlow），可在 此处 获得，我们将使用它作为我们 TFX 管道示例的基础。

使用自定义 TFX 组件进行图正则化

为了在此任务的 TFX 中构建图正则化的 NSL 模型，我们将使用 自定义 Python 函数 方法定义三个自定义组件。以下是我们示例中使用这些自定义组件的 TFX 管道的示意图。为了简洁起见，我们跳过了通常在 Trainer 组件之后出现的组件，如 Evaluator、Pusher 等。

图 1：使用图正则化的文本分类的示例 TFX 管道

在此图中，只有自定义组件（以粉红色显示）和 Graph-regularized Trainer 组件具有与 NSL 相关的逻辑。值得注意的是，此处显示的自定义组件仅供说明，可能可以用其他方式构建功能等效的管道。现在，我们将更详细地描述每个自定义组件，并显示它们的代码片段。

IdentifyExamples

此自定义组件为每个训练示例分配一个唯一的 ID，该 ID 用于将每个训练示例与其从图中对应的邻居关联起来。

 
@component
def IdentifyExamples(
    orig_examples: InputArtifact[Examples],
    identified_examples: OutputArtifact[Examples],
    id_feature_name: Parameter[str],
    component_name: Parameter[str]
  ) -> None:

  # Compute the input and output URIs.
  ...

  # For each input split, update the TF.Examples to include a unique ID.
  with beam.Pipeline() as pipeline:
    (pipeline
     | 'ReadExamples' >> beam.io.ReadFromTFRecord(
         os.path.join(input_dir, '*'),
         coder=beam.coders.coders.ProtoCoder(tf.train.Example))
     | 'AddUniqueId' >> beam.Map(make_example_with_unique_id, id_feature_name)
     | 'WriteIdentifiedExamples' >> beam.io.WriteToTFRecord(
         file_path_prefix=os.path.join(output_dir, 'data_tfrecord'),
         coder=beam.coders.coders.ProtoCoder(tf.train.Example),
         file_name_suffix='.gz'))

  identified_examples.split_names = orig_examples.split_names
  return

make_example_with_unique_id() 函数更新给定示例以包含包含唯一 ID 的附加特征。

SynthesizeGraph

如上所述，在 IMDB 数据集中，没有显式图作为输入给出。因此，在我们演示图正则化之前，我们将构建一个图。对于此示例，我们将使用预训练的文本嵌入模型将电影评论中的原始文本转换为嵌入，然后使用生成的嵌入来构建图。

SynthesizeGraph 自定义组件处理我们示例的图构建，并注意到它定义了一个名为 SynthesizedGraph 的新 Artifact，这将是此自定义组件的输出。

 
"""Custom Artifact type"""
class SynthesizedGraph(tfx.types.artifact.Artifact):
  """Output artifact of the SynthesizeGraph component"""
  TYPE_NAME = 'SynthesizedGraphPath'
  PROPERTIES = {
      'span': standard_artifacts.SPAN_PROPERTY,
      'split_names': standard_artifacts.SPLIT_NAMES_PROPERTY,
  }

@component
def SynthesizeGraph(
    identified_examples: InputArtifact[Examples],
    synthesized_graph: OutputArtifact[SynthesizedGraph],
    similarity_threshold: Parameter[float],
    component_name: Parameter[str]
  ) -> None:

  # Compute the input and output URIs
  ...

  # We build a graph only based on the 'train' split which includes both
  # labeled and unlabeled examples.
  create_embeddings(train_input_examples_uri, output_graph_uri)
  build_graph(output_graph_uri, similarity_threshold)
  synthesized_graph.split_names = artifact_utils.encode_split_names(
      splits=['train'])
  return

create_embeddings() 函数涉及使用 TensorFlow Hub 上的一些预训练模型将电影评论中的文本转换为相应的嵌入。build_graph() 函数涉及调用 NSL 中的 build_graph() API。

GraphAugmentation

此自定义组件的目的是将示例特征（电影评论中的文本）与从嵌入构建的图结合起来，以生成增强的训练数据集。生成的训练示例将包括来自其对应邻居的特征。

@component
def GraphAugmentation(
    identified_examples: InputArtifact[Examples],
    synthesized_graph: InputArtifact[SynthesizedGraph],
    augmented_examples: OutputArtifact[Examples],
    num_neighbors: Parameter[int],
    component_name: Parameter[str]
  ) -> None:

  # Compute the input and output URIs
  ...

  # Separate out the labeled and unlabeled examples from the 'train' split.
  train_path, unsup_path = split_train_and_unsup(train_input_uri) 

  # Augment training data with neighbor features.
  nsl.tools.pack_nbrs(
    train_path, unsup_path, graph_path, output_path, add_undirected_edges=True,
    max_nbrs=num_neighbors
  )

  # Copy the 'test' examples from input to output without modification.
  ...

  augmented_examples.split_names = identified_examples.split_names
  return

split_train_and_unsup() 函数涉及将输入示例拆分为带标签和未带标签的示例，pack_nbrs() NSL API 创建增强的训练数据集。

Graph-regularized Trainer

现在我们已经实现了所有自定义组件，TFX 管道中对 NSL 的其余特定添加是在 Trainer 组件中。以下是对图正则化 Trainer 组件的简化视图。

 
 ...

  estimator = tf.estimator.Estimator(
       model_fn=feed_forward_model_fn, config=run_config, params=HPARAMS)
  
  # Create a graph regularization config.
  graph_reg_config = nsl.configs.make_graph_reg_config(
      max_neighbors=HPARAMS.num_neighbors,
      multiplier=HPARAMS.graph_regularization_multiplier,
      distance_type=HPARAMS.distance_type,
      sum_over_axis=-1)
  
  # Invoke the Graph Regularization Estimator wrapper to incorporate
  # graph-based regularization for training.
  graph_nsl_estimator = nsl.estimator.add_graph_regularization(
      estimator,
      embedding_fn,
      optimizer_fn=optimizer_fn,
      graph_reg_config=graph_reg_config)

 ...
 

如您所见，一旦创建了基础模型（在本例中为前馈神经网络），通过调用 NSL 包装器 API，可以轻松地将其转换为图正则化模型。

就这样！现在我们拥有构建 TFX 中的图正则化 NSL 模型所需的所有缺失部分。一个基于 colab 的教程演示了在 TFX 中端到端地完成此示例，可在 此处 获得。随意尝试它并根据需要进行自定义！

对抗性学习

如上文介绍中所述，神经结构化学习的另一个方面是对抗性学习，其中，不是使用来自图的显式邻居进行正则化，而是动态地以对抗的方式创建隐式邻居来混淆模型。因此，使用对抗性示例进行正则化是提高模型鲁棒性的有效方法。使用 NSL 进行对抗性学习可以轻松地集成到 TFX 管道中。它不需要任何自定义组件，只需要更新训练器组件以调用 NSL 中的对抗性正则化包装器 API。

总结

我们演示了如何使用自定义组件在 TFX 中构建具有 NSL 的图正则化模型。当然也可以通过其他方式构建图，以及以不同的方式构建整体管道。我们希望此示例为您的 NSL 工作流程提供基础。

其他链接

有关 NSL 的更多信息，请查看以下资源

• TFX 中的 NSL colab 教程
• NSL 网站
• NSL GitHub
• 更多 NSL 教程和视频

致谢

我们要感谢 Google 的神经结构化学习和 TFX 团队，以及 Aurélien Geron 对他们的支持和贡献。