Метод experimental_distribute_dataset
Метод experimental_distribute_dataset применяется к экземпляру стратегии
MirroredStrategy для распределения данных между доступными устройствами.
Первым параметром метод принимает датасет TensorFlow, который необходимо
распределить. Метод возвращает объект DistributedDataset, итерация
по которому автоматически разбивает данные между репликами модели.
Это позволяет эффективно использовать несколько GPU при обучении.
Синтаксис
strategy.experimental_distribute_dataset(dataset)
Пример
Давайте создадим стратегию MirroredStrategy и распределим
простой датасет между устройствами:
import tensorflow as tf
strategy = tf.distribute.MirroredStrategy()
dataset = tf.data.Dataset.from_tensor_slices([1, 2, 3, 4, 5, 6])
dataset = dataset.batch(2)
dist_dataset = strategy.experimental_distribute_dataset(dataset)
for batch in dist_dataset:
print(batch)
Результат выполнения кода:
tf.Tensor([1 2], shape=(2,), dtype=int32)
tf.Tensor([3 4], shape=(2,), dtype=int32)
tf.Tensor([5 6], shape=(2,), dtype=int32)
Пример
Давайте применим метод к датасету с изображениями и метками, а затем выполним шаг обучения:
<+python+>
import tensorflow as tf
strategy = tf.distribute.MirroredStrategy()
features = tf.constant([[1, 2, 3], [4, 5, 6], [7, 8, 9], [10, 11, 12]])
labels = tf.constant([0, 1, 0, 1])
dataset = tf.data.Dataset.from_tensor_slices((features, labels))
dataset = dataset.batch(2)
dist_dataset = strategy.experimental_distribute_dataset(dataset)
for x, y in dist_dataset:
print(x)
print(y)
<-python+>
Результат выполнения кода:
tf.Tensor(
[[1 2 3]
[4 5 6]], shape=(2, 3), dtype=int32)
tf.Tensor([0 1], shape=(2,), dtype=int32)
tf.Tensor(
[[ 7 8 9]
[10 11 12]], shape=(2, 3), dtype=int32)
tf.Tensor([0 1], shape=(2,), dtype=int32)
Смотрите также
-
класс
MirroredStrategy,
который реализует синхронное распределенное обучение -
метод
distribute_datasets_from_function,
который распределяет датасет, созданный функцией -
метод
run,
который выполняет функцию на каждой реплике -
метод
scope,
который открывает область видимости стратегии
