はじめに
以前の記事のアップデート版になります。
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を作っておきましょう。
#!/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で軽く試したいだけの人は気にしなくて大丈夫です。
学習
前回とちょこちょこ変わっています。まずコードを紹介して、次に差分について説明します。
#!/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
に変わっただけです。
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
の作り方が以下のように変わります。
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でも大量データをガンガン突っ込んでみましょう!