如何在 TFX 中使用 NSL 框架實(shí)現(xiàn)計(jì)算圖正則化

神經(jīng)架構(gòu)學(xué)習(xí) (Neural Structured Learning，NSL) 是 TensorFlow 中的一個(gè)框架，可以利用結(jié)構(gòu)化信號(hào)來(lái)訓(xùn)練神經(jīng)網(wǎng)絡(luò)。這種框架接受 (i) 顯式計(jì)算圖或 (ii) 隱式計(jì)算圖，處理結(jié)構(gòu)化輸入，并在模型訓(xùn)練過(guò)程中動(dòng)態(tài)生成鄰接點(diǎn)(Neighbors)。顯式計(jì)算圖的 NSL 通常用于基于神經(jīng)網(wǎng)絡(luò)的圖學(xué)習(xí)(Neural Graph Learning)，而隱式計(jì)算圖的 NSL 通常用于 對(duì)抗學(xué)習(xí)。這兩種技術(shù)均以 NSL 框架中的正則化形式實(shí)現(xiàn)。所以，它們只對(duì)訓(xùn)練工作流有影響，而工作流的模型保持不變。我們將在本文中探討如何在 TFX 中使用 NSL 框架實(shí)現(xiàn)計(jì)算圖正則化(Graph Regularization )。

神經(jīng)架構(gòu)學(xué)習(xí)
https://tensorflow.google.cn/neural_structured_learning

使用 NSL 構(gòu)建計(jì)算圖正則化模型的高級(jí)工作流包含以下步驟：

如果沒有可用的計(jì)算圖，則需要先構(gòu)建一個(gè)計(jì)算圖。

使用計(jì)算圖和輸入樣本特征擴(kuò)充訓(xùn)練數(shù)據(jù)。

使用擴(kuò)充的訓(xùn)練數(shù)據(jù)對(duì)給定模型進(jìn)行計(jì)算圖正則化。

這些步驟不會(huì)立即映射到現(xiàn)有的 TFX (TensorFlow Extended) 流水線組件上。但是，TFX 支持自定義組件，允許用戶在其 TFX 流水線中實(shí)現(xiàn)自定義處理。如需了解 TFX 中的自定義組件，請(qǐng)參閱這篇文章。

TFX
https://tensorflow.google.cn/tfx/guide#tfx_standard_components

自定義組件
https://tensorflow.google.cn/tfx/guide/understanding_custom_components

為了在 TFX 中創(chuàng)建一個(gè)包含上述步驟的計(jì)算圖正則化模型，我們將利用擴(kuò)展自定義 TFX 組件。

為展示如何使用 NSL，我們構(gòu)建了一個(gè)示例 TFX 流水線，對(duì) IMDB 數(shù)據(jù)集進(jìn)行情感分類。我們提供了一個(gè)基于 Colab 的教程，演示了如何使用 NSL 與原生 TensorFlow 來(lái)完成這項(xiàng)任務(wù)，我們以此作為示例 TFX 流水線的基礎(chǔ)。

IMDB 數(shù)據(jù)集
https://tensorflow.google.cn/datasets/catalog/imdb_reviews

Colab 教程
https://tensorflow.google.cn/neural_structured_learning/tutorials/graph_keras_lstm_imdb

自定義 TFX 組件的計(jì)算圖正則化

為了在 TFX 中構(gòu)建一個(gè)計(jì)算圖正則化的 NSL 模型來(lái)完成這項(xiàng)任務(wù)，我們將使用自定義 Python 函數(shù)方法自定義三個(gè)組件。以下是使用這些自定義組件實(shí)現(xiàn)我們示例的 TFX 流水線示意圖。為了簡(jiǎn)潔起見，我們省略了通常在 Trainer 之后的組件，例如 Evaluator、 Pusher 等。

自定義 Python 函數(shù)
https://tensorflow.google.cn/tfx/guide/custom_function_component

圖 1：TFX 流水線示例：使用計(jì)算圖正則化進(jìn)行文本分類

在此圖中，僅有自定義組件（粉色）與計(jì)算圖正則化的 Trainer組件具備 NSL 相關(guān)邏輯。值得注意的是，此處展示的自定義組件僅作例證，還可以通過(guò)其他方式構(gòu)建類似功能的流水線。接下來(lái)，我們進(jìn)一步詳細(xì)描述各個(gè)自定義組件，并展示相應(yīng)的代碼段。

IdentifyExamples

此自定義組件為每個(gè)訓(xùn)練樣本分配一個(gè)唯一的 ID，將每個(gè)訓(xùn)練樣本與其在計(jì)算圖中相應(yīng)的鄰接點(diǎn)關(guān)聯(lián)起來(lái)。

@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() 函數(shù)可以更新給定樣本，將包含唯一 ID 的額外特征包括在內(nèi)。

SynthesizeGraph

如上所述，在 IMDB 數(shù)據(jù)集中，沒有提供顯式計(jì)算圖作為輸入。因此，在演示計(jì)算圖正則化之前，我們將構(gòu)建一個(gè)計(jì)算圖。在此示例中，我們使用一個(gè)預(yù)訓(xùn)練的文本嵌入向量模型將電影評(píng)論中的原始文本轉(zhuǎn)換為嵌入向量，然后通過(guò)生成的嵌入向量構(gòu)建計(jì)算圖。

