🕌

TensorFlowの公式チュートリアル「tf.dataを使って画像をロードする」をもうちょっとスリムにして読む

2020/12/19に公開

Python

TensorFlowのチュートリアル「tf.dataを使って画像をロードする」は、けっこう長くて、コードを１つ１つ理解しないといけないので、もう少しスリムにしつつ噛み砕いたものを書きたいと思います。

Google Colaboratoryを使って書いていきますが、別の環境でも動作します。

TL;DR

スリムにしても長かったので、なるべく省略したい方は
（最終形）TFRecordを使って性能を上げる
までスキップしてご覧ください。これが一番性能の出る方法です。

画像を読み込んで表示する

画像を読み込んで表示させてみます。

ポイントは、tf.keras.utils.get_fileで、これを使えばURLから画像をダウンロードして解凍までしてくれます。あとは、Pythonのライブラリを用いて表示しているだけです。

NoteBook

import tensorflow as tf
import pathlib
import random
import IPython.display as display

# ファイルをURLの位置からダウンロードする。
# untar=Trueとなっているので、解凍までしてくれる。
# 戻り値は解凍後のディレクトリ名
data_root_orig = tf.keras.utils.get_file(origin='https://storage.googleapis.com/download.tensorflow.org/example_images/flower_photos.tgz',
                                        fname='flower_photos', untar=True)

# pathLibのインスタンスを作成。
# pathLibはファイルやディレクトリをオブジェクトとして操作できる
data_root = pathlib.Path(data_root_orig)

# サブディレクトリにあるファイルをすべて抽出する
all_image_paths = list(data_root.glob('*/*'))
# 抽出したファイルのパスをすべて文字列にする
all_image_paths = [str(path) for path in all_image_paths]

for n in range(3):
 # どれか一つをダンラムに選択する
 image_path = random.choice(all_image_paths)
 print(image_path)
 # 表示する
 display.display(display.Image(image_path))

こんな感じで画像が表示されます。

1. データセットを構築する

画像を機械学習に読み込ませるデータセットにするには、次のような要件を考慮する必要があります。

Tensor型になっていること
モデルに合わせてリサイズしていること
モデルに合わせて正規化、標準化されていること
ラベルと画像のペアになっていること
よくシャッフルされていること
バッチ化されていること
限りなく繰り返し読み込めること
データがなるべく早く読み込めること（CPUの処理がGCUの処理のボトルネックにならないこと）

ということで、１つ１つの方法を説明します。

画像をモデルが読める形に変換する

次の要件を満たす処理です。

Tensor型になっていること
モデルに合わせてリサイズしていること
モデルに合わせて正規化、標準化されていること

後の処理で使用するためload_and_preprocess_imageという関数にまとめています。

NoteBook

def load_and_preprocess_image(path):
  # ファイルを１つ読み込む。ファイルを表現した生データのテンソルが得られる
  image = tf.io.read_file(path)
  # 生データのテンソルを画像のテンソルに変換する。
  # これによりshape=(240,240,3)、dtype=uint8になる
  image = tf.image.decode_jpeg(image, channels=3)
  # モデルに合わせてリサイズする
  image = tf.image.resize(image, [192, 192])
  # モデルに合わせて正規化する（値を0〜1の範囲に収める処理）
  image /= 255.0
  return image

# 上の関数にファイルパスを1つ与えて画像を読み込む
img = load_and_preprocess_image(all_image_paths[0])

print(repr(img)[:100]+"...")

