TensorFlowでカスタムロスを作成して、TensorFlowでの処理の感覚を覚えよう #Python

本稿ではTensorFlowでの自作のLoss関数を作成することで、機械学習でポイントとなる、まとめてデータを処理するという感覚にちょっと親しんでみたいと思います。TensorFlowの初歩からちょっと踏み出してみましょう。

標準のMnistプログラム

まずは簡単なmnistの画像分類プログラムから開始します。これは皆さんも一度はやったことのある題材ではないでしょうか。さっそく取り組んでみます。

今回はTensorFlowに加えて、TensorFlow Datasetsをライブラリとして使用します。

>pip install tensorflow -U
>pip install tensorflow_datasets -U

mnist.py

import tensorflow as tf
import tensorflow_datasets as tfds

# データセット変換処理用関数定義
def transform(data):
    image = data['image']
    image = tf.cast(image, tf.float32)
    image = image / 255.0
    label = tf.one_hot(data['label'], 10)
    return image, label

# 学習モデル取得用関数
def get_model():
    input = tf.keras.layers.Input(shape=[28,28,1])
    x = tf.keras.layers.Conv2D(filters=64, kernel_size=(3,3), padding='same')(input)
    x = tf.keras.layers.Conv2D(filters=64, kernel_size=(3,3), padding='same')(x)
    x = tf.keras.layers.Conv2D(filters=32, kernel_size=(3,3), padding='same')(x)
    x = tf.keras.layers.Flatten()(x)
    x = tf.keras.layers.Dense(128,activation='relu')(x)
    output = tf.keras.layers.Dense(10,activation='softmax')(x)
    return tf.keras.Model(inputs=[input],outputs=[output])

# データセット読み込み
mnist = tfds.load('mnist',split='train').map(transform).batch(50)

# モデルの作成
model = get_model()

# モデルのコンパイル
model.compile(optimizer='rmsprop',loss='categorical_crossentropy',metrics=['accuracy'])

# モデルの学習
model.fit(mnist,batch_size=50,epochs=3)

簡単なロス関数(2乗誤差)

ここで使用したロス関数を自作のもので差し替えてみましょう。作成する関数はy_pred(現在の学習したモデルでの推論値)と、y_true(one-hot化された正解ラベル情報)を入力され、ロス値を返却するような以下の関数です。

def my_loss( y_pred:tf.Tensor , y_true:tf.Tensor )->tf.Tensor:

まずは簡単な例として、推論結果と教師ラベル(one-hot化されたもの）の2乗誤差を返却するカスタムロスを作成してみたいと思います。

ここでのポイントはloss関数ではy_pred,y_trueともに一度の処理のタイミングで、データがバッチサイズ個数分やってくることです。バッチサイズの単位でデータ処理が回っているということです。

y_pred（機械学習の推論結果）,y_true（教師情報）の関数の入力値は[バッチサイズ,画像サイズ,画像サイズ,チャンネル数]次元の配列として、まとめてやってきますし、ロス関数の返却値は[バッチサイズ,ロス値]次元で返却する必要があります。

つまり、y_pred, y_trueで入力されたデータをバッチサイズの個数のロス値へと集約するような処理になります。

こういった関数を実現するには当たり、他の言語で慣れてきた方は、for文などでデータの要素を取り出してロスを計算するのが第一感かもしれません。しかし、TensorFlowのような機械学習の世界では、極力、データ全体をまとめて、一度で扱うように記述する必要があります。(Pythonの中でのループ処理は非常に重たいです。）

たとえば、今回の例ではy_predとy_trueは10次元の値になっています。10回のループをしてしまいがちです。が、そうはしてはいけません。各要素ごとに差をとって2乗を計算する演算と、その結果を集約する演算で、２乗誤差を実現します。こうすることで、重たいPythonの中でループ処理をしないようにして、その計算全体をGPUやC言語のプログラムへ極力、大きな単位で処理を依頼することで、高速に処理を実行させることができます。

def my_loss(y_pred, y_true):
    # y_pred [[0.1,0.8,....,(10クラス分の予測値)],[....][....]..(バッチサイズ個分)]
    # y_true [[1.0,0.0,0.0....,(正解のone-hot値)],[....][....]..(バッチサイズ個分)]

    sd = tf.math.squared_difference(y_pred,y_true) #全ての要素ごとに差をとった２乗を演算

    rs = tf.math.reduce_sum(sd, axis=1) # 要素をデータごとに集計

    return rs

一般にreduce_xxxは要素を集計するメソッドです。データに含まれている各クラスごとの推論値の集計になるので、axis=1を指定します。

これをcompileして組み込んで利用します。

model.compile(optimizer='adam',loss=my_loss,metrics=['accuracy'])

コサイン類似度ロス

次の演習課題として、推論値と正解ラベルの内積値をベースにしたロスを作ってみましょう。内積は2つのベクトルが一致すると1.0になるので、1―内積値としてみます。

ここでもtf.math.reduce_sum()を使いデータごとの集計をおこないます。

def my_loss2(y_pred, y_true):
    # y_pred [[0.1,0.8,....,(10クラス分の予測値)],[....][....]..(バッチサイズ個分)]
    # y_true [[1.0,0.0,0.0....,(正解のone-hot値)],[....][....]..(バッチサイズ個分)]

    mul = tf.math.multiply(y_pred, y_true) # 要素ごとに掛け算
    rs = tf.math.reduce_sum(mul, axis=1) # データごとに集計
    return 1.0-rs

まとめ

TensorFlow(tf.keras)でカスタムロスを作成し、適用してみました。

カスタムロスを導入するには、データがバッチサイズの単位で処理されることを意識したTensor操作が必要です。

本稿ではバッチサイズ単位でロス値を返却する、データ処理のパターンとして、データの各要素ごとに計算し、reduce系の演算でデータごとにまとめる処理パターンを紹介しました。

これをベースに様々なロス関数作成に挑戦してみましょう。