nikkie-ftnextの日記

イベントレポートや読書メモを発信

TensorFlowのドキュメントを確認し、tf.py_functionが何をやっているのか理解を深めました

はじめに

頑張れば、何かがあるって、信じてる。nikkieです。
2019年12月末から自然言語処理のネタで毎週1本ブログを書いています。

先週、総括としてライブドアニュース(日本語テキスト)の分類に取り組んだところ、tf.py_function何をやっているのか分からないという課題が見えました。
そこで、今回はドキュメントを参考にtf.py_functionの理解を深めました。
先週から参照している記事「tf.data.Dataset apiでテキスト (自然言語処理) の前処理をする方法をまとめる - Qiita 」のコードの書き換えにも挑戦しています。

先週の週1ブログ駆動開発!

目次

動作環境

先週と同じ環境です。

$ sw_vers
ProductName:    Mac OS X
ProductVersion: 10.14.6
BuildVersion:   18G103
$ # 注:コードを動かすのにGPUは使っていません
$ python -V
Python 3.7.3
$ pip list  # grepを使って抜粋して表示
tensorflow               2.1.0
tensorflow-datasets      2.0.0
matplotlib               3.1.2
numpy                    1.18.1
Janome                   0.3.10

tf.py_functionとは何かの前に:tf.data.Dataset.mapとは何か

ドキュメントはこちら

必須引数map_funcに焦点を当てます1

引数map_funcは関数(オブジェクト)を受け取り、

Maps map_func across the elements of this dataset.
(データセットの要素全体に map_func を適用する)

mapメソッドがやることは以下の2つです:

  • returns a new dataset((map_funcが適用された) 新しいtf.data.Dataset(以下では「データセット」と表記)オブジェクトを返す)
  • same order as they appeared in the input(要素は入力に現れたのと同じ順番)

map_funcの入力

The input signature of map_func is determined by the structure of each element in this dataset.
map_funcの入力シグネチャは現データセットオブジェクトの各要素の構造によって決定される)

シグネチャ(signature)は馴染みがない言葉なのですが、ここでは「パラメーターとその型」と理解しました。2

ドキュメントにはmap_funcの入力を決める例が3つ続きます:

  1. tf.Tensor1つから構成されたデータセットオブジェクト3
  2. 2つのtf.Tensorからなるタプル(tf.Tensor, tf.Tensor)から構成されたデータセットオブジェクト
  3. 辞書から構成されたデータセットオブジェクト

1の例(dataset.map(lambda x: x + 1))ではxtf.Tensor1つに対応します。

In [2]: dataset = tf.data.Dataset.range(1, 6)

In [7]: list(dataset.as_numpy_iterator())
Out[7]: [1, 2, 3, 4, 5]

In [8]: dataset2 = dataset.map(lambda x: x + 1)

In [9]: list(dataset2.as_numpy_iterator())
Out[9]: [2, 3, 4, 5, 6]

今回知ったデータセットオブジェクトの便利なメソッド

(1) 上の例でもlist()と合わせて使いましたが、as_numpy_iteratorメソッドが便利です。
for inでデータセットオブジェクトを扱わなくても値を見ることができます。

(2)from_generatorメソッドでデータセットオブジェクトを作る方法を学びました。

In [22]: elements = [(1, "foo"), (2, "bar"), (3, "baz")]

In [23]: dataset = tf.data.Dataset.from_generator(lambda: elements, (tf.int32, tf.string))

from_generatorメソッドの

  • generator引数にlambda(引数なし)を渡し、
  • output_types引数で型を定義

しています。

今回使った引数なしのlambda(無名関数)の確認

In [32]: f = lambda: elements

In [33]: f()
Out[33]: [(1, 'foo'), (2, 'bar'), (3, 'baz')]  # elememntsが返る

なお、「dataset2 = tf.data.Dataset.from_tensor_slices(elements)でもいいのでは」と試したところ、

ValueError: Can't convert Python sequence with mixed types to Tensor.

(型が混在したシーケンスはtf.Tensorに変換できない)とのことです。
だからfrom_generatorメソッドがあるのですね。納得!

map_funcの出力

入力を見たので続いて出力を見ます。

The value or values returned by map_func determine the structure of each element in the returned dataset.
map_funcの返り値が、mapで返されるデータセットオブジェクトの各要素の型を決める)

ドキュメントには、タプルを返すいくつかの例が紹介されていました。

  1. 関数gの例:要素が全て定数で変換されたdatasetが返される
  2. 関数hの例:リストとnumpy arrayはともにtf.Tensorに変換される
  3. 関数iはネストされたタプルの例

gの例です。

In [36]: def g(x):
    ...:     return tf.constant(37.0), tf.constant(["Foo", "Bar", "Baz"])
    ...:

In [37]: g(1)  # 実引数の値によらず返り値は同じ
Out[37]:
(<tf.Tensor: shape=(), dtype=float32, numpy=37.0>,
 <tf.Tensor: shape=(3,), dtype=string, numpy=array([b'Foo', b'Bar', b'Baz'], dtype=object)>)

