tf.data.Datasetをunzip（メソッドzipの逆操作）する方法

結論

例えば、

import tensorflow as tf

a = tf.data.Dataset.range(0, 5)
b = tf.data.Dataset.range(5, 10)

ds = tf.data.Dataset.zip((a, b))

print(ds.element_spec)

出力

(TensorSpec(shape=(), dtype=tf.int64, name=None), TensorSpec(shape=(), dtype=tf.int64, name=None))

このようにzipされたDatasetに対して、zipの逆操作をしたい場合、

c = ds.map(lambda x, y: x)
d = ds.map(lambda x, y: y)

print(c.element_spec)
print(d.element_spec)

出力

TensorSpec(shape=(), dtype=tf.int64, name=None)
TensorSpec(shape=(), dtype=tf.int64, name=None)

このように、片方の要素を取り出すLambda式をmapすれば取り出すことができます。

環境

Python: 3.10.8
TensorFlow: 2.12.0

背景

TensorFlow公式の時系列予測チュートリアルの「 tf.data.Dataset を作成する」の節について、

def make_dataset(self, data):
  data = np.array(data, dtype=np.float32)
  ds = tf.keras.utils.timeseries_dataset_from_array(
      data=data,
      targets=None,
      sequence_length=self.total_window_size,
      sequence_stride=1,
      shuffle=True,
      batch_size=32,)

  ds = ds.map(self.split_window)

  return ds

WindowGenerator.make_dataset = make_dataset

ds.map(self.split_window)で返ってくるDatasetをunzipしたくなったのが発端です。

このチュートリアルを参考にしながら独自のモデルを作成しているのですが、もう一つ別で用意したDatasetを入力に加えたく、一度unzipしてからzipし直すに至りました。

p.s.

「もっとスマートにunzipする方法がある。」
「そもそもそんな手順踏まなくてもkerasのInput層にデータを流せる。」
等々、もっと賢い方法がある場合はご教示いただけるとありがたいです。