オブジェクトの広場はオージス総研グループのエンジニアによる技術発表サイトです

AI

はじめての自然言語処理

第8回 続・T5 によるテキスト生成の検証
オージス総研 技術部 アドバンストテクノロジセンター
鵜野 和也
2020年4月23日

前回の続編ということで、t5 (T5 の実装コード) の構成や Mesh Tensorflow について説明し、OSCAR データセットでの実験結果や複数 GPU 環境の Tensorflow Serving へのエクスポートと並列実行などを試していきます。

1. はじめに

本記事は前回の続編です。 未読の方は右上のリンクから第7回の記事に目を通して戻ってきて頂けると、より理解がしやすいと思います。今回は、Google の T5(Text-to-Text Transfer Transformer) 1によるテキスト生成について、その実装コード2 の構成について説明し、Wikipedia 日本語版よりも規模の大きい OSCAR データセット3を用いた事前学習の結果や、学習済みモデルを複数 GPU 環境の Tensorflow Serving 4 にエクスポートし並列動作させる方法について解説していきます。また、前回に引き続き、"T5" と記述すれば 1 の論文を、"t5" と記述すれば 2 のソースコードやライブラリを指すものとします。

それでは、まずは t5 の構成について紹介していきます。

2. t5 : Text-To-Text Transfer Transformer の実装

t5 は data, evaluation, models の3つのサブパッケージで構成されています。順にその概要を見ていきましょう。

t5.data

t5 には Task 及びその混合物である Mixture という概念があり、学習時には Task や Mixture の名前を指定して実行します。Task は Tensorflow Datasets(以後、単に“TFDS”) 5に用意されているデータセットやテキストファイルからデータを読み込み、テキストベースの前処理、 Sentencepiece によるトークナイズ、事前学習に用いるノイズの挿入等を行って、tf.data.Dataset として返します。t5.data ではこれらの処理の実装が提供されています。

これだけではよくわからないので、前回のコードスニペットを振り返ってみましょう。このコード自体は t5 に含まれる Wikipedia 英語版での学習タスクの定義をコピーして、'en' を 'ja' に書き換え、 Sentencepiece パスを修正しただけのものです。 Wikipedia 日本語版も TFDS に登録がある為、データセットの名称を合わせるだけで、簡単に組み込むことが出来ました。

SPM_PATH = "gs://somewhere/t5/sentencepiece/wikipedia_20190301_ja_v003.model"
import functools
from t5.data import preprocessors
from t5.data.utils import TaskRegistry
from t5.data.utils import MixtureRegistry
from t5.data.utils import TfdsTask

task_name_wikipedia_ja = "wikipedia_20190301.ja_v003_unsupervised"

TaskRegistry.add(
    task_name_wikipedia_ja,
    TfdsTask,
    tfds_name="wikipedia/20190301.ja:0.0.3",
    text_preprocessor=functools.partial(                                     # ①
        preprocessors.rekey, key_map={"inputs": None, "targets": "text"}),
    token_preprocessor=preprocessors.unsupervised,                           # ③
    sentencepiece_model_path=SPM_PATH,                                       # ②
    metric_fns=[])                                                           # ④

MixtureRegistry.add(task_name_wikipedia_ja, [(task_name_wikipedia_ja, 1.0)]) # ⑤

① テキストレベルの前処理としてキーの付け替え

TFDS の Wikipedia 日本語版には属性として、title (記事タイトル)、 text (記事本文)が含まれています。この text を属性名 targets に付け替えます。この targets で後続の処理が待っているので、ここで合わせ込むわけです。 inputs に None を設定していますが、これは使用しません。後述する処理で targets から inputs が生成されます。

② Sentencepiece モデルの指定

① の targets に対して、ここで指定したモデルを使用して Sentencepiece によるトークン化が行われます。

③ トークン化後のデータに対するノイズ挿入等の前処理

この Task は事前学習のタスクなので、②のトークン化後の前処理として preprocessors.unsupervised を指定しています。この関数は gin で preprocessors パラメータに登録されたプリプロセッサのリストを順に動かすだけです6。

せっかくなので、この辺りの構成についても確認してみましょう。前回の事前学習ではこの辺りの設定として gin ファイルの objectives/span_u_u.gin を指定していますが、このファイルは objectives/span.gin を include するだけなので、そちらの該当部分を抜粋します7。

preprocessors.unsupervised.preprocessors = [
@preprocessors.select_random_chunk,
@preprocessors.reduce_concat_tokens,
@preprocessors.split_tokens,
@preprocessors.denoise,
]

処理の順番としては、

  1. select_random_chunk で target 属性から指定長(65536トークン)をランダムに切り出し（指定長に満たない場合はそのまま）。
  2. reduce_concat_tokens で指定したバッチサイズ(128サンプル)で複数サンプルを連結。イメージ的には [128, 65536] ではなく [128 * 65536] になります(ただし 1. で指定長に満たないサンプルがあるのでその分短くなります)。
  3. split_tokens で連結した複数サンプルを指定した固定長(preprocessors.random_spans_helper とそのパラメータの設定値で計算された値。今回は 568)で再分割。ここまで全て targets 属性に対する処理です。
  4. denoise で 3. の target 属性にノイズを挿入した inputs 属性と、ノイズ部分の正解を示す targets 属性を出力。

日本語で書いても伝わる気がしないので、処理を把握しやすいよう最大長等を小さい値にして実行したコードサンプルを入れておきます。トークン化された文書が最終的な学習サンプルに改変される過程が分かるかと思います。

import tensorflow as tf
tf.compat.v1.enable_eager_execution()
from t5.data.preprocessors import select_random_chunk, reduce_concat_tokens, split_tokens

def print_and_copy_data(name, dataset):
  print("\n### %s ###" % name)
  data = []
  for example in dataset:
    data.append(example['targets'].numpy())
    print(" 'targets' : %s" % (data[-1]))
  return tf.data.Dataset.from_tensor_slices({'targets': data})

dataset = tf.data.Dataset.from_tensor_slices({'targets': [
              [10, 11, 12, 13, 14, 15, 16, 17, 18, 19], 
              [20, 21, 22, 23, 24, 25, 26, 27, 28, 29], 
              [30, 31, 32, 33, 34, 35, 36, 37, 38, 39]]}) 
dataset = print_and_copy_data("Tokenized sentences.", dataset)

#### Tokenized sentences. ###
# 'targets' : [10 11 12 13 14 15 16 17 18 19]
# 'targets' : [20 21 22 23 24 25 26 27 28 29]
# 'targets' : [30 31 32 33 34 35 36 37 38 39]

dataset = select_random_chunk(dataset, max_length=5)
dataset = print_and_copy_data("select_random_chunk()", dataset)

#### select_random_chunk() ###
# 'targets' : [15 16 17 18 19]
# 'targets' : [25 26 27 28 29]
# 'targets' : [30 31 32 33 34]

dataset = reduce_concat_tokens(dataset, batch_size=3)
dataset = print_and_copy_data("reduce_concat_tokens()", dataset)

#### reduce_concat_tokens() ###
# 'targets' : [15 16 17 18 19 25 26 27 28 29 30 31 32 33 34]

dataset = split_tokens(dataset, min_tokens_per_segment=None, max_tokens_per_segment=3)
print_and_copy_data("split_tokens()", dataset)

#### split_tokens() ###
# 'targets' : [15 16 17]
# 'targets' : [18 19 25]
# 'targets' : [26 27 28]
# 'targets' : [29 30 31]
# 'targets' : [32 33 34]

この時点で [18 19 25], [29 30 31] は記事の境界を跨っており、それを示す情報がありません。

しかし、[29 30 31] のように境界が記事の末端(29)と先頭(30)であるケースは、文の切れ目で文脈が変わっている形になり、普通によくあることなのでさほど問題ありません。[18 19 25] は記事の境界が文の途中(25)になっていますが、 殆どの記事は select_random_chunk() に設定した最大長(65536)よりも短い為、文が途中で切られるということはなく、このようにトークンの並びが不自然になるのは頻度的にかなりのレアケースなので、気にするほどではないとのことです。

最後の denoise はこんな感じです。

from t5.data.sentencepiece_vocabulary import SentencePieceVocabulary
from t5.data.utils import DEFAULT_SPM_PATH
vocabulary = SentencePieceVocabulary(DEFAULT_SPM_PATH)