In [38]: type(g(1))
Out[38]: tuple

In [39]: result = dataset.map(g)

In [40]: for d in result:  # 値だけの確認であれば list(result.as_numpy_iterator()) が使える
    ...:     print(d)
    ...:
(<tf.Tensor: shape=(), dtype=float32, numpy=37.0>, <tf.Tensor: shape=(3,), dtype=string, numpy=array([b'Foo', b'Bar', b'Baz'], dtype=object)>)
(<tf.Tensor: shape=(), dtype=float32, numpy=37.0>, <tf.Tensor: shape=(3,), dtype=string, numpy=array([b'Foo', b'Bar', b'Baz'], dtype=object)>)
(<tf.Tensor: shape=(), dtype=float32, numpy=37.0>, <tf.Tensor: shape=(3,), dtype=string, numpy=array([b'Foo', b'Bar', b'Baz'], dtype=object)>)

In [41]: result.element_spec
Out[41]:
(TensorSpec(shape=(), dtype=tf.float32, name=None),
 TensorSpec(shape=(3,), dtype=tf.string, name=None))

tf.Tensor2つからなるタプルがデータセットオブジェクトの1つ1つとなっていることが確認できます(40のセル参照)。

セル41で使っているelement_spec属性は便利そうです。

本題:tf.py_functionとは何か

mapのドキュメントを読んでいくと、map_funcPythonのコードを渡す方法について言及されています4

To use Python code inside of the function you have two options

  1. Rely on AutoGraph to convert Python code into an equivalent graph computation(AutoGraphを用いてPythonコードを等価なグラフ計算に変換する)
    • AutoGraphですべてのPythonコードを変換できるとは限らない
  2. Use tf.py_function
    • 任意のPythonコードを書ける
    • 1に比べて一般にパフォーマンスは劣る

mapのドキュメントからリンクされているtf.py_functionを見ていきます。

引数について

  • funcPythonの関数
  • inptf.Tensorのリスト
    • funcinpと型が対応するtf.Tensorのリストを受け付ける
  • Tout:TensorFlowのデータ型のリストorタプル。funcの返り値を示す(Noneが返るときは空のリストとする)
    • funcToutの値に対応する型のtf.Tensorのリストを返す

返り値は

A list of Tensor or a single Tensor which func computes
funcが計算するtf.Tensorのリストまたはtf.Tensor単体)

tf.py_functionの使い方を試す

mapのドキュメントにあるupper_case_fnの例で確認します(書き換えたいQiita記事との読み替えが減るように、関数名の先頭をアンダースコアで始めています)。

In [75]: d = tf.data.Dataset.from_tensor_slices(['hello', 'world'])

In [76]: def _upper_case_fn(t):
    ...:     return t.numpy().decode().upper()
    ...:

In [77]: d = d.map(lambda x: tf.py_function(func=_upper_case_fn, inp=[x], Tout=tf.string))

In [78]: list(d.as_numpy_iterator())
Out[78]: [b'HELLO', b'WORLD']

mapmap_funcに渡す関数オブジェクトをlambda以外で作ってみます。
lambdaを使っていないだけで、関数の処理内容は変わりません(tf.py_functionの適用結果を返す)。

In [79]: def upper_case_func(x):
    ...:     return tf.py_function(func=_upper_case_fn, inp=[x], Tout=tf.string)
    ...:

In [80]: d = tf.data.Dataset.from_tensor_slices(['hello', 'world'])

In [81]: list(d.map(upper_case_func).as_numpy_iterator())
Out[81]: [b'HELLO', b'WORLD']

というわけで、mapメソッドのmap_func引数には

  • tf.py_functionを返すlambda
  • または、tf.py_functionを返す関数

のどちらかを渡せばよさそうです。

自然言語処理tf.py_functionを使ったコードの書き換えに挑む

tf.data.Dataset apiでテキスト (自然言語処理) の前処理をする方法をまとめる - Qiita 中のコードを書き換えてみます。
例えば、janomeを使って原形を取り出す部分は、関数内関数を多用しているために書き手以外にわかりにくいコードという印象です。
もっとわかりやすく書けないか、自分の理解を深めるために書き直します。

書き換え前のコード

def janome_tokenizer():
    token_filters = [
        CompoundNounFilter(),
        POSStopFilter(["記号", "助詞", "助動詞", "名詞,非自立", "名詞,代名詞"]),
        LowerCaseFilter(),
        ExtractAttributeFilter("surface"),
    ]
    tokenizer = Tokenizer()
    analyzer = Analyzer(tokenizer=tokenizer, token_filters=token_filters)
    analyze_function = analyzer.analyze

    def _tokenizer(text_tensor, label):
        text_str = text_tensor.numpy().decode("utf-8")
        tokenized_text = " ".join(list(analyze_function(text_str)))
        return tokenized_text, label

    return _tokenizer


def tokenize_map_fn(tokenizer):
    def _tokenize_map_fn(text_tensor, label):
        return tf.py_function(
            tokenizer, inp=[text_tensor, label], Tout=(tf.string, tf.int64)
        )

    return _tokenize_map_fn


