TFRecordからデータが読み込めない

Question

#環境

• Ubuntu 18.04
• Jupyter Notebook
• Python 3.6.8
• Tensorflow 1.12.0

解決したい問題

TFRecordに学習データを保存したが、それをパースできない

詳細

画像データをTFRecordに保存しようと考え、以下のようなコードを実行しました。

writer = tf.python_io.TFRecordWriter('training.tfrecord')
for image_path, label in zip(X_train, y_train):     # (X_train, y_train) = (画像ファイルのパス, 整数のラベル)
    image = cv2.imread(image_path)
    image = cv2.resize(image, (150, 150)) / 255.0    # 画像サイズを揃えて、正規化してtfrecordに入れる
    ex = tf.train.Example(
        features = tf.train.Features(
            feature={
                'image' : tf.train.Feature(float_list = tf.train.FloatList(value=image.ravel())),
                'label' : tf.train.Feature(int64_list = tf.train.Int64List(value=[label]))
            }
        )
    )
    writer.write(ex.SerializeToString())
writer.close()

これを最も簡単な方法として、次のような方法で取り出したところうまく行きました。

for record in tf.python_io.tf_record_iterator('test.tfrecord'):
    example = tf.train.Example()
    example.ParseFromString(record)
    
    img = example.features.feature['image'].float_list.value
    label = example.features.feature['label'].int64_list.value[0]

最後に取り出されたデータを画像に戻して表示すると

のようにうまく行っています。

しかしながら、これをモデルに入れるためにDataset APIを用いる形にするとうまく行きません。具体的には、

# 関数をいくつか定義しておく
def _parse_function(example_proto):
    features = {
        'label' : tf.FixedLenFeature((), tf.int64),
        'image' : tf.FixedLenFeature((), tf.float32)
    }
    parsed_features = tf.parse_single_example(example_proto, features)
    
    return parsed_features['image'], parsed_features['label']

def read_image(images, labels):
    label = tf.cast(labels, tf.int32)
    images = tf.cast(images, tf.float32)
    image = tf.reshape(images, [150, 150, 3])

# データの読み込み
dataset = tf.data.TFRecordDataset('training.tfrecord')
dataset = dataset.map(_parse_function)
dataset = dataset.map(read_image)    # <- ここでエラー

read_imageを適用する段階でエラーが出ます。

ValueError: Cannot reshape a tensor with 1 elements to shape [150,150,3] (67500 elements) for 'Reshape' (op: 'Reshape') with input shapes: [], [3] and with input tensors computed as partial shapes: input[1] = [150,150,3].

これは画像をベクトルから[150, 150, 3]に変えられないということを言っていると思ったので、_parse_functionを適用直後の段階でdatasetの中身を確認してみると

<MapDataset shapes: ((), ()), types: (tf.float32, tf.int64)>

となっており、中身が全く入っていないことがわかりました。データをシャッフルしたりバッチ化して読み込みたいので、うまく行っている方ではなくこちらの方法を用いたいと考えているのですが、どのようにしたらデータを読み込むことができるでしょうか？

よろしくおねがいいたします。