taka5hi’s blog

統計と機械学習の話題をメインに記事を書いています。

TFRecord、DataSet API を Keras で使う

機械学習では、時にはメモリに収まりきらないほどの大量のデータを扱う必要があります。
データを準備・加工する処理がボトルネックにならないようにするためには、例えば以下のような工夫が必要になります。

  • 複数のデータをまとめて単一のファイルに保存することでファイルI/Oにかかるオーバーヘッドを減らす
  • データをファイルから読み出してミニバッチを作成する一連のパイプラインを別スレッドで行う

TensorFlowでは、前者に対しては TFRecord というバイナリ形式、後者に対しては DataSet API という強力な仕組みをサポートしています。
普段、Keras を使うことが多いのですが、Keras でも TensorFlow の便利な仕組みを使えないかと思いまとめてみました。

前提

下記のようなケースを想定します。

  • 訓練データについては TFRecord を DataSet API を使って処理する
  • 検証、テストデータはメモリ内に読み込み済みのものを使う
  • Keras のバックエンドとして、Tensorflow を使う

大量の訓練データについては DataSet API を使い、数の少ない検証・テストデータに対してはメモリ内から直接参照するというのは、現実にもありがちな構成ではないかと思います。

サンプルコード

定数定義

まずは必要な定数を定義します。

BATCH_SIZE = 100
EPOCHS = 5
STEPS_PER_EPOCH = 100

TFRecord の作成

次に MNIST データを使って、TFRecord を作ります。
ここについては、TensorFlow で TFRecord を作る処理と変わりませんので説明は省略します。

from keras.datasets import mnist
from keras.utils import to_categorical
import tensorflow as tf

def save_to_tfrecord(datas, labels, filename):
    with tf.python_io.TFRecordWriter(filename) as w:
        for data, label in zip(datas, labels):
            data = data.reshape(-1).tobytes()
            features = tf.train.Features(feature={
                'data': tf.train.Feature(bytes_list=tf.train.BytesList(value=[data])),
                'label': tf.train.Feature(int64_list=tf.train.Int64List(value=label.astype(dtype='int64')))
            })

            example = tf.train.Example(features=features)
            w.write(example.SerializeToString())

# MNIST データのロード
(x_train, y_train), _ = mnist.load_data()

# 訓練データの加工と保存
train_labels = to_categorical(y_train)
save_to_tfrecord(x_train, train_labels, "train.tfrecord")

TFRecord の読み込み、DataSet API を使ったデータ準備

次に、TFRecord を読み込んで、訓練データを準備するところです。
ここも、TensorFlow でのやり方通りに行って下さい。

def parse_example(example):
    features = tf.parse_single_example(
        example,
        features={
            'data': tf.FixedLenFeature([], dtype=tf.string),
            'label': tf.FixedLenFeature((10,), dtype=tf.int64)
        })
    
    img = features['data']
    img = tf.decode_raw(img, tf.uint8)
    
    label = features['label']
    
    return img, label

dataset = tf.data.TFRecordDataset(['train.tfrecord']).map(parse_example)
dataset = dataset.repeat(-1).batch(BATCH_SIZE)
iterator = dataset.make_one_shot_iterator()

# 訓練データの加工・準備
train_imgs, train_labels = iterator.get_next()
train_imgs = tf.cast(train_imgs, dtype=tf.float32) / 255
train_imgs = tf.reshape(train_imgs, (-1, 28, 28, 1))
train_labels = tf.cast(train_labels, dtype=tf.float32)
train_labels = tf.reshape(train_labels, (-1, 10))

モデル構築

訓練用のモデル、テスト・検証用のモデルという2つのモデルを構築することがポイントとなります。
そして、訓練用のモデルでは、明示的に Tensor を渡してモデルを構築します。
一方、テスト・検証用のモデルは通常通りに作れば大丈夫です。ただし、訓練用のモデルと重みを共有することに注意してください。

from keras import layers
from keras import models
from keras.engine.network import Network

# 基本的な CNN の構築
input_layer = layers.Input(shape=(28, 28, 1))
output = layers.Conv2D(32, (3, 3), activation='relu')(input_layer)
output = layers.MaxPooling2D((2, 2))(output)
output = layers.Conv2D(64, (3, 3), activation='relu')(output)
output = layers.MaxPooling2D((2, 2))(output)
output = layers.Conv2D(64, (3, 3), activation='relu')(output)
output = layers.Flatten()(output)
output = layers.Dense(64, activation='relu')(output)
output = layers.Dense(10, activation='softmax')(output)

# 2つのモデルを作成するため共通部分をまとめて取り回しをしやすくする
common_network = Network(input_layer, output)

# 訓練用と学習用に重みを共有する二つのモデルを作る
train_in_layer = layers.Input(tensor=train_imgs)    # tensor を指定
test_in_layer = layers.Input(shape=(28, 28, 1))

train_output = common_network(train_in_layer)
test_output = common_network(test_in_layer)

train_model = models.Model(inputs=train_in_layer, outputs=train_output)
test_model = models.Model(inputs=test_in_layer, outputs=test_output)

train_model.compile(optimizer='rmsprop',
                    loss='categorical_crossentropy',
                    metrics=['accuracy'],
                    target_tensors=[train_labels])    # target_tensor を指定
test_model.compile(optimizer='rmsprop', 
                   loss='categorical_crossentropy',
                   metrics=['accuracy'])

訓練

すでに入力と正解データを Tensor として指定しているので、fit メソッドの引数としてはデータを与える必要がありません。

epoch_history = train_model.fit(epochs=EPOCHS, steps_per_epoch=STEPS_PER_EPOCH)

検証・テスト

検証・テストは、test_model を使って行います。

_, (x_test, y_test) = mnist.load_data()

test_imgs = x_test.reshape((*x_test.shape, 1))
test_imgs = test_imgs.astype('float32') / 255
test_labels = to_categorical(y_test)

test_loss, test_acc = test_model.evaluate(test_imgs, test_labels)

まとめ

サンプルコードが少なく、結構はまったのですが、
「訓練用と検証・テスト用の2つのモデルを構築する」というポイントさえ分かってしまえば、意外と簡単に実現できるようです。
ただ、k分割交差検証を実現しようと思うと、さらにモデルを複数作る必要がるため多少煩雑にはなってくるかもしれません。
(k 組のデータセットができるので、それ毎に訓練用もモデルを作っておく必要があるため)