datasets = datasets.map(tokenize_map_fn(janome_tokenizer()))
  • tokenize_map_fn(janome_tokenizer())mapの実引数)で_tokenize_map_fn関数が返る
  • mapで適用する_tokenize_map_fnの引数text_tensor, labelは、データセットの各要素と一致
    • mapで渡す関数はデータセットの構成要素に対応する数の引数を受け取れる
    • 1つの関数でまとめようとすると、引数はtext_tensor, labelとする必要がある(tf.py_functioninp引数への指定で使う)

1つの関数にまとめてみた

def tokenize_map_fn():
    token_filters = [
        CompoundNounFilter(),
        POSStopFilter(["記号", "助詞", "助動詞", "名詞,非自立", "名詞,代名詞"]),
        LowerCaseFilter(),
        ExtractAttributeFilter("surface"),
    ]
    tokenizer = Tokenizer()
    analyzer = Analyzer(tokenizer=tokenizer, token_filters=token_filters)
    analyze_function = analyzer.analyze

    def _tokenizer(text_tensor, label):
        text_str = text_tensor.numpy().decode("utf-8")
        tokenized_text = " ".join(list(analyze_function(text_str)))
        return tokenized_text, label

    def _tokenize_map_fn(text_tensor, label):
        return tf.py_function(
            _tokenizer, inp=[text_tensor, label], Tout=(tf.string, tf.int64),
        )

    return _tokenize_map_fn


all_tokenized_data = all_labeled_data.map(tokenize_map_fn())
  • mapに渡るのはtokenize_map_fn()、すなわち_tokenize_map_fnという関数オブジェクト
    • _tokenize_map_fnはデータセットの要素を引数に受け取れる(def _tokenize_map_fn(text_tensor, label):
    • _tokenize_map_fntf.py_functionを適用したオブジェクトを返す
  • 関数tokenize_map_fnjanomeのAnalyzerを作った後、それを使った関数_tokenizerを用意し、_tokenize_map_fn関数の中に設定している

動作するがWarningが出る例

def tokenize_map_fn(text_tensor, label):
    token_filters = [
        CompoundNounFilter(),
        POSStopFilter(["記号", "助詞", "助動詞", "名詞,非自立", "名詞,代名詞"]),
        LowerCaseFilter(),
        ExtractAttributeFilter("surface"),
    ]
    tokenizer = Tokenizer()
    analyzer = Analyzer(tokenizer=tokenizer, token_filters=token_filters)
    analyze_function = analyzer.analyze

    def _tokenizer(text_tensor, label):
        text_str = text_tensor.numpy().decode("utf-8")
        tokenized_text = " ".join(list(analyze_function(text_str)))
        return tokenized_text, label

    return tf.py_function(
        _tokenizer, inp=[text_tensor, label], Tout=(tf.string, tf.int64),
    )


all_tokenized_data = all_labeled_data.map(tokenize_map_fn)

mapに直接渡せるようにtokenize_map_fntext_tensor, labelを引数で渡せるようにしました。
ですが、以下のようなWarningが大量に出ます。

WARNING:tensorflow:AutoGraph could not transform <bound method Tokenizer.tokenize of <janome.tokenizer.Tokenizer object at 0x13e673860>> and will run it as-is.
WARNING:tensorflow:Entity <bound method CompoundNounFilter.apply of <janome.tokenfilter.CompoundNounFilter object at 0x13e6732b0>> appears to be a generator function. It will not be converted by AutoGraph.

おそらくPythonのコードが全て、tf.py_functionで包まれていないのが原因ではないかと見ています。

感想

maptf.py_functionのドキュメントからだいぶ理解は深まりました。
ですが、実務レベルのコードでの適用方法は、書き換えでwarningに出会ったこともあり、まだ自信がありません(時間切れ感もあります)。

先週感じた課題の中からTensorFlowを今後使っていく上で影響しそうなtf.py_functionを選択しました。
まだまだ課題は残っていますし、新しい調査事項も出てきました。
Qiita記事のencodeの部分、私のコードではset_shapeが見られ、これを削除するとエラーになります。
これが何をするものなのか気になっています。

課題は山積していますが、いったんそれらは脇に置いて(時間が余った時にやるように優先度を下げて)、次週からは自然言語処理の基礎を固めるために、「禁書にすべき」という声もあるあの本に取り組みます。


  1. 今回はtf.py_functionがよく分からない問題に取り組んでおり、tf.py_functionが関わってくるのがmap_funcのためです。

  2. Signature (functions) (シグネチャ (関数)) - MDN Web Docs 用語集: ウェブ関連用語の定義 | MDN

  3. ドキュメントの上部「Source Datasets」の記載から、データセットオブジェクトをfor inで取り出したときの1つ1つがtf.Tensorと理解しています

  4. ここまでの例で単純なlambdaは渡せることが確認できたので、グラフとして扱う(tf.data traces the function and executes it as a graph)中でPythonのコードを使う方法のようです