tf.Module

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

tf.Module(
    name=None
)

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

class Dense(tf.Module):
  def __init__(self, in_features, out_features, name=None):
    super(Dense, self).__init__(name=name)
    self.w = tf.Variable(
      tf.random.normal([in_features, out_features]), name='w')
    self.b = tf.Variable(tf.zeros([out_features]), name='b')
  def __call__(self, x):
    y = tf.matmul(x, self.w) + self.b
    return tf.nn.relu(y)
 

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

d = Dense(in_features=3, out_features=2)
d(tf.ones([1, 3]))
<tf.Tensor: shape=(1, 2), dtype=float32, numpy=..., dtype=float32)>

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

d.variables
    (<tf.Variable 'b:0' shape=(2,) dtype=float32, numpy=...,
    dtype=float32)>,
    <tf.Variable 'w:0' shape=(3, 2) dtype=float32, numpy=..., dtype=float32)>)

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

Все классы tf.Module имеют связанный tf.name_scope, который можно использовать для группировки операций в TensorBoard и создания иерархии имен переменных, что может помочь в отладке. Мы рекомендуем использовать область имен при создании вложенных подмодулей/параметров или для методов forward, график которых вы можете просмотреть в 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
list(a.submodules) == [b, c]
True
list(b.submodules) == [c]
True
list(c.submodules) == []
True

Методы

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], 3]))
    return tf.matmul(x, self.w)

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

mod = MyModule()
mod(tf.ones([1, 2]))
<tf.Tensor: shape=(1, 3), dtype=float32, numpy=..., dtype=float32)>
mod.w
<tf.Variable 'my_module/Variable:0' shape=(2, 3) dtype=float32,
numpy=..., dtype=float32)>