# 出力結果
# <tf.Tensor: shape=(192, 192, 3), dtype=float32, numpy=
array([[[0.22071078, 0.27156863, 0.05196078],...

ポイントは、read_fileで画像を読み込んだ時点で、Tensor型になっていることです。あとは、画像に変換したりリサイズしたりしています。

tf.data.Datasetを構築する

tf.data.Datasetを用いると、画像をシャッフルしたりバッチに切り出したりしやすくなります。つまり残りの要件を満たすためにはtf.data.Datasetを使う必要があります。

画像のパス名をtf.data.Datasetにするには、from_tensor_slicesを使って画像パス名のリストをスライスします。すると、画像のパス名をイテレートできるTensor型のコレクションに変換してくれます。

NoteBook

path_ds = tf.data.Dataset.from_tensor_slices(all_image_paths)

for i in path_ds.take(3):
  print(i)
  
# 出力結果
# tf.Tensor(b'/root/.keras/datasets/flower_photos/dandelion/16863587471_cc3a6ffb29_m.jpg', shape=(), dtype=string)
# tf.Tensor(b'/root/.keras/datasets/flower_photos/dandelion/18996957833_0bd71fbbd4_m.jpg', shape=(), dtype=string)
# tf.Tensor(b'/root/.keras/datasets/flower_photos/dandelion/7401173270_ebaf04c9b0_n.jpg', shape=(), dtype=string)

得られたパス名のスライスを、さきほど作ったload_and_preprocess_imageにmapメソッドで渡すと、イメージを読み込んでくれます。

NoteBook

# AUTOTUNEはGPUの処理とCPUの処理の配分を動的に設定してくれるパラメータ
AUTOTUNE = tf.data.experimental.AUTOTUNE
image_ds = path_ds.map(load_and_preprocess_image, num_parallel_calls=AUTOTUNE)

for image in image_ds.take(3):
  print(image)
  
# 出力結果
# 
# tf.Tensor(
# [[[0.22071078 0.27156863 0.05196078]
#   [0.22052696 0.26666668 0.05091912]
#   [0.22420344 0.2639093  0.05459559]
#   ...

ラベルと画像のペアを作る

次の要件を満たす処理です。

ラベルと画像のペアになっていること

まずは、画像すべてのラベル名をリストにします。

NoteBook

# ディレクトリ名からラベル名を得る
label_names = sorted(item.name for item in data_root.glob('*/') if item.is_dir())
# label_names =  ['daisy', 'dandelion', 'roses', 'sunflowers', 'tulips']
# ラベルに番号をつけ辞書型に登録する
label_to_index = dict((name, index) for index,name in enumerate(label_names))
# label_to_index = {'daisy': 0, 'dandelion': 1, 'roses': 2, 'sunflowers': 3, 'tulips': 4}
# すべての画像データのパスからディレクトリ名を元にラベル番号を得てリストに入れる
all_image_labels = [label_to_index[pathlib.Path(path).parent.name]
                    for path in all_image_paths]
# all_image_labels = [1, 1, 1, 1, 1, 1, 1, 1, 1, 1　・・・・
# all_image_labelsをスライスしてデータセットにする
label_ds = tf.data.Dataset.from_tensor_slices(tf.cast(all_image_labels, tf.int64))

for d in label_ds.take(3):
  print(d)
# 出力結果
# 
# tf.Tensor(1, shape=(), dtype=int64)
# tf.Tensor(1, shape=(), dtype=int64)
# tf.Tensor(1, shape=(), dtype=int64)

ラベルのデータセットと画像のデータセットをzipしてペアを作ります。

NoteBook

image_label_ds = tf.data.Dataset.zip((image_ds, label_ds))
for d in image_label_ds.take(3):
  print(d)
  
# 出力結果
# (<tf.Tensor: shape=(192, 192, 3), dtype=float32, numpy=
# array([[[0.22071078, 0.27156863, 0.05196078],
#         [0.22052696, 0.26666668, 0.05091912],
#         [0.22420344, 0.2639093 , 0.05459559],
# ・・・, <tf.Tensor: shape=(), dtype=int64, numpy=1>)

2つのデータセットが組み合わさっているのが分かると思います。

シャッフル、バッチ化、繰り返し、早い読み込み

残りの要件の処理を一気に作ります。

よくシャッフルされていること
バッチ化されていること
限りなく繰り返し読み込めること
データがなるべく早く読み込めること（CPUの処理がGCUの処理のボトルネックにならないこと）

これらはデータセットにしたあとはかなり直感的に使えます。

NoteBook

BATCH_SIZE = 32

# イメージの数を数える
image_count = len(all_image_paths)

# 全体をシャッフル
ds = image_label_ds.shuffle(buffer_size=image_count)
# 繰り返す
ds = ds.repeat()
# バッチサイズで切り出す
ds = ds.batch(BATCH_SIZE)
# モデル訓練中にバックグラウンドで読み込む
ds = ds.prefetch(buffer_size=AUTOTUNE)

# データセットをイテレータにする
it = iter(ds)
# 1つ取り出す
d = next(it)
print(d)

# 出力結果
#  
# (<tf.Tensor: shape=(32, 192, 192, 3), dtype=float32, numpy=
# array([[[[9.94923472e-01, 1.00000000e+00, 9.98845041e-01],
#          [9.92156863e-01, 1.00000000e+00, 9.96078432e-01],
#          [9.93187726e-01, 9.96037543e-01, 9.94612634e-01],
# ・・・・, <tf.Tensor: shape=(32,), dtype=int64, numpy=
# array([3, 4, 1, 0, 1, 0, 3, 3, 2, 1, 3, 0, 0, 4, 1, 1, 4, 0, 4, 3, 0, 2,
#        1, 3, 1, 1, 3, 0, 2, 1, 2, 3])>)

シャッフルされてバッチサイズ分になったデータが返っているのが分かると思います。
このあと、d = next(it)とprint(d)をずっと繰り返しても永遠に取得し続けることが出来ます。これはrepeat()を指定しているためです。

なお、shuffleとrepeatとbatchの順番によってエポック境界を超えてシャッフルされるかどうかなどの動作が違います。詳しくは公式のチュートリアルをご覧ください。

また、データセットをシャッフルしてリピートすると、リピートのタイミングで待ち時間が発生しますが、シャッフルとリピートを組み合わせたshuffle_and_repeatを使うと回避することができます。

NoteBook

ds = image_label_ds.apply(
  tf.data.experimental.shuffle_and_repeat(buffer_size=image_count))
ds = ds.batch(BATCH_SIZE)
ds = ds.prefetch(buffer_size=AUTOTUNE)

2. モデルに読み込ませる

画像データをMobileNet v2をベースに転移学習させます。

データセットは、先程作ったものを使いますが、Mobile Netにあわせて入力を[-1,1]の範囲に正規化しています。

NoteBook

# 値を[-1,1]の範囲にする関数
def change_range(image,label):
    return 2*image-1, label

# シャッフルとリピート（先程出たものと同じ）
ds = image_label_ds.apply(
  tf.data.experimental.shuffle_and_repeat(buffer_size=image_count))
ds = ds.batch(BATCH_SIZE)
ds = ds.prefetch(buffer_size=AUTOTUNE)
# 値を[-1,1]の範囲にする
keras_ds = ds.map(change_range)

次に、一気にモデルを作って学習までさせてみます。

NoteBook

# MobileNet v2を取得
mobile_net = tf.keras.applications.MobileNetV2(input_shape=(192, 192, 3), include_top=False)
# MobileNetの重みを訓練不可にする。
# 転移学習なので元のモデルは固定し、後から加えたレイヤーだけを学習させる
mobile_net.trainable=False

# レイヤーを追加する
model = tf.keras.Sequential([
    mobile_net,
    tf.keras.layers.GlobalAveragePooling2D(),
    tf.keras.layers.Dense(len(label_names))])

# モデルをコンパイルする
model.compile(optimizer=tf.keras.optimizers.Adam(), 
              loss='sparse_categorical_crossentropy',
              metrics=["accuracy"])

# モデルを訓練する
model.fit(ds, epochs=1, steps_per_epoch=3)  

# 出力結果
#
# 3/3 [==============================] - 11s 804ms/step - loss: 9.1035 - accuracy: 0.1185
# <tensorflow.python.keras.callbacks.History at 0x7f3e95b9bc88>

学習ができています。

3. 読み込み性能を上げる

読み込み性能を上げていきます。

キャッシュを使って性能を上げる方法や、TFRecordという読み込みに最適化されたファイルフォーマットを利用して性能を上げる方法などを説明していきます。

最初に、時間計測用のtimeitという関数を作っておきます。
この関数は、データセットを読み込むのにかかった時間を計測してくれます。計測に使うだけなので細かい説明は入れていません。

NoteBook

import time

# 1エポックあたりのステップ数を決定する
# すべての画像を1エポックで読み切れるステップ数を決定している
steps_per_epoch=tf.math.ceil(len(all_image_paths)/BATCH_SIZE).numpy()
default_timeit_steps = 2*steps_per_epoch+1

def timeit(ds, steps=default_timeit_steps):
  overall_start = time.time()
  # Fetch a single batch to prime the pipeline (fill the shuffle buffer),
  # before starting the timer
  it = iter(ds.take(steps+1))
  next(it)

  start = time.time()
  for i,(images,labels) in enumerate(it):
    if i%10 == 0:
      print('.',end='')
  print()
  end = time.time()

  duration = end-start
  print("{} batches: {} s".format(steps, duration))
  print("{:0.5f} Images/s".format(BATCH_SIZE*steps/duration))
  print("Total time: {}s".format(end-overall_start))  
# 出力結果
#

性能を上げる工夫をしていないときの読み込み時間を調べておきます。

NoteBook

timeit(ds)

# 出力結果
#
# ........................
# 231.0 batches: 11.89171576499939 s
# 621.60921 Images/s
# Total time: 17.746938943862915s

キャッシュを使って性能を上げる

cacheメソッドを使うとキャッシュを使うことができます。

NoteBook

ds = image_label_ds.cache()
ds = ds.apply(
    tf.data.experimental.shuffle_and_repeat(buffer_size=image_count))
ds = ds.batch(BATCH_SIZE).prefetch(buffer_size=AUTOTUNE)

timeit(ds)

# 出力結果
#
# ........................
# 231.0 batches: 1.0274419784545898 s
# 7194.56685 Images/s
# Total time: 6.322047233581543s

読み込みが早くなっていますね。

次にキャッシュをファイルに保存してみます。ファイルに保存することで、メモリに収まらないデータもキャッシュすることができます。さらに、２回目の読み込みはファイルからキャッシュを取得するので高速になります。

NoteBook

ds = image_label_ds.cache(filename='./cache.tf-data')
ds = ds.apply(
    tf.data.experimental.shuffle_and_repeat(buffer_size=image_count))
ds = ds.batch(BATCH_SIZE).prefetch(1)

timeit(ds)
# 2回目を実行
timeit(ds)

# 出力結果
#
# ........................
# 231.0 batches: 3.2389090061187744 s
# 2282.24998 Images/s
# Total time: 13.806905269622803s
# ........................
# 231.0 batches: 3.145420789718628 s
# 2350.08302 Images/s
# Total time: 4.079197406768799s

2回目はだいぶ早くなっていることがわかります。

TFRecordを使う

TFRecordはTensorFlowがファイルを読み込みやすいように最適化されたフォーマットです。特にリモートストレージのようなところからファイルを読み込むときは、この形式で保存しておくと読み込みが格段に早くなります。

まずは、TFRecordを保存し、それを読み込んでデータセットを構築します。
TFRecordへの保存処理が間に挟まるだけで、それ以外はこれまでとほぼ同じコードです。

NoteBook

# 画像の前処理（画像形式の変換、リサイズ、値を0〜1の範囲にする）
def preprocess_image(image):
  image = tf.image.decode_jpeg(image, channels=3)
  image = tf.image.resize(image, [192, 192])
  image /= 255.0

  return image

# 画像を読み込む
image_ds = tf.data.Dataset.from_tensor_slices(all_image_paths).map(tf.io.read_file)
# TFRecordのWriterを作成
tfrec = tf.data.experimental.TFRecordWriter('images.tfrec')
# 書き込み
tfrec.write(image_ds)

# 画像をTFRecordのデータセットから読み込んで前処理
image_ds = tf.data.TFRecordDataset('images.tfrec').map(preprocess_image)

# (ここから下はこれまでと全く同じ)
# イメージとラベルをペアにする
ds = tf.data.Dataset.zip((image_ds, label_ds))
# シャッフルとリピートを設定する
ds = ds.apply(
    tf.data.experimental.shuffle_and_repeat(buffer_size=image_count))
# バッチにする
ds=ds.batch(BATCH_SIZE).prefetch(AUTOTUNE)

timeit(ds)
  
# 出力結果
#
# ........................
# 231.0 batches: 14.31936502456665 s
# 516.22401 Images/s
# Total time: 20.784823894500732s

この段階ではキャッシュを使っていないので早くはありません。

（最終形）TFRecordを使って性能を上げる

TFRecordとキャッシュを使って高速化します。
この方法がこれまでに紹介した中で一番性能を上げることができます。

方法としては、前処理した画像をシリアライズして、TFRecordのファイルに保存します。
これを読み込んでパースし、キャッシュを使いながらデータセットを読み込むことで飛躍的に速度が向上します。

NoteBook

# データセットを作る
paths_ds = tf.data.Dataset.from_tensor_slices(all_image_paths)
image_ds = paths_ds.map(load_and_preprocess_image)

# シリアライズする。文字列のテンソルになる
ds = image_ds.map(tf.io.serialize_tensor)

# TFRecordの形式で保存する
tfrec = tf.data.experimental.TFRecordWriter('images.tfrec')
tfrec.write(ds)

# TFRecordからデータを読み込む
ds = tf.data.TFRecordDataset('images.tfrec')

# パース用の関数
def parse(x):
  result = tf.io.parse_tensor(x, out_type=tf.float32)
  result = tf.reshape(result, [192, 192, 3])
  return result

# データセットをパースする
# TFRecordにはテンソルの型や形状（shape）が保存されていないので、それらを復元する
ds = ds.map(parse, num_parallel_calls=AUTOTUNE)

# (ここから下はこれまでと全く同じ)
# データセットをラベルとペアにして読み込みを実行する
ds = tf.data.Dataset.zip((ds, label_ds))
# キャシュ化する処理はチュートリアルに書いていなかったが、
# これを入れないと処理速度が出ない。チュートリアルのバグ？
ds = ds.cache(filename='./cache.tf-data')
ds = ds.apply(
  tf.data.experimental.shuffle_and_repeat(buffer_size=image_count))
ds=ds.batch(BATCH_SIZE).prefetch(AUTOTUNE)

timeit(ds)

# 出力結果
#
# ........................
# 231.0 batches: 3.1774957180023193 s
# 2326.36033 Images/s
# Total time: 5.068001985549927s

公式よりも処理が遅くなっていますが、キャッシュにファイルを指定しているためだと思われます。ファイルをしていしなければ公式と同じ速度(2.6秒)が得られました。

4. 最後に

noteでは「モバイルアプリエンジニアのためのTensorFlow 2.x 入門」というのを連載しています。
モバイルアプリエンジニアの方以外でも活用できるように書いていますのでぜひご覧ください。

TwitterではiOSの開発や機械学習についてツイートしています。

https://twitter.com/jugemjugemjugem

最後に、若干宣伝ぽくて恐縮ですが、私はフリーランスエンジニアをしております。
機械学習をiPhoneデバイス上で動作させるといったお仕事もできますので、お気軽にご相談下さい。

連絡先名：TokyoYoshida
連絡先： yoshidaforpublic@gmail.com

TL;DR

画像を読み込んで表示する

1. データセットを構築する

画像をモデルが読める形に変換する

tf.data.Datasetを構築する

ラベルと画像のペアを作る

シャッフル、バッチ化、繰り返し、早い読み込み

2. モデルに読み込ませる

3. 読み込み性能を上げる

キャッシュを使って性能を上げる

TFRecordを使う

（最終形）TFRecordを使って性能を上げる

4. 最後に

Discussion