[TensorFlow 2.x対応版] TensorFlow (Keras) で TFRecord & DataSetを使って大量のデータを学習させる方法

はじめに

以前の記事のアップデート版になります。
TensorFlow & Keras で TFRecord & DataSetを使って大量のデータを学習させる方法 - Qiita

やりたいことは一点、「メモリに入り切らない巨大データを学習させる効率的な方法がほしい」ということです。
CPUのデータ読み込みとGPUの計算が並行処理できるような方法ですね。
DataSet APIを使って、特定のフォーマットで保存したデータからの学習を効率的に行っていきます。

TensorFlow 2がリリースされ、以前のバージョンの記事と比べてモジュールの名前が変わったり、一部の処理が簡単に書けるようになったりしました。この記事では以前との差分を中心に、TensorFlow 2での書き方を紹介します。また、ついでにKerasもTensorFlowに内包されているものを使うように変えてみます。

事前準備

この記事ではLinux (Ubuntu 18.04) 上で Python 3.6.9 + TensorFlow 2.1.0 を使います。

TensorFlow 1.15/2.1から、CPUバージョンとGPUバージョンのpipパッケージが統合されました。よってCPUで手軽に試したい人もGPUで本格的に回したい人も

pip3 install tensorflow==2.1.0

これでOKです。なお、GPUを使いたい方はCUDA 10.1のセットアップが必要ですのでご注意ください。
GPU support | TensorFlow

データ準備

TensorFlowで効率よく計算させるための独自データフォーマット (TFRecord) があります。DataSet APIを使って、既存のデータからTFRecordを作っておきましょう。

data2tfrecord.py

#!/usr/bin/env python3

import numpy as np
import tensorflow as tf
from tensorflow.keras.datasets import mnist

def feature_float_list(l):
    return tf.train.Feature(float_list=tf.train.FloatList(value=l))

def record2example(r):
    return tf.train.Example(features=tf.train.Features(feature={
        "x": feature_float_list(r[0:-1]),
        "y": feature_float_list([r[-1]])
    }))

filename_train = "train.tfrecords"
filename_test  = "test.tfrecords"

# === MNISTデータを読み込む ===
# 簡単のため、学習中の検証データに評価データと同じものを使うとする
(x_train, y_train), (x_test, y_test) = mnist.load_data()
print("x_train   : ", x_train.shape) # x_train   :  (60000, 28, 28)
print("y_train   : ", y_train.shape) # y_train   :  (60000,)
print("x_test    : ", x_test.shape)  # x_test    :  (10000, 28, 28)
print("y_test    : ", y_test.shape)  # y_test    :  (10000,)

# 前処理をする
# 画素は[0, 1]のfloat32型に変換する
# さらに、TFRecord化のために、特徴量を1次元にしておく（行がレコードに対応）
x_train = x_train.reshape((-1, 28*28)).astype("float32") / 255.0
x_test  = x_test.reshape((-1, 28*28)).astype("float32") / 255.0
# ラベルもfloat32型にする
y_train = y_train.reshape((-1, 1)).astype("float32")
y_test  = y_test.reshape((-1, 1)).astype("float32")
# TFRecord化するために、特徴量とラベルを結合する
data_train = np.c_[x_train, y_train]
data_test = np.c_[x_test,  y_test]

# 実際には、学習したいデータを同じ形式に変換して作る。
# 全データがメモリに乗り切らない場合は、以下の書き込みフェーズで
# 少しずつ作って書き込むことを繰り返せばよい。

# 学習データをTFRecordに書き込む
with tf.io.TFRecordWriter(filename_train) as writer:
    for r in data_train:
        ex = record2example(r)
        writer.write(ex.SerializeToString())

# 評価データをTFRecordに書き込む
with tf.io.TFRecordWriter(filename_test) as writer:
    for r in data_test:
        ex = record2example(r)
        writer.write(ex.SerializeToString())

ほとんど前回と同じなのですが、TensorFlowのバージョンアップに伴って tensorflow.python_io というパッケージがなくなり、TFRecord関係の関数は tensorflow.io に入りました。
また、TensorFlow内包のKerasを使うようにしたので import が変わっていますが、MNISTデータセットの読み込み方法自体は変わっていませんでした。

なお、GPU計算用のライブラリが整っていないとCUDAがらみのWARNING（libcublasが見つからないなど）が出ますが、CPUで軽く試したいだけの人は気にしなくて大丈夫です。

学習

前回とちょこちょこ変わっています。まずコードを紹介して、次に差分について説明します。

train.py

#!/usr/bin/env python3

import numpy as np
import tensorflow as tf
from tensorflow.keras.layers import Dense, Input
from tensorflow.keras.optimizers import RMSprop
from tensorflow.keras.callbacks import ModelCheckpoint
from tensorflow.keras.models import Model

# 学習設定
batch_size = 32
epochs = 10
# 特徴量の設定
num_classes = 10    # ラベルの種類。0-9の10種類
feature_dim = 28*28 # 特徴量の次元。簡単のため1次元のままで扱う
# 学習・評価データ件数。事前に調べておく。
# 複数のTFRecordを使う場合、以下の件数は全ファイルの合計になることに注意。
num_records_train = 60000
num_records_test  = 10000
# 1エポックあたりのミニバッチ数。学習時に使う。
steps_per_epoch_train = (num_records_train-1) // batch_size + 1
steps_per_epoch_test  = (num_records_test-1) // batch_size + 1

