tf.keras.dtensor.experimental.layout_map_scope

Применяет макет к всем tf.Переменным, созданным в области.

@contextlib.contextmanager
tf.keras.dtensor.experimental.layout_map_scope(
    layout_map
)

Создаёт область, в которой все tf.Переменные, созданные в этой области, будут лениво инициализированы и инициализированы позже соответствующим макетом, когда путь к объекту в модели будет стабильным/окончательным.

Обратите внимание, что отображение макета будет использовать имена объекта/атрибута в качестве ключа для сопоставления переменной с макетом.

Для моделей, созданных с помощью наследования, используется полное имя объекта/атрибута в качестве ключа. Для моделей Functional/Sequential, так как слоям внутри модели не присваивается осмысленное имя атрибута, мы используем layer.name в качестве ключа для слоя, за которым следует имя атрибута. Keras гарантирует уникальность имен среди слоёв во всех моделях Functional/Sequential.

Ниже приведены примеры, демонстрирующие имена объектов переменных для различных типов моделей Keras:

layout_map = layout_map_lib.LayoutMap(mesh=self.mesh)
layout_map['d1.kernel'] = layout_1
layout_map['d1.bias'] = layout_2
layout_map['d2.kernel'] = layout_3
layout_map['d2.bias'] = layout_4

## Subclassed model
class SubclassModel(tf.keras.Model):

  def __init__(self, name=None):
    super().__init__(name=name)
    self.d1 = tf.keras.layers.Dense(1000)
    self.d2 = tf.keras.layers.Dense(1000)

  def call(self, inputs):
    x = self.d1(inputs)
    return self.d2(x)

with layout_map_scope(layout_map):
  model = SubclassModel()
# Triggering the creation of weights within or outside of the scope works
inputs = tf.zeros((10, 10))
results = model(inputs)

model.d1.kernel.layout == layout_1
model.d1.bias.layout == layout_2
model.d2.kernel.layout == layout_3
model.d2.bias.layout == layout_4

## Functional model
with layout_map_scope(layout_map):
  inputs = tf.keras.Input((10,), batch_size=10)
  x = tf.keras.layers.Dense(20, name='d1')(inputs)
  output = tf.keras.layers.Dense(30, name='d2')(x)

  model = tf.keras.Model(inputs, output)

d1 = model.layers[1]
d2 = model.layers[2]

d1.kernel.layout == layout_1
d1.bias.layout == layout_2
d1.kernel.layout == layout_3
d1.bias.layout == layout_4

## Sequential model
with layout_map_scope(layout_map):
  model = tf.keras.Sequential([
      tf.keras.layers.Dense(20, name='d1', input_shape=(10,)),
      tf.keras.layers.Dense(30, name='d2')
  ])

d1 = model.layers[0]
d2 = model.layers[1]

d1.kernel.layout == layout_1
d1.bias.layout == layout_2
d1.kernel.layout == layout_3
d1.bias.layout == layout_4