from t5.data.preprocessors import noise_span_to_unique_sentinel
from t5.data.preprocessors import random_spans_noise_mask
from t5.data.preprocessors import nonnoise_span_to_unique_sentinel
from t5.data.preprocessors import denoise
import functools

inputs_fn = noise_span_to_unique_sentinel
noise_mask_fn = functools.partial(random_spans_noise_mask, mean_noise_span_length=3.0)
targets_fn = nonnoise_span_to_unique_sentinel

dataset = denoise(dataset, vocabulary, noise_density=0.15, noise_mask_fn=noise_mask_fn, 
    inputs_fn=inputs_fn, targets_fn=targets_fn)

for example in dataset:
    print(" 'inputs' : %s,  'targets' : %s" % (example['inputs'].numpy(), example['targets'].numpy()))

#  'inputs' : [   15    16 31999],  'targets' : [31999    17]
#  'inputs' : [   18    19 31999],  'targets' : [31999    25]
#  'inputs' : [   26    27 31999],  'targets' : [31999    28]
#  'inputs' : [   29    30 31999],  'targets' : [31999    31]
#  'inputs' : [   32    33 31999],  'targets' : [31999    34]

最後の inputs と targets で Text-to-Text の学習をします。

objectives/span.gin には、ここに示した各preprocessorの細かい設定が記述されているので確認して見てください。

④ メトリクス関数の指定

間が空きすぎて④がなんだか覚えてませんよね。TaskRegistory.add(..., metric_fns=[]) の metric_fns の部分です。事前学習では使用しないので空リストを指定しています。次節で補足します。

⑤ Mixture の指定

これは MixtureRegistry.add(task_name_wikipedia_ja, [(task_name_wikipedia_ja, 1.0)]) の部分です。Mixture はマルチタスク学習の為の仕組みですが、前回の実験ではマルチタスク学習を行わなかったので、Wikipedia日本語版の Task を 1.0 のレートで登録しています。

マルチタスク学習ですと、例えば Task として glue と squad があり [('glue', 50.0), ('squad', 100.0)] と指定した場合は、glue と squad の学習サンプルを 1:2 の混合比で混ぜて学習することになります。またレートはタスク名から float を返す関数で指定することも可能です。

次はメトリクス関数についてです。

t5.metrics

t5.metrics には検証モード( eval.gin )での実行時に出力される events.* ファイルをパースしてメトリクス関数を適用する処理や各種メトリクス関数の実装が含まれています。

定義されているメトリクス関数としては以下のようなものがあります。

import t5.evaluation.metrics as metrics
with open(metrics.__file__, 'r') as f:
  lines = f.readlines()
  for line in [line.strip()[4:-1] for line in lines if line.startswith("def ")]:
    print(line)

# bleu(targets, predictions)
# rouge(targets, predictions, score_keys=None)
# span_qa(targets, predictions)
# qa(targets, predictions)
# accuracy(targets, predictions)
# sequence_accuracy(targets, predictions)
# pearson_corrcoef(targets, predictions)
# spearman_corrcoef(targets, predictions)
# matthews_corrcoef(targets, predictions)
# mean_multiclass_f1(num_classes)
# exact_match(targets, predictions)
# f1_score_with_invalid(targets, predictions)
# mean_group_metric(metric_fn, group_key="group", value_key="value")
# multirc_f1_over_all_answers(targets, predictions)

自然言語処理でよく使いそうなメトリクスが用意されていますが、日本語は基本的に考慮外なので必要な場合は前回記事を参考にして自前で用意する必要があります8。

t5.models

t5.models には前述の t5.data, t5.metrics とモデル実装の繋ぎになるコード、各種実験に対応する gin ファイルが含まれます。T5 のモデルは Mesh Tensorflow で実装されており、Encoder-Decoder だけでなく Language model, Prefix LM を含めた全ての構成は Mesh Tensorflow の実装と gin の設定で実現されます。

モデル実装の構成についても前回の gin ファイルを確認してみましょう。この辺りの設定として models/bert_v1.gin を指定しています。

# -*-Python-*-
# with relative attention instead of positional embedding
include 'models/bi_bert_base.gin'

transformer.Unitransformer.positional_embedding = False
# relative attention is implemented as a bias (independent of context)
# this bias differs across attention heads but is shared across different
#  layers in a stack (to save time)
transformer_layers.SelfAttention.relative_attention_type = "bias_shared"

このファイルには “Positional Encoding” に関する設定が記載されています。BERT 的な絶対位置でのエンコーディングをやめて(positional_embedding = False)、Self-Attention での Key - Query の相対位置ベースでの学習とパラメータ共有(各ヘッドで独立した値を全レイヤで共有)に関する指定が記述されています。

models/bi_bert_base.gin が include されてますが、こんな感じです。

# -*-Python-*-
# same hyperparameters as bert_base and xlnet_base

utils.run.model_type = "bitransformer"
d_model = 768
num_layers = 12
d_ff = 3072
num_heads = 12
d_kv = 64

こちらは Encoder-Decoder構成( bitransformer )とスタックのレイヤ数及び各次元の設定です。 BERTBASE に準じた値となっています。

モデル実装の中核は Mesh Tensorflow ですので、ここからは Mesh Tensorflow について見ていきましょう。

3. Mesh Tensorflow

Mesh Tensorflow 9 は分散深層学習向けライブラリで Tensorflow 上のレイヤとして実装されています。

分散深層学習の方式としては、データ並行、モデル並行が考えられます。以下は バッチサイズ b、入力ユニット数 dio、隠れ層ユニット数 dh、出力ユニット数 dio の全結合ネットワークをデータ並行、モデル並行で分割した場合(分割数 n = 2)のイメージ図です。

  • データ並行は全てのパラメータ(w, v)をプロセッサ[0] と [1] に複製して配置、それぞれのプロセッサでバッチの半分のデータを受け持ちます。
  • モデル並行はパラメータ(w, v)の半分をプロセッサ[0]、残り半分を[1]に分割配置し、バッチ中の入力データ1件の計算を複数プロセッサで行います。

data_and_model_parallel

Mesh Tensorflow はこの両方に対応していますが、データ並行だけなら必ずしも必要ではありません。Mesh Tensorflow の主たる目的は「単一の TPU や GPU に全パラメータが格納できないサイズのモデルを作りたい(=モデル並行)」ということです。 OpenAI GPT-2 などの例に見られるように Transformer は規模の拡大が精度に効く傾向があります。そこで、

  • どんどんモデルを巨大にしていったら単一プロセッサにパラメータが格納できなくなった。
  • でも、パラメータが複数プロセッサに配置されるコードを地道に with tf.device('/GPU:X'): とかで指定すると気が狂いそう。
  • モデル並行をエレガントに記述できるライブラリが欲しい！

という流れで Mesh Tensorflow が開発されました。次は Mesh Tensorflow の基本的な考え方を見ていきましょう。

Mesh Tensorflow の基本概念

Mesh Tensorflow の Mesh とは複数のプロセッサ(TPU, GPU, CPU)を論理的なN次元配列に見立てたものです。この Mesh の各要素(=プロセッサ)に対して、以下のように計算処理を分散配置することで、データ並行／モデル並行処理を実現します。

  • Tensor と Mesh の各次元に名前を付ける。
  • Tensor の次元と Mesh の次元をレイアウトルールで対応付ける。
  • Mesh を構成するN次元配列の各要素にプロセッサを対応づける。

Tensor の各次元はレイアウトルールに従い分割または複製され、Mesh を構成するプロセッサで計算されます。

このように Tensor とプロセッサの間に Mesh を挟むことで、同一プログラムを異なる構成のデバイスで実行できる SPMD (Single Program Multiple Device) を実現しています。そういえば、前回の実験でも TPU 8コアで事前学習して、GPU 1コアでファインチューニングしましたね。

具体例として、先程の全結合ネットワークにデータ並行とモデル並行を併用したイメージを考えてみます。図中のプロセッサを “gpu:0"～"gpu:3"と考えてみてください。

