tf.Module

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

tf.Module(
    name=None
)

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

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

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

d = Dense(input_dim=3, output_size=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_dim=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

module = MLP(input_size=5, sizes=[5, 5])
module.variables
(<tf.Variable 'mlp/b:0' shape=(5,) dtype=float32, numpy=..., dtype=float32)>,
<tf.Variable 'mlp/w:0' shape=(5, 5) dtype=float32, numpy=...,
   dtype=float32)>,
<tf.Variable 'mlp/b:0' shape=(5,) dtype=float32, numpy=..., dtype=float32)>,
<tf.Variable 'mlp/w:0' shape=(5, 5) dtype=float32, numpy=...,
   dtype=float32)>)

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)>