TFDS を使用して外部 tfrecord をロードする

サードパーティのツールによって生成されたtf.train.Exampleプロト ( .tfrecord 、 .riegeliなど内) があり、それを tfds API で直接ロードしたい場合は、このページが最適です。

.tfrecordファイルをロードするには、次の操作を行うだけです。

  • TFDS の命名規則に従ってください。
  • tfrecord ファイルにメタデータ ファイル ( dataset_info.json 、 features.json ) を追加します。

制限事項:

ファイルの命名規則

TFDS は、ファイル名のテンプレートの定義をサポートしているため、さまざまなファイル命名スキームを柔軟に使用できます。テンプレートはtfds.core.ShardedFileTemplateで表され、変数{DATASET} 、 {SPLIT} 、 {FILEFORMAT} 、 {SHARD_INDEX} 、 {NUM_SHARDS} 、および{SHARD_X_OF_Y}をサポートします。たとえば、TFDS のデフォルトのファイル命名スキームは、 {DATASET}-{SPLIT}.{FILEFORMAT}-{SHARD_X_OF_Y}です。 MNIST の場合、これはファイル名が次のようになることを意味します。

  • mnist-test.tfrecord-00000-of-00001
  • mnist-train.tfrecord-00000-of-00001

メタデータの追加

特徴構造を提供する

TFDS がtf.train.Exampleプロトをデコードできるようにするには、仕様に一致するtfds.features構造体を提供する必要があります。例えば：

features = tfds.features.FeaturesDict({
    'image':
        tfds.features.Image(
            shape=(256, 256, 3),
            doc='Picture taken by smartphone, downscaled.'),
    'label':
        tfds.features.ClassLabel(names=['dog', 'cat']),
    'objects':
        tfds.features.Sequence({
            'camera/K': tfds.features.Tensor(shape=(3,), dtype=tf.float32),
        }),
})

次のtf.train.Example仕様に対応します。

{
    'image': tf.io.FixedLenFeature(shape=(), dtype=tf.string),
    'label': tf.io.FixedLenFeature(shape=(), dtype=tf.int64),
    'objects/camera/K': tf.io.FixedLenSequenceFeature(shape=(3,), dtype=tf.int64),
}