example_of_data_and_model_parallel

  • Mesh は 次元が {"r”, “c”} となる 2 x 2 の構成です。Mesh の各要素に “gpu:0"～"gpu:3” を紐付けます。
  • Tensor の各次元には以下のように名前を付けます。
    • x : {“b”, “dio”}
    • h : {“b”, “dh”}
    • y : {“b”, “dio”}
    • w : {“dio”, “dh”}
    • v : {“dh”, “dio”}
  • Mesh と Tensor にレイアウトルール {“b”:“r”, “dh”:“c”} を適用します。

x, h, y の “b” 次元(バッチサイズ)を Mesh の “r” 次元に紐付けてデータ並行で2分割、 h, w, v の “dh” 次元(隠れ層ユニット数)を Mesh の “c” 次元に紐付けモデル並行で2分割です。 h はデータ並行、モデル並行が両方適用されて4分割され、それぞれが図中の青数字に対応する “gpu:0"～"gpu:3” に配置されます（"0,1"の表記は複製されて同じデータが “gpu:0” と “gpu:1” に置かれます）。

この Tensor と Mesh 間のレイアウトは自由自在という訳ではなく、いつくかルールやポイントがあります。

  • 同じ Tensor の複数の次元を Mesh の同じ次元で分割することはできない。 ex. {“b”:“r”; “dio”:“r”}
  • Tensor のある次元を Mesh の複数の次元で分割することはできない。ex. {“dio”:“r”; “dio”:“c”}
  • Tensor の次元サイズが Mesh の次元サイズで割り切れない分割はできない。ex. 上図で b=3の場合(b % r = 3 % 2 != 0)
  • 行列積のような計算量の大きい処理は Mesh の全ての次元に分割する。
  • 計算に対する通信の比率を抑える為、次元は大きい塊に分割する。

Mesh Tensorflow の解説としては、kuroko さんの SlideShare 10 が分かりやすいので興味のある方は見てみるとよいでしょう。

t5 での例

それでは t5 でどのように Mesh Tensorflow が利用されているか見てみましょう。

Mesh の定義

まずは Mesh からです。TPU を用いた事前学習では以下のようにパラメータを定義していました。

  "utils.tpu_mesh_shape.model_parallelism = 1",
  "utils.tpu_mesh_shape.tpu_topology = '2x2'",

このパラメータは、tpu_mesh_shape() 関数で以下のように処理されます11。ただ model_parallelism = 1 では面白くないので、model_parallelism = 2 でいってみましょう。

from mesh_tensorflow.transformer.utils import tpu_mesh_shape
tpu_mesh_shape(tpu_topology="2x2", model_parallelism=2)

# Shape[batch=4, model=2]

このように、 "batch" 次元が 4 、 "model" 次元が 2 の “4x2” の Mesh に変換されました。 tpu_topology="2x2" だから 8 コア12、モデル並行が 2 分割なので、残りがデータ並行となって 4 分割となります。また ensemble_parallelism パラメータを指定することで、 Mesh を {"ensemble", "batch", "model"} の3次元にすることもできます。

レイアウトルール

次はレイアウトルールです。Mesh Tensorflow の default.gin に既定のレイアウトルール13が定義されています。

utils.run.layout_rules = "ensemble:ensemble,batch:batch,d_ff:model,heads:model,vocab:model,experts:batch"

これを確認するとモデル並行が適用されるのは vocab, heads, d_ff の3つの次元のようです。 これらの名称の次元を持つ Tensor はモデル並行の対象になります。vocab は語彙数(32000)、 heads は Multi-Head Self-Attention のヘッド数(12)、 d_ff は Multi-Head Self-Attention を抜けた後の Feed Forward の隠れ層ユニット数(3072)に相当します。

計算グラフが複雑なので全て紹介できませんが、 d_ff で分割された様子を確認して見ましょう。 以下はデータ並行 1、モデル並行 2 の構成の計算グラフにおける Feed Foward の全結合とアクティベーションの部分を Tensorboard で可視化したものです。

model_parallel_graph

kernel_slice_0, kernel_slice_1 で d_ff の 3072 が 1536 に 2 分割されていることがわかります。また、データ並行、モデル並行の数を変えると kernel_slice_x の数が変化します。このことから同じ構造、パラメータのモデルでも Mesh やレイアウトルールの構成によって出来上がる計算グラフは異なってくるということがわかります。

Mesh とプロセッサの対応付け

Mesh のN次元配列とプロセッサの対応づけは TPU の場合は（きっちり確認したわけではありませんが） TPU アドレスなどから自動で計算されるようです。 GPU の場合は mesh_tensorflow.transformer.utils.tpu_estimator_model_fn() の mesh_devices に ['gpu:0', 'gpu:1', 'gpu:2', 'gpu:3'] のように流し込めれば良いです。 t5 の場合は、学習や推論を行うときは mesh_tensorflow.transformer.utils.run() の mesh_devices に gin で設定します。エクスポート時はもう一捻りいるので後述します。

さて、ここからは前回の後に追加で行った検証についてご紹介していきます。

4. OSCAR データセットでの事前学習

まずは、Wikipedia 日本語版よりも大規模なデータセットで事前学習を行った結果です。データセットとしては OSCAR 14 を使いました。

oscar_dataset

OSCAR セットは T5 の事前学習で用いられた C4 同様に Common Crawl に対して記述言語の識別やフィルタリングを行った多言語のデータセットになっています。言語毎に、重複排除をしたもの／していないものの 2 種類のファイルが公開されています。今回は重複排除してある ja_dedup.txt.gz を使いました。C4 には及びませんが圧縮状態で 37.6GB、展開すると 106GB のサイズがあります。

まずは、 OSCAR データセットを TFDS に追加してみましょう。

OSCAR データセットを TFDS へ追加

OSCAR データセットを T5 の事前学習で使う為に、まず TFDS 経由で使えるようにしておきましょう。 もちろん圧縮ファイルを展開して TextLineTask で読み込んでも良いのですが、 TPU で回すことを考えると GCS 上に TFRecord 形式で配置したいところです。そうであれば、 TFDS に組み込んでしまうのがてっとり早そうです。

※ Colab で動かした風に書いていますが、初回実行時はファイルが VM のストレージ上に展開される為、ディスクが足りません。

  • 初回は GCE で余裕をもったストレージサイズでインスタンスを立てて実行して下さい15。展開ファイル(106GB)が収まれば大丈夫なハズですが、128GB だとちょっと不安な感じでしょうか。
  • GCE である必要はないのですが、通信料金を考えると GCE がおすすめです。ローカルストレージと GCS の間で圧縮ファイル、展開ファイル、TFRecord と段階を踏んで複数回のやり取りになる為、結構な転送量になります。
  • 一度、実行すると処理済みの TFRecord ファイルが GCS 上に保存されるので 2 回目からは Colab で実行できます。
  • コード的には「手元でとりあえず動けばいいや」です(「Google の人がそのうちちゃんとしたやつ作るでしょ」的な。。。)。
  • Colab での環境使い捨て前提なので /usr/local/lib/python3.6/dist-packages/tensorflow_datasets/ の下を無邪気に書き換えています。

まずは、Tensorflow のバージョン指定、Googleアカウントの認証、t5 のインストールです。

%tensorflow_version 1.x 
from google.colab import auth
auth.authenticate_user()
!pip install t5[gcp]

OSCAR データセットを TFDS に組み込む為のコードは以下のようになります。

%%bash
cat <<EOF > oscar.py
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import math
import tensorflow as tf
import tensorflow_datasets.public_api as tfds
from tensorflow_datasets.core.download import checksums

_CITATION = "@inproceedings{ortizsuarez:hal-02148693, URL = {https://hal.inria.fr/hal-02148693}, }"
_DESCRIPTION = ("Japanese subset of OSCAR or Open Super-large Crawled ALMAnaCH coRpus.")

_OSCAR_JA_CHECKSUM = (40358865810, "146b1933606db6f97e36c3fbd2fad47cae5443e2e37c4ffe0912ab90f8c45c57")
_OSCAR_JA_URL = "https://traces1.inria.fr/oscar/files/Compressed/ja_dedup.txt.gz"

class OscarJa(tfds.core.GeneratorBasedBuilder):
  """A TFDS crude implementation of japanese subset of OSCAR dataset."""
  VERSION = tfds.core.Version('1.0.0', experiments={tfds.core.Experiment.S3: False, })

  def _info(self):
    return tfds.core.DatasetInfo(
        builder=self,
        description=_DESCRIPTION,
        features=tfds.features.FeaturesDict({
            "text": tfds.features.Text(), }),
        supervised_keys=None,
        homepage="https://traces1.inria.fr/oscar/",
        citation=_CITATION,
    )

  def _split_generators(self, dl_manager):
    # Add checksum for OSCAR datasetx to pre-defined checksums.
    checksums.get_all_sizes_checksums()[_OSCAR_JA_URL] = _OSCAR_JA_CHECKSUM
    downloaded_files = dl_manager.download_and_extract({ 'ja': _OSCAR_JA_URL })
    return [
        tfds.core.SplitGenerator(
            name = tfds.Split.TRAIN,
            num_shards= int(math.ceil(_OSCAR_JA_CHECKSUM[0] / (128 * 2**20))),  
            gen_kwargs = {"filepaths": downloaded_files},
        ),
    ]

  def _generate_examples(self, filepaths):
    with tf.io.gfile.GFile(filepaths["ja"]) as text_file:
      for line in text_file:
        yield {'text': line.strip()}
EOF

以下のようにしてテンプレートを生成と TFDS への登録を行います。 

!python /usr/local/lib/python3.6/dist-packages/tensorflow_datasets/scripts/create_new_dataset.py --dataset oscar_ja  --type text

生成したテンプレートを先ほどのコードで差し替えます。

!mv ./oscar_ja.py /usr/local/lib/python3.6/dist-packages/tensorflow_datasets/text/

追加したファイルが認識されていることを確認しましょう。

import tensorflow_datasets as tfds
import tensorflow_datasets.core.registered as r
"oscar_ja" in r.list_builders()

# True

TFDS を使って OSCAR データセットをロードします16。

ds, ds_info = tfds.load(name='oscar_ja:1.0.0', 
          data_dir = 'gs://somewhere/oscar/ja',
          shuffle_files=False,
          download=True,
          try_gcs=False,
          with_info=True)

 print(ds_info)

 # tfds.core.DatasetInfo(
 #   name='oscar_ja',
 #   version=1.0.0,
 #   description='Japanese subset of OSCAR or Open Super-large Crawled ALMAnaCH coRpus.',
 #   homepage='https://traces1.inria.fr/oscar/',
 #   features=FeaturesDict({
 #       'text': Text(shape=(), dtype=tf.string),
 #   }),
 #   total_num_examples=217755988,
 #   splits={
 #       'train': 217755988,
 #   },
 #   supervised_keys=None,
 #   citation="""@inproceedings{ortizsuarez:hal-02148693, URL = {https://hal.inria.fr/hal-02148693}, }""",
 #   redistribution_info=,
 #)

せっかくですから、データの中身を確認してみましょう。

train_ds = ds["train"].batch(10)
batches =  train_ds.take(1)
for batch in tfds.as_numpy(batches):
 raw_texts = batch["text"]
 for i in range(batch["text"].shape[0]):
   print("%3d, %s" % (i, raw_texts[i].decode("utf-8")))

# 0, 2018/2/13 ★★★★★, ★★★★★★★★, ★女, ★男, ★★★★, ★★★★★, ★★★★★★★, ★★★★★, ★★★★, ★★★★★...
# 1, 結果としては全敗と、ｲｴﾛｰﾁｰﾑは悔しい思いをしました。今回非常に痛感したのが、練習量や内容、練習に取り組む姿勢、それらの違い...
# 2, 個人的には、オリンピックに興味はなく、 スポーツにも特に興味はない。 というか西が丘のプールに通う子の多かった 近隣住民にと...
# 3, 一方、ことさら燃費の改善だけを追わなかった開発陣の意図は、走れば走るほど身にしみる。スラッと明快な吹け上がり、前：後＝50：...
# 4, Mさん：私の場合、英語はまだ理解できたけど、スウェーデン語でディスカッションとか始まることがあって、もうそういうときはまっ...
# 5, それで、途中でこっちでルールを捻じ曲げる(例えば比較の対象としてこっちから具体的な数値を出しちゃう)とか、こっち側に「おり...
# 6, 合資会社　山口長生堂（やまぐち薬局）薬剤師 / 有限会社　大里商事調剤薬局　薬剤師 / 緑の里　調剤薬局有限会社調剤薬局　薬剤師 ...
# 7, 11位フェルスタッペン以降、フェリペ・ナッサー（ザウバー）、ペレス、エリクソン、ヒュルケンベルグ、バトン、ウィル・スティーブ...
# 8, 平和通１丁目駅のカウンター席のあるお店が10件見つかりました。特にレストランのお店が多いです。この条件に合致するお店は大街道...
# 9, サイフェルト　おっしゃる通りです。そして、非常に大切な点は、子どもの音楽教育はすでに胎内から始まっているということです。

上記を見てわかるように、 OSCAR は様々なサイトから取得したデータを行単位でシャッフルした形式となっています。ですので、Wikipedia 日本語版のデータセットと比較すると 1 サンプル辺りの文長はかなり短くなります。また取得元を示すメタデータは提供されていません。

データの内容に着目すると行番号 0, 6 のように、このまま事前学習に含めてしまうのが好ましくないと思われるサンプルも散見されました(「★」の部分は筆者の自主規制です)。また筆者が目視確認した限りでは Wikipedia 日本語版の文章も含まれているようです。ただし Wikipedia 全体の何%がカバーされているかは確認できていません。データの質という観点では対象ドメインが広範囲をカバー出来ているが、広告や内容的に不適切なものもそれなりに含まれているので注意が必要といったところでしょうか。

このようなサンプルに対し、今回の実験では簡易的な対応として「。」で終了しないサンプルを排除するという対応をしています。目視でざっと確認しただけですが、それなりに不適切なデータを排除できている様子でした17。

量的な側面として Wikipedia 日本語版と比較したところ以下のようになりました。

size_comparison

有効サンプル数は「。」で終了するサンプルの数、トークン数は Wikipedia、OSCAR それぞれに個別の Sentencepiece モデルを学習してトークン化した結果を集計したものです。ただし OSCAR は有効とした 141,678,084 サンプルのみを Sentencepiece モデルの学習及びトークン数の集計に用いています。

「。」で終了するデータのみを対象とすることで、サンプル数は元の 65% 程度になってしまいましたが、それでもトークン数としては Wikipedia の 10倍の規模になります。逆にサンプル毎の平均トークン長は 1/10 程度になっているので長い文脈を十分に学習できるか？という懸念があります。

なにはともあれ、より大規模なデータセットが使えるようになったので実験して結果を比較してみました。手順は前回と異なる部分について記述しているので、詳細は 前回記事 を参照ください。

事前学習の実行

まずは Wikipedia の時と同様に Task, Mixture を作って登録しておく必要があります。

SPM_PATH = "gs://somewhere/t5/sentencepiece/oscar_ja_146b193.model"

import tensorflow as tf
def exclude_inappropriate(dataset):
  def my_fn(x):
    return {'text' : tf.strings.regex_replace(x['text'], '^.*[^。]$', '')}
  return dataset.map(my_fn, num_parallel_calls=tf.data.experimental.AUTOTUNE)

import functools
from t5.data import preprocessors
from t5.data.utils import TaskRegistry
from t5.data.utils import MixtureRegistry
from t5.data.utils import TfdsTask

task_name_oscar_ja = "oscar_ja_146b193_unsupervised"

TaskRegistry.add(
    task_name_oscar_ja,
    TfdsTask,
    tfds_name="oscar_ja:1.0.0",
    text_preprocessor=[exclude_inappropriate, functools.partial(
        preprocessors.rekey, key_map={"inputs": None, "targets": "text"})],
    token_preprocessor=preprocessors.unsupervised,
    sentencepiece_model_path=SPM_PATH,
    metric_fns=[])

MixtureRegistry.add(task_name_oscar_ja, [(task_name_oscar_ja, 1.0)])

処理的に前回と異なるのは Sentencepiece モデルのパスと、「。」で終了しないサンプルを除外する exclude_inappropriate() を定義して text_preprocessor に組み込んだくらいです。他は Mixture の名称が異なる以外は前回と同じ設定で事前学習を行いました。

ファインチューニングの実行

ファインチューニングの実行は Wikipedia と変わりありません。

結果比較

OSCAR データセットで事前学習したモデルを前回の実験にならって livedoor ニュースコーパス、やさしい日本語(拡張)コーパスを用いてファインチューニングしたので見ていきましょう。なお、事前学習のデータセットと Sentencepiece のモデルが異なる以外の条件は同じにしてあります。

まずは、 livedoor ニュースコーパスでの文章分類です。例によって交差検証していないので、割り引いて考えないといけませんが、数値的には 1 pt ほど向上しています。

ldcc_wikipedia_oscar

また、Wikipedia での 0が連続で出力される問題(“000000…”)も今回のモデルでは解消されました。

!cat infer_ouputs.txt-535288

# 7
# 2
# 5
# 6
# 0

こちらが、やさしい日本語変換の BLEU と SARI です。こちらは誤差程度の違いで向上も劣化もしていない感じです。

snow_wikipedia_oscar

気持ち的には、もっと明確にスコアが上がって「やはりデータ量が大きいと違いますね。」と言いたいところだったので少し残念な結果になってしまいました。この結果について感覚的にですが少し考えてみます。

まず、T5 の論文では事前学習データセットの対象ドメインと後続タスクの対象が近いと精度が向上するようでした。 livedoor News コーパスの分類では精度が上がりましたが、感覚的には文章が Wikipedia よりも OSCAR に近い気がしなくもないです。

次にスコアが伸びない理由としては事前学習サンプルの平均長が影響した可能性があります。 事前学習でのシーケンス長は 512 です。2章のノイズ挿入等の前処理で示したように、今回の事前学習ではサンプルを連結し、できるだけ学習シーケンスの 512 を埋めるような挙動をします。OSCAR の平均トークン数は 83 です。面倒なので、ざっくり計算ですが 512 / 83 = 6.17 なので、事前学習に投入されるサンプル毎に [29 30 31] の 29 30 のような元データのサンプル境界を跨ぐケースが 6 回発生することになります。 C4 や Wikipedia のように元データのサンプルが長ければ、2章で述べたように「文の切れ目で文脈が変わっている形になり、普通によくあることなのでさほど問題ありません。」なのですが、サンプル毎に平均 6 回だと文脈をとらえる能力に影響がでないか心配です18。やさしい日本語(拡張)コーパスはサンプルの文長が短いので、長い文脈をとらえる力は必要ないのですが、脈絡のない文脈変化がこれだけ混入すると短い文脈をとらえる力にも影響があるかもしれません。

人間でも会話していて突然に話の文脈が変わる人がいます。仮に田中さん(仮名)としておきます。私はもう慣れたので戸惑うことはないのですが、私は田中さん(仮名)と会話するとき長い文脈での整合性を重視しないようになりました。これと同じようなことが Transformer でも起こりそうです。

最後に学習トークン数です。 OSCAR の事前学習でも学習トークン数は 235 で固定です。Wikipedia に比べて全体のトークン数は10倍なので、もっと回してあげるとスコアが向上する可能性はあるかもしれません。

さて、ここからは学習したモデルをエクスポートして Tensorflow Serving にデプロイしてみましょう。

5. T5 を Tensorflow Serving へデプロイする

それでは、T5 の学習済みモデルを Tensorflow Serving へデプロイしてみましょう。事前学習は TPU 8コア、ファインチューニングは GPU 1コアで行ってきました。せっかくの Mesh Tensorflow ですので、 GPU 8コアの環境にデプロイしてモデル並行、データ並行を併用する形にしてみます。

※ 今回の検証で使用している t5 のバージョンは以下のとおりです。

!pip list | grep -e t5
# t5                       0.3.0

SavedModel のエクスポート

まずは checkpoint を SavedModel の形式でエクスポートしていきましょう。

エクスポートする際は t5_mesh_transformer コマンドを mode="export" で起動するのですが、3つほどポイントがあります。

Mesh シェイプの指定

Mesh の構成としてはデータ並行 2 、モデル並行 4 とし GPU 8コアで並列実行とします。GPU で実行する場合は tpu_mesh_shape() を使わず--gin_param="MtfModel.mesh_shape = 'batch:2,model:4'" のように Mesh シェイプを直接記述できます。

Sentencepiece のモデルとビームサーチのパラメータ

Sentencepiece のモデルやビームサーチのパラメータを指定するコマンドライン引数は用意されていません。仕方がないので、以下のようなコードを用意して --module_import でコードを書き換えることにしました。

%%bash
cat <<EOF > param_set_sentencepiece_and_beam_search.py
from t5.models.mtf_model import MtfModel
export_org = MtfModel.export
def export_with_param_overrides(self, export_dir, checkpoint_step):
  export_org(self, export_dir, checkpoint_step, 
         beam_size=4, temperature=0.0,
         sentencepiece_model_path="gs://somewhere/t5/sentencepiece/wikipedia_20190301_ja_v003.model")
MtfModel.export = export_with_param_overrides
EOF

Mesh とプロセッサの対応付け

次にMesh とプロセッサの対応付けで示したとおり、Mesh の各要素に GPU を割り当てなければなりません。

t5 0.3.0 では t5.models.mtf_model.MtfModel クラスのコンストラクタに mesh_devices パラメータが追加されている為、t5_mesh_transformer コマンドの引数に gin で、

--gin_param="MtfModel.mesh_devices = ['gpu:0', 'gpu:1', 'gpu:2', 'gpu:3', 'gpu:4', 'gpu:5', 'gpu:6', 'gpu:7']" \

のように指定すればＯＫなのですが、どうも t5.models.mtf_model.MtfModel.estimator() にはエクスポート時に Mesh が Shape([]) になる問題があるようで、--module_import で以下のようにコードを差し替えました。 mesh_shape の行を修正しています(コメント行が修正前の記述です)。

%%bash
cat <<EOF > fix_MtfModel_estimator.py
from mesh_tensorflow.transformer import utils
from t5.models.mtf_model import MtfModel
import mesh_tensorflow as mtf

# mesh_shape=mtf.Shape([]) if disable_tpu else self._mesh_shape,

def estimator_fixed(self, vocabulary, init_checkpoint=None, disable_tpu=False):
  return utils.get_estimator(
    model_type=self._model_type,
    vocabulary=vocabulary,
    layout_rules=self._layout_rules,
    mesh_shape=self._mesh_shape,
    mesh_devices=self._mesh_devices,
    model_dir=self._model_dir,
    batch_size=self.batch_size,
    sequence_length=self._sequence_length,
     autostack=self._autostack,
    learning_rate_schedule=self._learning_rate_schedule,
    keep_checkpoint_max=self._keep_checkpoint_max,
    save_checkpoints_steps=self._save_checkpoints_steps,
    optimizer=self._optimizer,
    predict_fn=self._predict_fn,
    variable_filter=self._variable_filter,
    ensemble_inputs=self._ensemble_inputs,
    use_tpu=None if disable_tpu else self._tpu,
    tpu_job_name=self._tpu_job_name,          
    iterations_per_loop=self._iterations_per_loop,
    cluster=self._cluster,
    init_checkpoint=init_checkpoint)

MtfModel.estimator = estimator_fixed
EOF

以上の点を踏まえてエクスポートのコマンドは以下のようになりました。

!export PYTHONPATH=${PYTHONPATH}:. && \
  FINE_TUNED_MODEL_DIR='gs://somewhere/t5/wikipedia_20190301.ja_v003/snow_t15_23' && \
  EXPORT_MODEL_DIR='gs://somewhere/t5/wikipedia_20190301.ja_v003/snow_t15_23/exportx8_gpux8' && \
  INPUT_SEQ_LEN=64 && \
  TARGET_SEQ_LEN=64 && \
  MESH_SHAPE='batch:2,model:4' && \
  MESH_DEVICES="['gpu:0', 'gpu:1', 'gpu:2', 'gpu:3', 'gpu:4', 'gpu:5', 'gpu:6', 'gpu:7']"&& \
  BATCH_SIZE_METHOD='tokens_per_batch' && \
  BATCH_SIZE_VALUE=128 && \
  \
  echo "FINE_TUNED_MODEL_DIR=$FINE_TUNED_MODEL_DIR" && \
  echo "EXPORT_MODEL_DIR=$EXPORT_MODEL_DIR" && \
  echo "INPUT_SEQ_LEN=$INPUT_SEQ_LEN" && \
  echo "TARGET_SEQ_LEN=$TARGET_SEQ_LEN" && \
  echo "MESH_SHAPE=$MESH_SHAPE" && \
  echo "MESH_DEVICES=$MESH_DEVICES" && \
  echo "BATCH_SIZE_METHOD=$BATCH_SIZE_METHOD" && \
  echo "BATCH_SIZE_VALUE=$BATCH_SIZE_VALUE" && \
  \
  t5_mesh_transformer \
  --model_dir="$FINE_TUNED_MODEL_DIR" \
  --module_import="param_set_sentencepiece_and_beam_search" \
  --module_import="fix_MtfModel_estimator" \
  --use_model_api \
  --mode="export" \
  --gin_param="MtfModel.mesh_shape = '$MESH_SHAPE'" \
  --gin_param="MtfModel.batch_size=('$BATCH_SIZE_METHOD', $BATCH_SIZE_VALUE)" \
  --gin_param="MtfModel.sequence_length = {'inputs': $INPUT_SEQ_LEN, 'targets': $TARGET_SEQ_LEN}" \
  --gin_param="MtfModel.mesh_devices=$MESH_DEVICES" \
  --export_dir="$EXPORT_MODEL_DIR"

検証や推論時と違うのは、差し替えコードを --module_import で読み込むところ、 --mode に export を渡すところ、 --use_model_api を指定するところになります。

実行すると、やたらと以下のような警告が出力されますが、「指定されたデバイスにグラフを配置しようとしたけど、そんなデバイスないよ」という意味で、モデル自体はちゃんと指定したとおりにエクスポートされているので大丈夫です。

2020-03-13 06:56:56.188584: W tensorflow/core/common_runtime/colocation_graph.cc:983] Failed to place the graph without changing the devices of some resources. Some of the operations (that had to be colocated with resource generating operations) are not supported on the resources' devices. Current candidate devices are [
  /job:localhost/replica:0/task:0/device:GPU:0
  /job:localhost/replica:0/task:0/device:CPU:0].

