tf.Module

Базовый класс модуля нейронной сети.

Наследуется от: Checkpointable

tf.Module(
    name=None
)

Модуль — это именованный контейнер для tf.Variable, других tf.Module и функций, которые применяются к входным данным пользователя. Например, слой плотности в нейронной сети может быть реализован как tf.Module:

class Dense(tf.Module):
  def __init__(self, in_features, output_features, name=None):
    super(Dense, self).__init__(name=name)
    self.w = tf.Variable(
        tf.random.normal([input_features, output_features]), name='w')
    self.b = tf.Variable(tf.zeros([output_features]), name='b')

  def __call__(self, x):
    y = tf.matmul(x, self.w) + self.b
    return tf.nn.relu(y)

Вы можете использовать слой Dense, как ожидается:

d = Dense(input_features=64, output_features=10)
d(tf.ones([100, 64]))
#==> <tf.Tensor: ...>

Подклассированием tf.Module вместо object любые tf.Variable или tf.Module экземпляры, назначенные свойствам объекта, можно собрать с помощью свойства variables, trainable_variables или submodules:

d.variables
#==> (<tf.Variable 'b:0' ...>, <tf.Variable 'w:0' ...>)

Подклассы tf.Module также могут использовать метод _flatten, который можно использовать для отслеживания любых других типов.

У всех классов tf.Module есть связанный tf.name_scope, который можно использовать для группировки операций в TensorBoard и создания иерархий для имён переменных, что может помочь в отладке. Мы рекомендуем использовать область имён при создании вложенных подмодулей/параметров или для методов вперёд, график которых вы хотите просмотреть в TensorBoard. Вы можете явно войти в область имён с помощью with self.name_scope: или вы можете аннотировать методы (кроме __init__) с помощью @tf.Module.with_name_scope.

class MLP(tf.Module):
  def __init__(self, input_size, sizes, name=None):
    super(MLP, self).__init__(name=name)
    self.layers = []
    with self.name_scope:
      for size in sizes:
        self.layers.append(Dense(input_size=input_size, output_size=size))
        input_size = size

  @tf.Module.with_name_scope
  def __call__(self, x):
    for layer in self.layers:
      x = layer(x)
    return x

a = tf.Module()
b = tf.Module()
c = tf.Module()
a.b = b
b.c = c
assert list(a.submodules) == [b, c]
assert list(b.submodules) == [c]
assert list(c.submodules) == []

Методы

with_name_scope

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

@classmethod
with_name_scope(
    method
)

Декоратор для автоматического входа в область имён модуля.

class MyModule(tf.Module):
  @tf.Module.with_name_scope
  def __call__(self, x):
    if not hasattr(self, 'w'):
      self.w = tf.Variable(tf.random.normal([x.shape[1], 64]))
    return tf.matmul(x, self.w)

Использование указанного выше модуля создаст tf.Variable и tf.Tensor, имена которых включают имя модуля:

mod = MyModule()
mod(tf.ones([8, 32]))
# ==> <tf.Tensor: ...>
mod.w
# ==> <tf.Variable ...'my_module/w:0'>