32
23

Delete article

Deleted articles cannot be recovered.

Draft of this article would be also deleted.

Are you sure you want to delete this article?

More than 3 years have passed since last update.

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

Last updated at Posted at 2020-03-12

はじめに

以前の記事のアップデート版になります。
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_iotf.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でも大量データをガンガン突っ込んでみましょう!

32
23
0

Register as a new user and use Qiita more conveniently

  1. You get articles that match your needs
  2. You can efficiently read back useful information
  3. You can use dark theme
What you can do with signing up
32
23

Delete article

Deleted articles cannot be recovered.

Draft of this article would be also deleted.

Are you sure you want to delete this article?