SavedModel のインタフェース

それでは、エクスポートされたモデルのインタフェースを見てみましょう。

!gsutil ls gs://somewhere/t5/wikipedia_20190301.ja_v003/snow_t15_23/exportx8_gpux8/

# gs://somewhere/t5/wikipedia_20190301.ja_v003/snow_t15_23/exportx8_gpux8/
# gs://somewhere/t5/wikipedia_20190301.ja_v003/snow_t15_23/exportx8_gpux8/1584082501/

!saved_model_cli show --dir gs://somewhere/t5/wikipedia_20190301.ja_v003/snow_t15_23/exportx8_gpux8/1584082501/ \
  --tag_set serve --signature_def serving_default

# The given SavedModel SignatureDef contains the following input(s):
#   inputs['input'] tensor_info:
#       dtype: DT_STRING
#       shape: (-1)
#       name: inputs:0
# The given SavedModel SignatureDef contains the following output(s):
#   outputs['inputs'] tensor_info:
#       dtype: DT_STRING
#       shape: (2)
#       name: SentenceTokenizer/SentenceTokenizer/SentencepieceDetokenizeOp:0
#   outputs['outputs'] tensor_info:
#       dtype: DT_STRING
#       shape: (2)
#       name: SentenceTokenizer_1/SentenceTokenizer/SentencepieceDetokenizeOp:0
# Method name is: tensorflow/serving/predict

