Как интегрировать tf.data.dataset с rayTune для распределенного обучения

Question

Использование tensorflow-cpu==2.9.3, petastorm==0.12.1 на Python 3.7

Я создал tf.data.Dataset, используя petastorm для набора данных обучения и проверки.

• ds_train (DatasetV1Adapter; думаю, это старая версия tf.data.dataset)
• ds_valid (адаптер набора данныхV1)

Первая пробная версия: следуйте документу rayTune https://docs.ray.io/en/latest/tune/faq.html#how-can-i-use-large-datasets-in-tune .

      tuner = tune.Tuner(
    tune.with_resources(tune.with_parameters(
       train_model, ds_train=ds_train, ds_valid=ds_valid),
    resources={'cpu':1})
)

def train_model(config, ds_train, ds_valid):
    model = <simple deep learning model>
    history = model.fit(x=ds_train, validation_data=ds_valid)
    return history

Это выводитraytune tensorflow.python.framework.errors_impl.InvalidArgumentERror: Cannot convert a Tensor of dtype variant to a Numpy array

Второе испытание: создание tf.data.dataset в функции train_model работает, однако оно потребляет больше памяти, поскольку каждому рабочему rayTune необходимо генерировать tf.data.dataset, что приводит к ошибке OOM.

Третье испытание: сохранено ds_train, ds_valid с использованиемtf.data.experimentai.save(ds_train, path)затем в train_model каждый работник rayTune мог просто загрузить tf.data.dataset, а затем использовать его черезtf.data.experimental.load(path).

Это лишает возможностиpetastormкоторый использует AWS S3 в качестве исходных данных для tf.data.Dataset, поэтому каждый раз, когда размер tf.data.dataset увеличивается, размер локального диска также должен увеличиваться.

Каковы наилучшие методы проведения распределенного обучения с использованием RayTune в tf.data.Dataset?