【Tensorflow2】データをTFRecord形式に書き込む方法

この記事を読んで分かること

TFRecordファイルの作成方法

今回は、機械学習で扱いやすいように、データをシリアライズしてTFRecordファイルを作成する方法を解説していきます。

TFRecordを作る
TFRecordの応用
まとめ
参考文献

TFRecordを作る

TFRecordを作るには、「データの整形」→「tf.train.Exampleに変換」→「TFRecordファイルに書き込み」の順に行っていきます。

featureなど、TFRecordの構造についてはこちらの記事で解説しています。

ここでは、画像データを例にTFRecordファイルを作成していきます。

データの整形

tf.train.Exampleには1次元配列しか保存できないので、画像データをバイナリ化していきます。

画像はネコと犬の写真1枚ずつの計2枚です。

open関数で画像を開けば、バイト列が取得できます。

import tensorflow as tf

cat_image_string = open("cat.jpg", 'rb').read()
cat_image_string

実行結果(クリックして表示)

b'\xff\xd8\xff\xe0\x00\x10JFIF\x00\x01\x01\x00\x00\x01\x00\x01\x00\x00\xff\xdb\x00C\x00\x03\x02\x02\x02\x02\x02\x03\x02\x02\x02\・・・

ラベルは(ネコ:0、犬:1)として設定します。

cat_label = 0

画像のサイズとチャネル数も変数に格納しておきます。

cat_image_shape = tf.io.decode_jpeg(cat_image_string).shape
cat_image_shape

実行結果(クリックして表示)

TensorShape([1300, 1920, 3])

犬の画像でも同様の処理をして、データの整形は完了です。

dog_image_string = open("dog.jpg", 'rb').read()
dog_label = 1
dog_image_shape = tf.io.decode_jpeg(dog_image_string).shape

tf.train.Exampleに変換

次は、整形したデータをtf.train.Exampleオブジェクトの形式にします。

画像サイズやラベルは「tf.train.Int64List」に格納し、バイナリ化した画像は「tf.train.BytesList」として保存します。

cat_feature = {
    'height': tf.train.Feature(int64_list=tf.train.Int64List(value=[cat_image_shape[0]])),
    'width': tf.train.Feature(int64_list=tf.train.Int64List(value=[cat_image_shape[1]])),
    'depth': tf.train.Feature(int64_list=tf.train.Int64List(value=[cat_image_shape[2]])),
    'label': tf.train.Feature(int64_list=tf.train.Int64List(value=[cat_label])),
    'image_raw': tf.train.Feature(bytes_list=tf.train.BytesList(value=[cat_image_string])),
}

cat_example = tf.train.Example(features=tf.train.Features(feature=cat_feature))
cat_example

実行結果(クリックして表示)

features {
  feature {
    key: "depth"
    value {
      int64_list {
        value: 3
      }
    }
  }
  feature {
    key: "height"
    value {
      int64_list {
        value: 1300
      }
    }
  }
  feature {
    key: "image_raw"
    value {
      bytes_list {
        value: "\377\330\377\340\000・・・
      }
    }
  }
  feature {
    key: "label"
    value {
      int64_list {
        value: 0
      }
    }
  }
  feature {
    key: "width"
    value {
      int64_list {
        value: 1920
      }
    }
  }
}

同様に犬のデータもtf.train.Exampleオブジェクトの形式にします。

dog_feature = {
    'height': tf.train.Feature(int64_list=tf.train.Int64List(value=[dog_image_shape[0]])),
    'width': tf.train.Feature(int64_list=tf.train.Int64List(value=[dog_image_shape[1]])),
    'depth': tf.train.Feature(int64_list=tf.train.Int64List(value=[dog_image_shape[2]])),
    'label': tf.train.Feature(int64_list=tf.train.Int64List(value=[dog_label])),
    'image_raw': tf.train.Feature(bytes_list=tf.train.BytesList(value=[dog_image_string])),
}

dog_example = tf.train.Example(features=tf.train.Features(feature=dog_feature))

TFRecordファイルに書き込み

tf.train.ExampleオブジェクトをTFRecordファイルに書き込んでいきます。

record_file = 'cat_dog.tfrecords'
with tf.io.TFRecordWriter(record_file) as writer:
    writer.write(cat_example.SerializeToString())
    writer.write(dog_example.SerializeToString())

cat_dog.tfrecordsファイルが作成され、画像データをTFRecordにすることができました。

作成したTFRecordファイルの読み込み方法はこちらで解説しています。