入力側から見ていきましょう。

#   inputs['input'] tensor_info:
#       dtype: DT_STRING
#       shape: (-1)
#       name: inputs:0

これは Text-to-Text への入力文を束ねて渡せば良さそうですね。出力側はどうなっているでしょう？

#   outputs['inputs'] tensor_info:
#       dtype: DT_STRING
#       shape: (2)
#       name: SentenceTokenizer/SentenceTokenizer/SentencepieceDetokenizeOp:0
#   outputs['outputs'] tensor_info:
#       dtype: DT_STRING
#       shape: (2)
#       name: SentenceTokenizer_1/SentenceTokenizer/SentencepieceDetokenizeOp:0

出力としては、Text-to-Text への入力文と出力文をセットで返してくれるようです。ここで、 shape: (2) の 2 はデータ並行の 2 ですね。

saved_model_cli は SavedModel を実行する機能があるので試してみます。

'吾輩は猫である'.encode('utf-8')
# b'\xe5\x90\xbe\xe8\xbc\xa9\xe3\x81\xaf\xe7\x8c\xab\xe3\x81\xa7\xe3\x81\x82\xe3\x82\x8b'

これを以下のようにして実行します。

!saved_model_cli run \
  --dir gs://somewhere/t5/wikipedia_20190301.ja_v003/snow_t15_23/exportx8_gpux8/1584082501/ \
  --tag_set serve \
  --signature_def serving_default \
  --input_exprs="input=[b'\xe5\x90\xbe\xe8\xbc\xa9\xe3\x81\xaf\xe7\x8c\xab\xe3\x81\xa7\xe3\x81\x82\xe3\x82\x8b']"