SynthesizeGraph自定義組件負(fù)責(zé)處理計(jì)算圖構(gòu)建，請(qǐng)注意，它定義了一個(gè)新的Artifact，名為 SynthesizedGraph，作為此自定義組件的輸出。

"""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() 函數(shù)通過(guò) TensorFlow Hub 上的一些預(yù)訓(xùn)練模型將電影評(píng)論中的文本轉(zhuǎn)換為相應(yīng)的嵌入向量。build_graph() 函數(shù)調(diào)用 NSL 中的 build_graph()API。

TensorFlow Hub
https://tensorflow.google.cn/hub

build_graph()
https://tensorflow.google.cn/neural_structured_learning/api_docs/python/nsl/tools/build_graph

GraphAugmentation

此自定義組件的目的在于將樣本特征（電影評(píng)論中的文本）與通過(guò)嵌入向量構(gòu)建的計(jì)算圖結(jié)合起來(lái)，生成一個(gè)擴(kuò)充的訓(xùn)練數(shù)據(jù)集。由此得出的訓(xùn)練樣本也將包括其相應(yīng)鄰接點(diǎn)的特征。

@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() 函數(shù)將輸入樣本拆分成帶標(biāo)簽和無(wú)標(biāo)簽的樣本，pack_nbrs() NSL API 創(chuàng)建擴(kuò)充的訓(xùn)練數(shù)據(jù)集。

pack_nbrs()
https://tensorflow.google.cn/neural_structured_learning/api_docs/python/nsl/tools/pack_nbrs

計(jì)算圖正則化的 Trainer

我們目前所有的自定義組件都已實(shí)現(xiàn)，TFX 流水線的Trainer 組件中增加了其他 NSL 相關(guān)的內(nèi)容。下方展示了一個(gè)計(jì)算圖正則化 Trainer 組件的簡(jiǎn)化視圖。

... 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) ...

如您所見，創(chuàng)建了基礎(chǔ)模型后（本例中指一個(gè)前饋神經(jīng)網(wǎng)絡(luò)），就可以通過(guò)調(diào)用 NSL 封裝容器 API 將其直接轉(zhuǎn)換為計(jì)算圖正則化模型。

一切就這么簡(jiǎn)單?，F(xiàn)在，我們已補(bǔ)充完在 TFX 中構(gòu)建計(jì)算圖正則化 NSL 模型所需的步驟。此處提供了一個(gè)基于 Colab 的教程，在 TFX 中端到端地演示這個(gè)示例。不妨嘗試一下，并可根據(jù)您的需要進(jìn)行自定義。

此處
https://tensorflow.google.cn/tfx/tutorials/tfx/neural_structured_learning

對(duì)抗學(xué)習(xí)

如前文簡(jiǎn)介中所述，神經(jīng)架構(gòu)學(xué)習(xí)的另一個(gè)方面是對(duì)抗學(xué)習(xí)，即不使用計(jì)算圖中的顯式鄰接點(diǎn)來(lái)進(jìn)行正則化，而是動(dòng)態(tài)、對(duì)抗性地創(chuàng)建隱式鄰接點(diǎn)來(lái)迷惑模型。

因此，使用對(duì)抗樣本進(jìn)行正則化是提高模型魯棒性的有效方式。使用神經(jīng)架構(gòu)學(xué)習(xí)的對(duì)抗學(xué)習(xí)可以輕松集成到 TFX 流水線中。無(wú)需任何自定義組件，只需要更新 Trainer 組件即可在神經(jīng)架構(gòu)學(xué)習(xí)中調(diào)用對(duì)抗正則化封裝容器 API。

總結(jié)

我們演示了如何利用自定義組件在 TFX 中使用神經(jīng)架構(gòu)學(xué)習(xí)構(gòu)建計(jì)算圖正則化模型。當(dāng)然，也可以用其他方式構(gòu)建計(jì)算圖，或者按照另外的方式來(lái)構(gòu)建整體流水線。

我們希望這個(gè)示例能夠?yàn)槟鷺?gòu)建自己的神經(jīng)架構(gòu)學(xué)習(xí)工作流提供幫助。

相關(guān)鏈接

有關(guān)神經(jīng)架構(gòu)學(xué)習(xí)的更多信息，請(qǐng)查閱以下資源：

TFX 中的 NSL Colab 教程
https://tensorflow.google.cn/tfx/tutorials/tfx/neural_structured_learning

NSL 網(wǎng)站
https://tensorflow.google.cn/neural_structured_learning

NSL GitHub
https://github.com/tensorflow/neural-structured-learning

更多 NSL 教程與視頻
https://github.com/tensorflow/neural-structured-learning#videos-and-colab-tutorials

致謝：

我們特此鳴謝 Google 神經(jīng)架構(gòu)學(xué)習(xí)團(tuán)隊(duì)、TFX 團(tuán)隊(duì)以及 Aurélien Geron 的支持與貢獻(xiàn)。

責(zé)任編輯：xj

原文標(biāo)題：TensorFlow Extended + 神經(jīng)架構(gòu)學(xué)習(xí)：構(gòu)建計(jì)算圖正則化模型

文章出處：【微信公眾號(hào)：TensorFlow】歡迎添加關(guān)注！文章轉(zhuǎn)載請(qǐng)注明出處。