# 1件のTFRecordをデコード
def parse_example(example):
    features = tf.io.parse_single_example(
        example,
        features={
            # リストを読み込む場合は次元数を指定する
            "x": tf.io.FixedLenFeature([feature_dim], dtype=tf.float32),
            "y": tf.io.FixedLenFeature([], dtype=tf.float32)
        })
    x = features["x"]
    y = features["y"]
    return x, y

# === TFRecordファイルのデータを学習・評価用に準備 ===

dataset_train = tf.data.TFRecordDataset(["train.tfrecords"]) \
    .map(parse_example) \
    .shuffle(batch_size * 100) \
    .batch(batch_size).repeat(-1)
# 上で複数のTFRecordファイルを用いる場合は、ファイル名のリストを指定する。
# dataset_train = tf.data.TFRecordDataset(["train.tfrecords.{}".format(i) for i in range(10)]) \

dataset_test = tf.data.TFRecordDataset(["test.tfrecords"]) \
    .map(parse_example) \
    .batch(batch_size)

# === モデル定義 ===
# 今回は512次元の中間層を1層だけ指定している
layer_input = Input(shape=(feature_dim,))
fc1 = Dense(512, activation="relu")(layer_input)
layer_output = Dense(num_classes, activation="softmax")(fc1)
model = Model(layer_input, layer_output)
model.summary()

# ラベルがカテゴリ変数の場合でも loss="sparse_categorical_crossentropy" で学習できる
# ラベルをone-hotベクトル化した場合は、loss="categorical_crossentropy" になる
model.compile(
    loss="sparse_categorical_crossentropy",
    optimizer=RMSprop(),
    metrics=["accuracy"])

# === 学習 ===

# 途中のモデルを保存しておく
cp_cb = ModelCheckpoint(
    filepath="weights.{epoch:02d}-{loss:.4f}-{val_loss:.4f}.hdf5",
    monitor="val_loss",
    verbose=1,
    save_best_only=True,
    mode="auto")
model.fit(
    x=dataset_train,
    epochs=epochs,
    verbose=1,
    steps_per_epoch=steps_per_epoch_train,
    validation_data=dataset_test,
    validation_steps=steps_per_epoch_test,
    callbacks=[cp_cb])

前回との差分

tensorflow.keras.Model.fit() が、学習データにDataSetを取れるように変わっていました。
tf.keras.Model | TensorFlow Core v2.1.0

x: Input data. It could be:
（中略）
A tf.data dataset. Should return a tuple of either (inputs, targets) or (inputs, targets, sample_weights).

以前は、DataSetから学習するときには Input レイヤーにデータを差し込むしかなかったので、重みを共有する2つのモデルを学習用・評価用にそれぞれ作るという面倒な手順がありました。TensorFlow 2.x（に内包のKeras）では Model.fit() にDataSetを与えることができるので、モデルは1個で大丈夫です。make_one_shot_iterator() でイテレータを自分で作るといったことも不要になりました。やったね！

さらに、tensorflow.keras.Model.fit() の validation_data 引数にも評価用のDataSetを与えることができるようになっていました。よって、評価用のコールバックを自分で作る必要もなくなりました（評価時のプログレスバーが出なくなってしまいますが…。自分で学習ループを書けという話ですね）。

複数のTFRecordを用いたパフォーマンス改善

複数ファイルを並列に読み込ませるようにすることで、GPUの使用率を上げられる（＝学習を高速化できる）場合があります。

前回と同じ要領で学習データを分割して書き込みます。前回との違いは tf.python_io が tf.io に変わっただけです。

data2tfrecord.py（一部）

for i in range(10):
    with tf.io.TFRecordWriter(filename_train + "." + str(i)) as writer:
        for r in data_train[i::10]:
            ex = record2example(r)
            writer.write(ex.SerializeToString())

学習時は、dataset_train の作り方が以下のように変わります。

train.py（一部）

dataset_train = tf.data.Dataset.from_tensor_slices(["train.tfrecords.{}".format(i) for i in range(10)]) \
    .interleave(
        lambda filename: tf.data.TFRecordDataset(filename).map(parse_example, num_parallel_calls=1),
        cycle_length=10) \
    .shuffle(batch_size * 100) \
    .batch(batch_size) \
    .prefetch(1) \
    .repeat(-1)

前回の記事でいうところの tf.contrib.data.parallel_interleave() （後に tf.data.experimental.parallel_interleave()）に相当する機能が、正式にDataSetのメソッドとして取り込まれましたので、記述が少し簡単になりました。ただし sloppy=False 相当の動作をするので、sloppy=True 相当の動作にするには with_options() でオプション指定する必要があるようです。
tf.data.experimental.parallel_interleave | TensorFlow Core v2.1.0

みんなもTensorFlow 2に移行してみよう

多少の変更点はありますが、概して簡単に書ける方向に変わっているので、あまり恐れなくてもよい気がしました。
コアのパフォーマンスも改善されていることが期待できますし（本当か？）、TensorFlow 2でも大量データをガンガン突っ込んでみましょう！