# 2020-03-12 01:38:38.250573: I tensorflow/stream_executor/platform/default/dso_loader.cc:44] Successfully opened dynamic library libcuda.so.1
# 2020-03-12 01:38:38.284616: I tensorflow/stream_executor/cuda/cuda_gpu_executor.cc:983] successful NUMA node ...
# ...
# tensorflow.python.framework.errors_impl.NotFoundError: Op type not registered 'SentencepieceOp' in binary running on 3c34d5ca0220. 
# Make sure the Op and Kernel are registered in the binary running in this process....

ログが長いので省略したり改行したりしていますが、Op type not registered 'SentencepieceOp' in binary が原因でエラーになりました。

SavedModel には文章で入力しているので、 SavedModel の中で Sentencepiece のトークン化、脱トークン化が行われます。この部分を担うのが Tensorflow Text 19 なのですが、このライブラリが saved_model_cli にリンクされていないことが原因です。

そして、この事情は Tensorflow Serving でも同じなのでした。。。

Tensorflow Serving のビルド

前述のとおり、Tensorflow Serving には SentencepieceOp がリンクされていません。現状では以下の issue で議論されているように自前で Tensorflow Serving をリビルドするしかないようです。

この issue のコメント20で SentencepieceOp をリンクした Tenserflow Serving のコンテナイメージをビルドする話がでてきますが、これは GPUのサポートがありません。コメントの内容を参考に GPU サポートありでコンテナをビルドしたところ動作したので紹介しておきます。必要なファイルは以下の4つです。Dockerfile.gpu 以外は先ほどのコメント中に掲載されています。

  • Dockerfile.gpu
  • tftext.patch1
  • tftext.patch2
  • workspace.bzl

Dockerfile.gpu の内容は少し長いですが以下のとおりです。

FROM nvidia/cuda:10.0-base-ubuntu16.04 as base_build

ARG TF_SERVING_VERSION_GIT_BRANCH=r2.1
ARG TF_SERVING_VERSION_GIT_COMMIT=head

ENV CUDNN_VERSION=7.4.1.5
ENV TF_TENSORRT_VERSION=5.0.2

