tf.keras.layers.FlaxLayer

Слой Keras, оборачивающий модуль Flax.

Наследует от: JaxLayer, Layer, Operation

tf.keras.layers.FlaxLayer(
    module, method=None, variables=None, **kwargs
)

Этот слой позволяет использовать компоненты Flax в виде экземпляров flax.linen.Module в Keras при использовании JAX в качестве бэкенда Keras.

Метод модуля, используемый для прямого прохода, можно указать с помощью аргумента method; по умолчанию используется __call__. Этот метод должен принимать следующие аргументы с указанными именами:

• self, если метод привязан к модулю, как в случае значения по умолчанию __call__, и module в противном случае, чтобы передать модуль.
• inputs: входные данные модели — массив JAX или PyTree массивов.
• training (необязательно): аргумент, указывающий, работает ли модель в режиме обучения или вывода; в режиме обучения передаётся True.

FlaxLayer автоматически обрабатывает нетренируемое состояние модели и необходимые RNG. Обратите внимание, что параметру mutable метода flax.linen.Module.apply() присваивается значение DenyList(["params"]). Таким образом предполагается, что все переменные вне коллекции "params" являются нетренируемыми весами.

В этом примере показано, как создать FlaxLayer на основе Module Flax, используя метод __call__ по умолчанию и не указывая аргумент обучения:

class MyFlaxModule(flax.linen.Module):
    @flax.linen.compact
    def __call__(self, inputs):
        x = inputs
        x = flax.linen.Conv(features=32, kernel_size=(3, 3))(x)
        x = flax.linen.relu(x)
        x = flax.linen.avg_pool(x, window_shape=(2, 2), strides=(2, 2))
        x = x.reshape((x.shape[0], -1))  # flatten
        x = flax.linen.Dense(features=200)(x)
        x = flax.linen.relu(x)
        x = flax.linen.Dense(features=10)(x)
        x = flax.linen.softmax(x)
        return x

flax_module = MyFlaxModule()
keras_layer = FlaxLayer(flax_module)

В этом примере показано, как обернуть метод модуля, чтобы он соответствовал требуемой сигнатуре. Это позволяет использовать несколько входных аргументов и аргумент обучения с другим именем и значениями. Кроме того, здесь показано, как использовать функцию, не привязанную к модулю.

class MyFlaxModule(flax.linen.Module):
    @flax.linen.compact
    def forward(self, input1, input2, deterministic):
        ...
        return outputs

def my_flax_module_wrapper(module, inputs, training):
    input1, input2 = inputs
    return module.forward(input1, input2, not training)

flax_module = MyFlaxModule()
keras_layer = FlaxLayer(
    module=flax_module,
    method=my_flax_module_wrapper,
)

Методы

from_config

Посмотреть исходный код

@classmethod
from_config(
    config
)

Создаёт слой из его конфигурации.

Этот метод является обратным по отношению к get_config и позволяет создать такой же слой из словаря конфигурации. Он не обрабатывает связи слоя (это делает Network) и веса (это делает set_weights).

symbolic_call

Посмотреть исходный код

symbolic_call(
    *args, **kwargs
)