特徴を指定すると、TFDS が画像、ビデオなどを自動的にデコードできるようになります。他の TFDS データセットと同様に、特徴のメタデータ (例: ラベル名など) がユーザーに公開されます (例: info.features['label'].names ）。

生成パイプラインを制御する場合

TFDS の外部でデータセットを生成しても、生成パイプラインを制御する場合は、 tfds.features.FeatureConnector.serialize_exampleを使用して、データをdict[np.ndarray]からtf.train.Example proto bytesにエンコードできます。

with tf.io.TFRecordWriter('path/to/file.tfrecord') as writer:
  for ex in all_exs:
    ex_bytes = features.serialize_example(data)
    writer.write(ex_bytes)

これにより、TFDS との機能互換性が保証されます。

同様に、proto ( example ) をデコードするためにfeature.deserialize_example存在します。

生成パイプラインを制御しない場合

tfds.featuresがtf.train.Exampleでどのように表されるかを確認したい場合は、colab でこれを調べることができます。

  • tfds.features tf.train.Exampleの人間が判読できる構造に変換するには、 features.get_serialized_info()を呼び出すことができます。
  • tf.io.parse_single_exampleに渡された正確なFixedLenFeature ,...仕様を取得するには、 spec = features.tf_example_specを使用できます。

分割に関する統計を取得する

TFDS では、各シャード内のサンプルの正確な数を知る必要があります。これは、 len(ds)やサブプリット API : split='train[75%:]'などの機能に必要です。

  • この情報があれば、 tfds.core.SplitInfoのリストを明示的に作成して、次のセクションに進むことができます。

    split_infos = [
    tfds.core.SplitInfo(
        name='train',
        shard_lengths=[1024, ...],  # Num of examples in shard0, shard1,...
        num_bytes=0,  # Total size of your dataset (if unknown, set to 0)
    ),
    tfds.core.SplitInfo(name='test', ...),
]

  • この情報がわからない場合は、 compute_split_info.pyスクリプトを使用して (またはtfds.folder_dataset.compute_split_infoを使用した独自のスクリプトで) 計算できます。指定されたディレクトリ上のすべてのシャードを読み取り、情報を計算するビーム パイプラインを起動します。

メタデータ ファイルを追加する

データセットに沿って適切なメタデータ ファイルを自動的に追加するには、 tfds.folder_dataset.write_metadataを使用します。

tfds.folder_dataset.write_metadata(
    data_dir='/path/to/my/dataset/1.0.0/',
    features=features,
    # Pass the `out_dir` argument of compute_split_info (see section above)
    # You can also explicitly pass a list of `tfds.core.SplitInfo`.
    split_infos='/path/to/my/dataset/1.0.0/',
    # Pass a custom file name template or use None for the default TFDS
    # file name template.
    filename_template='{SPLIT}-{SHARD_X_OF_Y}.{FILEFORMAT}',

    # Optionally, additional DatasetInfo metadata can be provided
    # See:
    # https://www.tensorflow.org/datasets/api_docs/python/tfds/core/DatasetInfo
    description="""Multi-line description."""
    homepage='http://my-project.org',
    supervised_keys=('image', 'label'),
    citation="""BibTex citation.""",
)

データセット ディレクトリで関数が一度呼び出されると、メタデータ ファイル ( dataset_info.jsonなど) が追加され、TFDS を使用してデータセットをロードする準備が整います (次のセクションを参照)。

TFDS を使用してデータセットをロードする

フォルダーから直接

メタデータが生成されたら、標準の TFDS API ( tfds.builderなど) でtfds.core.DatasetBuilderを返すtfds.builder_from_directoryを使用してデータセットをロードできます。

builder = tfds.builder_from_directory('~/path/to/my_dataset/3.0.0/')

# Metadata are available as usual
builder.info.splits['train'].num_examples

# Construct the tf.data.Dataset pipeline
ds = builder.as_dataset(split='train[75%:]')
for ex in ds:
  ...

複数のフォルダーから直接

複数のフォルダーからデータを読み込むことも可能です。これは、たとえば強化学習において、複数のエージェントがそれぞれ個別のデータセットを生成しており、それらすべてを一緒にロードしたい場合に発生する可能性があります。他の使用例としては、新しいデータセットが定期的に生成される場合 (たとえば、1 日あたり新しいデータセット)、日付範囲からデータをロードする場合があります。

複数のフォルダーからデータをロードするには、 tfds.builder_from_directoriesを使用します。これは、標準の TFDS API ( tfds.builderなど) でtfds.core.DatasetBuilderを返します。

builder = tfds.builder_from_directories(builder_dirs=[
    '~/path/my_dataset/agent1/1.0.0/',
    '~/path/my_dataset/agent2/1.0.0/',
    '~/path/my_dataset/agent3/1.0.0/',
])

# Metadata are available as usual
builder.info.splits['train'].num_examples

# Construct the tf.data.Dataset pipeline
ds = builder.as_dataset(split='train[75%:]')
for ex in ds:
  ...

フォルダー構造 (オプション)

TFDS との互換性を高めるために、データを<data_dir>/<dataset_name>[/<dataset_config>]/<dataset_version>として編成できます。例えば：

data_dir/
    dataset0/
        1.0.0/
        1.0.1/
    dataset1/
        config0/
            2.0.0/
        config1/
            2.0.0/

これにより、 data_dir/指定するだけで、データセットがtfds.load / tfds.builder API と互換性を持つようになります。

ds0 = tfds.load('dataset0', data_dir='data_dir/')
ds1 = tfds.load('dataset1/config0', data_dir='data_dir/')