RUN apt-get update && apt-get install -y --no-install-recommends \
        automake \
        build-essential \
        ca-certificates \
        cuda-command-line-tools-10-0 \
        cuda-cublas-dev-10-0 \
        cuda-cudart-dev-10-0 \
        cuda-cufft-dev-10-0 \
        cuda-curand-dev-10-0 \
        cuda-cusolver-dev-10-0 \
        cuda-cusparse-dev-10-0 \
        curl \
        git \
        libfreetype6-dev \
        libpng12-dev \
        libtool \
        libcudnn7=${CUDNN_VERSION}-1+cuda10.0 \
        libcudnn7-dev=${CUDNN_VERSION}-1+cuda10.0 \
        libcurl3-dev \
        libzmq3-dev \
        mlocate \
        openjdk-8-jdk\
        openjdk-8-jre-headless \
        pkg-config \
        python-dev \
        software-properties-common \
        swig \
        unzip \
        wget \
        zip \
        zlib1g-dev \
        && \
    apt-get clean && \
    rm -rf /var/lib/apt/lists/* && \
    find /usr/local/cuda-10.0/lib64/ -type f -name 'lib*_static.a' -not -name 'libcudart_static.a' -delete && \
    rm /usr/lib/x86_64-linux-gnu/libcudnn_static_v7.a

# The 'apt-get install' of the nvinfer-runtime-trt-repo-* library adds a new
# list which contains libnvinfer library, so it needs another 'apt-get update'
# to retrieve that list before it can actually install the library.
RUN apt-get update && \
    apt-get install --no-install-recommends \
        nvinfer-runtime-trt-repo-ubuntu1604-${TF_TENSORRT_VERSION}-ga-cuda10.0 && \
    apt-get update && \
    apt-get install --no-install-recommends \
        libnvinfer5=${TF_TENSORRT_VERSION}-1+cuda10.0 \
        libnvinfer-dev=${TF_TENSORRT_VERSION}-1+cuda10.0 && \
    apt-get clean && \
    rm -rf /var/lib/apt/lists/* && \
    rm /usr/lib/x86_64-linux-gnu/libnvinfer_static.a && \
    rm /usr/lib/x86_64-linux-gnu/libnvinfer_plugin_static.a && \
    rm /usr/lib/x86_64-linux-gnu/libnvcaffe_parser* && \
    rm /usr/lib/x86_64-linux-gnu/libnvparsers*

RUN curl -fSsL -O https://bootstrap.pypa.io/get-pip.py && \
    python get-pip.py && \
    rm get-pip.py

RUN pip --no-cache-dir install \
    future>=0.17.1 \
    grpcio \
    h5py \
    keras_applications>=1.0.8 \
    keras_preprocessing>=1.1.0 \
    mock \
    numpy \
    requests

# Set up Bazel
ENV BAZEL_VERSION 0.24.1
WORKDIR /
RUN mkdir /bazel && \
    cd /bazel && \
    curl -H "User-Agent: Mozilla/5.0 (X11; Linux x86_64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/57.0.2987.133 Safari/537.36" -fSsL -O https://github.com/bazelbuild/bazel/releases/download/$BAZEL_VERSION/bazel-$BAZEL_VERSION-installer-linux-x86_64.sh && \
    curl -H "User-Agent: Mozilla/5.0 (X11; Linux x86_64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/57.0.2987.133 Safari/537.36" -fSsL -o /bazel/LICENSE.txt https://raw.githubusercontent.com/bazelbuild/bazel/master/LICENSE && \
    chmod +x bazel-*.sh && \
    ./bazel-$BAZEL_VERSION-installer-linux-x86_64.sh && \
    cd / && \
    rm -f /bazel/bazel-$BAZEL_VERSION-installer-linux-x86_64.sh

# Build TensorFlow with the CUDA configuration
ENV CI_BUILD_PYTHON python
ENV LD_LIBRARY_PATH /usr/local/cuda/extras/CUPTI/lib64:$LD_LIBRARY_PATH
ENV TF_NEED_CUDA 1
ENV TF_NEED_TENSORRT 1
ENV TENSORRT_INSTALL_PATH=/usr/lib/x86_64-linux-gnu
ENV TF_CUDA_COMPUTE_CAPABILITIES=3.0,3.5,5.2,6.0,6.1,7.0
ENV TF_CUDA_VERSION=10.0
ENV TF_CUDNN_VERSION=7

# Fix paths so that CUDNN can be found: https://github.com/tensorflow/tensorflow/issues/8264
WORKDIR /
RUN mkdir /usr/lib/x86_64-linux-gnu/include/ && \
  ln -s /usr/include/cudnn.h /usr/local/cuda/include/cudnn.h && \
  ln -s /usr/lib/x86_64-linux-gnu/libcudnn.so /usr/local/cuda/lib64/libcudnn.so && \
  ln -s /usr/lib/x86_64-linux-gnu/libcudnn.so.${TF_CUDNN_VERSION} /usr/local/cuda/lib64/libcudnn.so.${TF_CUDNN_VERSION}

# For backward compatibility we need this line. After 1.13 we can safely remove
# it.
ENV TF_NCCL_VERSION=

# Set TMP for nvidia build environment
ENV TMP="/tmp"

# Download TF Serving sources (optionally at specific commit).
WORKDIR /tensorflow-serving
RUN git clone --branch=${TF_SERVING_VERSION_GIT_BRANCH} https://github.com/tensorflow/serving . && \
    git remote add upstream https://github.com/tensorflow/serving.git && \
    if [ "${TF_SERVING_VERSION_GIT_COMMIT}" != "head" ]; then git checkout ${TF_SERVING_VERSION_GIT_COMMIT} ; fi

COPY workspace.bzl ./tensorflow_serving/workspace.bzl

COPY tftext.patch1 ./third_party/tf_text/tftext.patch1
COPY tftext.patch2 ./third_party/tf_text/tftext.patch2

RUN sed -i.bak '/@org_tensorflow_text\/\/tensorflow_text:wordpiece_tokenizer_cc/a\    "@org_tensorflow_text\/\/tensorflow_text:sentencepiece_tokenizer_cc",' \
    tensorflow_serving/model_servers/BUILD

FROM base_build as binary_build
# Build, and install TensorFlow Serving
ARG TF_SERVING_BUILD_OPTIONS="--config=nativeopt"
RUN echo "Building with build options: ${TF_SERVING_BUILD_OPTIONS}"
ARG TF_SERVING_BAZEL_OPTIONS=""
RUN echo "Building with Bazel options: ${TF_SERVING_BAZEL_OPTIONS}"

RUN ln -s /usr/local/cuda/lib64/stubs/libcuda.so /usr/local/cuda/lib64/stubs/libcuda.so.1 && \
    LD_LIBRARY_PATH=/usr/local/cuda/lib64/stubs:${LD_LIBRARY_PATH} \
    bazel build --color=yes --curses=yes --config=cuda --copt="-fPIC"\
    ${TF_SERVING_BAZEL_OPTIONS} \
    --verbose_failures \
    --output_filter=DONT_MATCH_ANYTHING \
    ${TF_SERVING_BUILD_OPTIONS} \
    tensorflow_serving/model_servers:tensorflow_model_server && \
    cp bazel-bin/tensorflow_serving/model_servers/tensorflow_model_server \
    /usr/local/bin/ && \
    rm /usr/local/cuda/lib64/stubs/libcuda.so.1

CMD ["/bin/bash"]

ビルドする環境は GPU は必要ありません。複数CPUコアがあると並行してビルドしてくれるので、コア数の多い環境でビルドすると良いでしょう。前述の 4 ファイルをカレントディレクトリに置いた状態で以下のコマンドでビルドします21。

docker build -f Dockerfile.gpu -t tensorflow-serving-with-latest-tensorflow-text-gpu .

Tensorflow Serving の起動と gRPC での呼び出し

ビルド出来たら、コンテナのイメージと SavedModel を GPU のある環境に移し以下のようにして起動します。筆者は AWS の p2.8xlarge と “Deep Learning Base AMI (Ubuntu 16.04) Version 22.0 (ami-0f009bf0d2e29c5a9)” で確認しました。

※ 以下、"$“ で始まる行は ホストOSの Ubuntu、”#“ はコンテナのプロンプトです。

$ docker run
  --name serving_text \ 
  --gpus all \
  --entrypoint /usr/local/bin/tensorflow_model_server \
  -p 8500:8500 \
  -p 8501:8501 \ 
  -v /home/ubuntu/exportx8_gpux8:/models/t5-snow \
  -t tensorflow-serving-with-latest-tensorflow-text-gpu:latest \  
  --port=8500 \
  --rest_api_port=8501 \
  --model_name=t5-snow \
  --model_base_path=/models/t5-snow

同じ p2.8xlarge インスタンス上でクライアントになるコンテナを以下のように起動します。

$ docker run -it -v /home/ubuntu:/work --name serving_client --link serving_text:serving tensorflow/tensorflow:1.15.2-gpu-py3

クライアント側のコンテナで必要なモジュールをインストールします。

# pip install tensorflow-serving-api==1.15.0
# pip install t5

クライアント側のコードは以下のとおりです。第三回の BERT の時は REST インタフェースを使ったので、今回は gRPC を使ってみましょう。

grpc_predict.py

#!/usr/bin/env python3
# coding: utf-8

import grpc
import numpy as np
import tensorflow as tf
from tensorflow_serving.apis import prediction_service_pb2_grpc, predict_pb2
from tensorflow.core.framework import types_pb2
from t5.data.sentencepiece_vocabulary import SentencePieceVocabulary
import sys
args = sys.argv

text="sample"
host="localhost"

if len(args) == 3:
  text = args[1]
  host = args[2]
else :
  print("Usage: python3 grpc_predict.py text_to_convert serving_host")
  sys.exit(1)

SPM_PATH = "./wikipedia_20190301_ja_v003.model"
vocabulary = SentencePieceVocabulary(SPM_PATH)

channel = grpc.insecure_channel("%s:8500" % host)
predict_stub = prediction_service_pb2_grpc.PredictionServiceStub(channel)

request = predict_pb2.PredictRequest()  
request.model_spec.name = 't5-snow' 
request.model_spec.signature_name = 'serving_default'
encoded_texts = [text.encode('utf-8')]
request.inputs['input'].CopyFrom(tf.make_tensor_proto(encoded_texts))

result = predict_stub.Predict(request, 60.0)

inputs = tf.make_ndarray(result.outputs["inputs"])
outputs = tf.make_ndarray(result.outputs["outputs"])

print("text       : %s" % (text))

print("inputs[0]  : %s" % (inputs[0].decode("utf-8", "backslashreplace")))
print("inputs[1]  : %s" % (inputs[1].decode("utf-8", "backslashreplace")))
print("outputs[0] : %s" % (outputs[0].decode("utf-8", "backslashreplace")))
print("outputs[1] : %s" % (outputs[1].decode("utf-8", "backslashreplace")))

実際に動かしたときの出力です。「吾輩」→「あなた」と（間違ってますけど）やさしい日本語に変換されているのが確認できました。

# python3 grpc_predict.py 吾輩は猫である serving  

WARNING:tensorflow:From grpc_predict.py:33: The name tf.make_tensor_proto is deprecated. Please use tf.compat.v1.make_tensor_proto instead.

text       : 吾輩は猫である
inputs[0]  : 吾輩は猫である
inputs[1]  :
outputs[0] : あなたは猫である。
outputs[1] : 彼女に会うためには、お金が必要だ。

データ並行数が 2 の構成で 1 件だけ入力したので、 input[1] に ""（空文字列）がパディングされています。output[1] は "" を入力として文を出力させたら、こうなりましたということで。ちょっと世知辛いですね。もう少し深いことを言ってほしかったです。。。

連続的に gRPC 呼び出しをしている最中の GPU の使用状況を nvidia-smi で確認してみました。

Fri Mar 13 07:19:15 2020
+-----------------------------------------------------------------------------+
| NVIDIA-SMI 440.33.01    Driver Version: 440.33.01    CUDA Version: 10.2     |
|-------------------------------+----------------------+----------------------+
| GPU  Name        Persistence-M| Bus-Id        Disp.A | Volatile Uncorr. ECC |
| Fan  Temp  Perf  Pwr:Usage/Cap|         Memory-Usage | GPU-Util  Compute M. |
|===============================+======================+======================|
|   0  Tesla K80           On   | 00000000:00:17.0 Off |                    0 |
| N/A   68C    P0    64W / 149W |  10904MiB / 11441MiB |     40%      Default |
+-------------------------------+----------------------+----------------------+
|   1  Tesla K80           On   | 00000000:00:18.0 Off |                    0 |
| N/A   53C    P0    76W / 149W |  10904MiB / 11441MiB |     32%      Default |
+-------------------------------+----------------------+----------------------+
|   2  Tesla K80           On   | 00000000:00:19.0 Off |                    0 |
| N/A   71C    P0    63W / 149W |  10904MiB / 11441MiB |     36%      Default |
+-------------------------------+----------------------+----------------------+
|   3  Tesla K80           On   | 00000000:00:1A.0 Off |                    0 |
| N/A   59C    P0    75W / 149W |  10904MiB / 11441MiB |     32%      Default |
+-------------------------------+----------------------+----------------------+
|   4  Tesla K80           On   | 00000000:00:1B.0 Off |                    0 |
| N/A   73C    P0    63W / 149W |  10904MiB / 11441MiB |     36%      Default |
+-------------------------------+----------------------+----------------------+
|   5  Tesla K80           On   | 00000000:00:1C.0 Off |                    0 |
| N/A   54C    P0    75W / 149W |  10903MiB / 11441MiB |     31%      Default |
+-------------------------------+----------------------+----------------------+
|   6  Tesla K80           On   | 00000000:00:1D.0 Off |                    0 |
| N/A   73C    P0    64W / 149W |  10904MiB / 11441MiB |     32%      Default |
+-------------------------------+----------------------+----------------------+
|   7  Tesla K80           On   | 00000000:00:1E.0 Off |                    0 |
| N/A   58C    P0    75W / 149W |  10904MiB / 11441MiB |     33%      Default |
+-------------------------------+----------------------+----------------------+

+-----------------------------------------------------------------------------+
| Processes:                                                       GPU Memory |
|  GPU       PID   Type   Process name                             Usage      |
|=============================================================================|
|    0      5367      C   /usr/local/bin/tensorflow_model_server     10886MiB |
|    1      5367      C   /usr/local/bin/tensorflow_model_server     10886MiB |
|    2      5367      C   /usr/local/bin/tensorflow_model_server     10886MiB |
|    3      5367      C   /usr/local/bin/tensorflow_model_server     10886MiB |
|    4      5367      C   /usr/local/bin/tensorflow_model_server     10886MiB |
|    5      5367      C   /usr/local/bin/tensorflow_model_server     10886MiB |
|    6      5367      C   /usr/local/bin/tensorflow_model_server     10886MiB |
|    7      5367      C   /usr/local/bin/tensorflow_model_server     10886MiB |
+-----------------------------------------------------------------------------+

めでたく GPU 8コアに処理が分散されることが確認できました。

6. おわりに

今回は前回に引き続きの T5 ということで、事前学習データを変えてみたり、Mesh Tensorflow の話をしてみたりとやや詳細な話題を取り上げました。 次回は、もう少しお手軽に Hagging Face Transformers の日本語事前学習済み BERT を使って Sentence-BERT を試してみたいと思います。

  1. https://arxiv.org/abs/1910.10683 ↩

  2. https://github.com/google-research/text-to-text-transfer-transformer< ↩

  3. https://traces1.inria.fr/oscar/ ↩

  4. https://github.com/tensorflow/serving ↩

  5. https://github.com/tensorflow/datasets  ↩

  6. t5 は処理の順序や内容を gin で組み換え可能として、論文中の様々な目的関数に柔軟に対応できるようになっています。  ↩

  7. https://github.com/google-research/text-to-text-transfer-transformer/blob/3c30c15b0216048ef8de1948f3d46de502602a0b/t5/models/gin/objectives/span.gin#L8-L13 ↩

  8. 中身は sklearn, allennlp 等のライブラリのラッパーなのでさほど難しくはないと思います。 ↩

  9. https://github.com/tensorflow/mesh  ↩

  10. https://www.slideshare.net/ssuserb63bc9/mesh-tensorflow  ↩

  11. GPU の場合は utils.run.mesh_shape="model:3, batch:2" のように直接指定するようにドキュメントに記述されていますが、この関数を使うこともできるようです。  ↩

  12. この "2x2” が 8 コアになる理屈は理解できてません。  ↩

  13. https://github.com/tensorflow/mesh/blob/8c43808806b0f34101045257820e7fb3e37bd312/mesh_tensorflow/transformer/gin/defaults.gin#L64  ↩

  14. https://traces1.inria.fr/oscar/  ↩

  15. GCS へのアクセス方法は https://www.tensorflow.org/datasets/gcs を参考にしてください。  ↩

  16. 事前に GCE 等で事前実行していない場合は、この時点でダウンロード、展開、フォーマット変換、GCSへの保存が行われますが、 Colab のインスタンスストレージの空き容量が足らないので失敗してしまいます。  ↩

  17. そういえば C4 データセットも “.”, “!”, “?”, “”” 等で終了しないサンプルを除去してましたね。 ↩

  18. じつは途中で「あれ？」とは思っていて。学習のシーケンス長を 512 から短くすることも考えたのですが、そうすると挿入するノイズの数とかスパン長とか調整必要そうで、Wikipedia との比較という意味ではどうかな？と思い同じ設定で実験しました（面倒だったともいいます）。  ↩

  19. https://github.com/tensorflow/text  ↩

  20. https://github.com/tensorflow/serving/issues/1490#issuecomment-571623401  ↩

  21. Dockerfile 中で Tensorflow Serving のコミットが head になっています。筆者がビルドしたコンテナでは d83512c6b5b2b8433 になっていました。  ↩