Введение в модули, слои и модели

Посмотреть на TensorFlow.org Запустить в Google Colab Посмотреть исходный код на GitHub Скачать блокнот

Для машинного обучения в TensorFlow вам, вероятно, потребуется определить, сохранить и восстановить модель.

Модель, абстрактно:

  • Функция, которая вычисляет что-то на тензорах ( прямой проход )
  • Некоторые переменные, которые могут быть обновлены в ответ на обучение

В этом руководстве вы углубитесь в Keras, чтобы увидеть, как определяются модели TensorFlow. Здесь рассматривается, как TensorFlow собирает переменные и модели, а также как они сохраняются и восстанавливаются.

Настраивать

import tensorflow as tf
from datetime import datetime

%load_ext tensorboard

Определение моделей и слоев в TensorFlow

Большинство моделей состоят из слоев. Слои — это функции с известной математической структурой, которые можно использовать повторно и которые имеют обучаемые переменные. В TensorFlow большинство высокоуровневых реализаций слоев и моделей, таких как Keras или Sonnet , построены на одном и том же базовом классе: tf.Module .

Вот пример очень простого tf.Module , работающего со скалярным тензором:

class SimpleModule(tf.Module):
  def __init__(self, name=None):
    super().__init__(name=name)
    self.a_variable = tf.Variable(5.0, name="train_me")
    self.non_trainable_variable = tf.Variable(5.0, trainable=False, name="do_not_train_me")
  def __call__(self, x):
    return self.a_variable * x + self.non_trainable_variable

simple_module = SimpleModule(name="simple")

simple_module(tf.constant(5.0))
<tf.Tensor: shape=(), dtype=float32, numpy=30.0>

Модули и, в более широком смысле, слои — это терминология глубокого изучения для «объектов»: у них есть внутреннее состояние и методы, которые используют это состояние.

В __call__ нет ничего особенного, кроме того, что он действует как callable Python ; вы можете вызывать свои модели с любыми функциями, какие пожелаете.

Вы можете включать и выключать обучаемость переменных по любой причине, в том числе замораживать слои и переменные во время тонкой настройки.

Путем создания подкласса tf.Module автоматически собираются любые экземпляры tf.Variable или tf.Module , назначенные свойствам этого объекта. Это позволяет сохранять и загружать переменные, а также создавать коллекции tf.Module s.

# All trainable variables
print("trainable variables:", simple_module.trainable_variables)
# Every variable
print("all variables:", simple_module.variables)
trainable variables: (<tf.Variable 'train_me:0' shape=() dtype=float32, numpy=5.0>,)
all variables: (<tf.Variable 'train_me:0' shape=() dtype=float32, numpy=5.0>, <tf.Variable 'do_not_train_me:0' shape=() dtype=float32, numpy=5.0>)
2021-10-26 01:29:45.284549: W tensorflow/python/util/util.cc:348] Sets are not currently considered sequences, but this may change in the future, so consider avoiding using them.

Это пример двухслойной линейно-слойной модели, состоящей из модулей.

Сначала плотный (линейный) слой:

class Dense(tf.Module):
  def __init__(self, in_features, out_features, name=None):
    super().__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)

А затем полная модель, которая создает два экземпляра слоя и применяет их:

class SequentialModule(tf.Module):
  def __init__(self, name=None):
    super().__init__(name=name)

    self.dense_1 = Dense(in_features=3, out_features=3)
    self.dense_2 = Dense(in_features=3, out_features=2)

  def __call__(self, x):
    x = self.dense_1(x)
    return self.dense_2(x)

# You have made a model!
my_model = SequentialModule(name="the_model")

# Call it, with random results
print("Model results:", my_model(tf.constant([[2.0, 2.0, 2.0]])))
Model results: tf.Tensor([[7.706234  3.0919805]], shape=(1, 2), dtype=float32)

tf.Module будут автоматически рекурсивно собирать любые назначенные ему экземпляры tf.Variable или tf.Module . Это позволяет вам управлять коллекциями tf.Module с одним экземпляром модели, а также сохранять и загружать целые модели.

print("Submodules:", my_model.submodules)
Submodules: (<__main__.Dense object at 0x7f7ab2391290>, <__main__.Dense object at 0x7f7b6869ea10>)
for var in my_model.variables:
  print(var, "\n")
<tf.Variable 'b:0' shape=(3,) dtype=float32, numpy=array([0., 0., 0.], dtype=float32)> 

<tf.Variable 'w:0' shape=(3, 3) dtype=float32, numpy=
array([[ 0.05711935,  0.22440144,  0.6370985 ],
       [ 0.3136791 , -1.7006774 ,  0.7256515 ],
       [ 0.16120772, -0.8412193 ,  0.5250952 ]], dtype=float32)> 

<tf.Variable 'b:0' shape=(2,) dtype=float32, numpy=array([0., 0.], dtype=float32)> 

<tf.Variable 'w:0' shape=(3, 2) dtype=float32, numpy=
array([[-0.5353216 ,  1.2815404 ],
       [ 0.62764466,  0.47087234],
       [ 2.19187   ,  0.45777202]], dtype=float32)>

Ожидание создания переменных

Возможно, вы заметили здесь, что вы должны определить как входные, так и выходные размеры слоя. Это значит, что переменная w имеет известную форму и может быть размещена.

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

class FlexibleDenseModule(tf.Module):
  # Note: No need for `in_features`
  def __init__(self, out_features, name=None):
    super().__init__(name=name)
    self.is_built = False
    self.out_features = out_features

  def __call__(self, x):
    # Create variables on first call.
    if not self.is_built:
      self.w = tf.Variable(
        tf.random.normal([x.shape[-1], self.out_features]), name='w')
      self.b = tf.Variable(tf.zeros([self.out_features]), name='b')
      self.is_built = True

    y = tf.matmul(x, self.w) + self.b
    return tf.nn.relu(y)
# Used in a module
class MySequentialModule(tf.Module):
  def __init__(self, name=None):
    super().__init__(name=name)

    self.dense_1 = FlexibleDenseModule(out_features=3)
    self.dense_2 = FlexibleDenseModule(out_features=2)

  def __call__(self, x):
    x = self.dense_1(x)
    return self.dense_2(x)

my_model = MySequentialModule(name="the_model")
print("Model results:", my_model(tf.constant([[2.0, 2.0, 2.0]])))
Model results: tf.Tensor([[4.0598335 0.       ]], shape=(1, 2), dtype=float32)

Эта гибкость является причиной того, что слоям TensorFlow часто нужно указывать только форму своих выходных данных, например, в tf.keras.layers.Dense , а не размер входных и выходных данных.

Сохранение весов

Вы можете сохранить tf.Module и как контрольную точку , и как SavedModel .

Контрольные точки — это просто веса (то есть значения набора переменных внутри модуля и его подмодулей):

chkp_path = "my_checkpoint"
checkpoint = tf.train.Checkpoint(model=my_model)
checkpoint.write(chkp_path)
'my_checkpoint'

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

ls my_checkpoint*
my_checkpoint.data-00000-of-00001  my_checkpoint.index

Вы можете заглянуть внутрь контрольной точки, чтобы убедиться, что вся коллекция переменных сохранена, отсортирована по объекту Python, который их содержит.

tf.train.list_variables(chkp_path)
[('_CHECKPOINTABLE_OBJECT_GRAPH', []),
 ('model/dense_1/b/.ATTRIBUTES/VARIABLE_VALUE', [3]),
 ('model/dense_1/w/.ATTRIBUTES/VARIABLE_VALUE', [3, 3]),
 ('model/dense_2/b/.ATTRIBUTES/VARIABLE_VALUE', [2]),
 ('model/dense_2/w/.ATTRIBUTES/VARIABLE_VALUE', [3, 2])]

При распределенном (многомашинном) обучении их можно сегментировать, поэтому они нумеруются (например, «00000-из-00001»). В этом случае, однако, есть только один осколок.

Когда вы загружаете модели обратно, вы перезаписываете значения в своем объекте Python.

new_model = MySequentialModule()
new_checkpoint = tf.train.Checkpoint(model=new_model)
new_checkpoint.restore("my_checkpoint")

# Should be the same result as above
new_model(tf.constant([[2.0, 2.0, 2.0]]))
<tf.Tensor: shape=(1, 2), dtype=float32, numpy=array([[4.0598335, 0.       ]], dtype=float32)>

Сохранение функций

TensorFlow может запускать модели без исходных объектов Python, как продемонстрировали TensorFlow Serving и TensorFlow Lite , даже если вы загружаете обученную модель из TensorFlow Hub .

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

Этот граф содержит операции, или ops , которые реализуют функцию.

Вы можете определить граф в приведенной выше модели, добавив декоратор @tf.function , чтобы указать, что этот код должен выполняться как граф.

class MySequentialModule(tf.Module):
  def __init__(self, name=None):
    super().__init__(name=name)

    self.dense_1 = Dense(in_features=3, out_features=3)
    self.dense_2 = Dense(in_features=3, out_features=2)

  @tf.function
  def __call__(self, x):
    x = self.dense_1(x)
    return self.dense_2(x)

# You have made a model with a graph!
my_model = MySequentialModule(name="the_model")

Сделанный вами модуль работает точно так же, как и раньше. Каждая уникальная подпись, переданная в функцию, создает отдельный граф. Для получения подробной информации см. руководство «Введение в графики и функции» .

print(my_model([[2.0, 2.0, 2.0]]))
print(my_model([[[2.0, 2.0, 2.0], [2.0, 2.0, 2.0]]]))
tf.Tensor([[0.62891716 0.        ]], shape=(1, 2), dtype=float32)
tf.Tensor(
[[[0.62891716 0.        ]
  [0.62891716 0.        ]]], shape=(1, 2, 2), dtype=float32)

Вы можете визуализировать график, отследив его в сводке TensorBoard.

# Set up logging.
stamp = datetime.now().strftime("%Y%m%d-%H%M%S")
logdir = "logs/func/%s" % stamp
writer = tf.summary.create_file_writer(logdir)

# Create a new model to get a fresh trace
# Otherwise the summary will not see the graph.
new_model = MySequentialModule()

# Bracket the function call with
# tf.summary.trace_on() and tf.summary.trace_export().
tf.summary.trace_on(graph=True)
tf.profiler.experimental.start(logdir)
# Call only one tf.function when tracing.
z = print(new_model(tf.constant([[2.0, 2.0, 2.0]])))
with writer.as_default():
  tf.summary.trace_export(
      name="my_func_trace",
      step=0,
      profiler_outdir=logdir)
tf.Tensor([[0.         0.01750386]], shape=(1, 2), dtype=float32)

Запустите TensorBoard, чтобы просмотреть полученную трассировку:

#docs_infra: no_execute
%tensorboard --logdir logs/func

Скриншот графика в TensorBoard

Создание SavedModel

Рекомендуемый способ обмена полностью обученными моделями — использовать SavedModel . SavedModel содержит как набор функций, так и набор весов.

Вы можете сохранить только что обученную модель следующим образом:

tf.saved_model.save(my_model, "the_saved_model")
INFO:tensorflow:Assets written to: the_saved_model/assets
# Inspect the SavedModel in the directory
ls -l the_saved_model
total 24
drwxr-sr-x 2 kbuilder kokoro  4096 Oct 26 01:29 assets
-rw-rw-r-- 1 kbuilder kokoro 14702 Oct 26 01:29 saved_model.pb
drwxr-sr-x 2 kbuilder kokoro  4096 Oct 26 01:29 variables
# The variables/ directory contains a checkpoint of the variables
ls -l the_saved_model/variables
total 8
-rw-rw-r-- 1 kbuilder kokoro 408 Oct 26 01:29 variables.data-00000-of-00001
-rw-rw-r-- 1 kbuilder kokoro 356 Oct 26 01:29 variables.index

Файл saved_model.pb — это буфер протокола, описывающий функционал tf.Graph .

Модели и слои могут быть загружены из этого представления без создания экземпляра создавшего его класса. Это желательно в ситуациях, когда у вас нет (или вы не хотите) интерпретатора Python, например, при работе в масштабе или на пограничном устройстве, или в ситуациях, когда исходный код Python недоступен или практичен для использования.

Вы можете загрузить модель как новый объект:

new_model = tf.saved_model.load("the_saved_model")

new_model , созданный из загрузки сохраненной модели, является внутренним объектом пользователя TensorFlow без каких-либо знаний о классе. Это не тип SequentialModule .

isinstance(new_model, SequentialModule)
False

Эта новая модель работает с уже определенными входными подписями. Вы не можете добавить больше подписей к модели, восстановленной таким образом.

print(my_model([[2.0, 2.0, 2.0]]))
print(my_model([[[2.0, 2.0, 2.0], [2.0, 2.0, 2.0]]]))
tf.Tensor([[0.62891716 0.        ]], shape=(1, 2), dtype=float32)
tf.Tensor(
[[[0.62891716 0.        ]
  [0.62891716 0.        ]]], shape=(1, 2, 2), dtype=float32)

Таким образом, используя SavedModel , вы можете сохранять веса и графики TensorFlow с помощью tf.Module , а затем снова загружать их.

Модели и слои Keras

Обратите внимание, что до этого момента Керас не упоминался. Вы можете создать свой собственный высокоуровневый API поверх tf.Module , и люди это сделали.

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

Слои Кераса

tf.keras.layers.Layer — это базовый класс всех слоев Keras, он наследуется от tf.Module .

Вы можете преобразовать модуль в слой Keras, просто заменив родителя, а затем изменив __call__ на call :

class MyDense(tf.keras.layers.Layer):
  # Adding **kwargs to support base Keras layer arguments
  def __init__(self, in_features, out_features, **kwargs):
    super().__init__(**kwargs)

    # This will soon move to the build step; see below
    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)

simple_layer = MyDense(name="simple", in_features=3, out_features=3)

Слои Keras имеют свой собственный __call__ , который выполняет некоторые бухгалтерские операции, описанные в следующем разделе, а затем вызывает call() . Вы не должны заметить никаких изменений в функциональности.

simple_layer([[2.0, 2.0, 2.0]])
<tf.Tensor: shape=(1, 3), dtype=float32, numpy=array([[0.      , 0.179402, 0.      ]], dtype=float32)>

Шаг build

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

Слои Keras поставляются с дополнительным этапом жизненного цикла, который дает вам больше гибкости в том, как вы определяете свои слои. Это определяется в функции build .

build вызывается ровно один раз и вызывается с формой ввода. Обычно он используется для создания переменных (весов).

Вы можете переписать слой MyDense выше, чтобы он был гибким в зависимости от размера его входных данных:

class FlexibleDense(tf.keras.layers.Layer):
  # Note the added `**kwargs`, as Keras supports many arguments
  def __init__(self, out_features, **kwargs):
    super().__init__(**kwargs)
    self.out_features = out_features

  def build(self, input_shape):  # Create the state of the layer (weights)
    self.w = tf.Variable(
      tf.random.normal([input_shape[-1], self.out_features]), name='w')
    self.b = tf.Variable(tf.zeros([self.out_features]), name='b')

  def call(self, inputs):  # Defines the computation from inputs to outputs
    return tf.matmul(inputs, self.w) + self.b

# Create the instance of the layer
flexible_dense = FlexibleDense(out_features=3)

На данный момент модель еще не построена, поэтому переменных нет:

flexible_dense.variables
[]

Вызов функции выделяет переменные соответствующего размера:

# Call it, with predictably random results
print("Model results:", flexible_dense(tf.constant([[2.0, 2.0, 2.0], [3.0, 3.0, 3.0]])))
Model results: tf.Tensor(
[[-1.6998017  1.6444504 -1.3103955]
 [-2.5497022  2.4666753 -1.9655929]], shape=(2, 3), dtype=float32)
flexible_dense.variables
[<tf.Variable 'flexible_dense/w:0' shape=(3, 3) dtype=float32, numpy=
 array([[ 1.277462  ,  0.5399406 , -0.301957  ],
        [-1.6277349 ,  0.7374014 , -1.7651852 ],
        [-0.49962795, -0.45511687,  1.4119445 ]], dtype=float32)>,
 <tf.Variable 'flexible_dense/b:0' shape=(3,) dtype=float32, numpy=array([0., 0., 0.], dtype=float32)>]

Поскольку build вызывается только один раз, входные данные будут отклонены, если входная форма несовместима с переменными слоя:

try:
  print("Model results:", flexible_dense(tf.constant([[2.0, 2.0, 2.0, 2.0]])))
except tf.errors.InvalidArgumentError as e:
  print("Failed:", e)
Failed: In[0] mismatch In[1] shape: 4 vs. 3: [1,4] [3,3] 0 0 [Op:MatMul]

Слои Keras имеют гораздо больше дополнительных функций, включая:

  • Дополнительные потери
  • Поддержка метрик
  • Встроенная поддержка дополнительного training аргумента, чтобы различать использование обучения и логического вывода.
  • get_config и from_config , которые позволяют точно хранить конфигурации, чтобы разрешить клонирование модели в Python.

Читайте о них в полном руководстве по пользовательским слоям и моделям.

Керас модели

Вы можете определить свою модель как вложенные слои Keras.

Однако Keras также предоставляет полнофункциональный класс модели под названием tf.keras.Model . Он наследуется от tf.keras.layers.Layer , поэтому модель Keras можно использовать, вкладывать и сохранять так же, как слои Keras. Модели Keras обладают дополнительными функциями, которые упрощают их обучение, оценку, загрузку, сохранение и даже обучение на нескольких машинах.

Вы можете определить SequentialModule сверху с почти идентичным кодом, снова преобразовав __call__ в call() и изменив родителя:

class MySequentialModel(tf.keras.Model):
  def __init__(self, name=None, **kwargs):
    super().__init__(**kwargs)

    self.dense_1 = FlexibleDense(out_features=3)
    self.dense_2 = FlexibleDense(out_features=2)
  def call(self, x):
    x = self.dense_1(x)
    return self.dense_2(x)

# You have made a Keras model!
my_sequential_model = MySequentialModel(name="the_model")

# Call it on a tensor, with random results
print("Model results:", my_sequential_model(tf.constant([[2.0, 2.0, 2.0]])))
Model results: tf.Tensor([[5.5604653 3.3511646]], shape=(1, 2), dtype=float32)

Доступны все те же функции, включая переменные отслеживания и подмодули.

my_sequential_model.variables
[<tf.Variable 'my_sequential_model/flexible_dense_1/w:0' shape=(3, 3) dtype=float32, numpy=
 array([[ 0.05627853, -0.9386015 , -0.77410126],
        [ 0.63149   ,  1.0802224 , -0.37785745],
        [-0.24788402, -1.1076807 , -0.5956209 ]], dtype=float32)>,
 <tf.Variable 'my_sequential_model/flexible_dense_1/b:0' shape=(3,) dtype=float32, numpy=array([0., 0., 0.], dtype=float32)>,
 <tf.Variable 'my_sequential_model/flexible_dense_2/w:0' shape=(3, 2) dtype=float32, numpy=
 array([[-0.93912166,  0.77979285],
        [ 1.4049559 , -1.9380962 ],
        [-2.6039495 ,  0.30885765]], dtype=float32)>,
 <tf.Variable 'my_sequential_model/flexible_dense_2/b:0' shape=(2,) dtype=float32, numpy=array([0., 0.], dtype=float32)>]
my_sequential_model.submodules
(<__main__.FlexibleDense at 0x7f7b48525550>,
 <__main__.FlexibleDense at 0x7f7b48508d10>)

Переопределение tf.keras.Model — это очень Pythonic подход к построению моделей TensorFlow. Если вы переносите модели из других фреймворков, это может быть очень просто.

Если вы строите модели, которые представляют собой простые сборки существующих слоев и входных данных, вы можете сэкономить время и место, используя функциональный API , который поставляется с дополнительными функциями, связанными с реконструкцией модели и архитектурой.

Вот та же модель с функциональным API:

inputs = tf.keras.Input(shape=[3,])

x = FlexibleDense(3)(inputs)
x = FlexibleDense(2)(x)

my_functional_model = tf.keras.Model(inputs=inputs, outputs=x)

my_functional_model.summary()
Model: "model"
_________________________________________________________________
Layer (type)                 Output Shape              Param #   
=================================================================
input_1 (InputLayer)         [(None, 3)]               0         
_________________________________________________________________
flexible_dense_3 (FlexibleDe (None, 3)                 12        
_________________________________________________________________
flexible_dense_4 (FlexibleDe (None, 2)                 8         
=================================================================
Total params: 20
Trainable params: 20
Non-trainable params: 0
_________________________________________________________________
my_functional_model(tf.constant([[2.0, 2.0, 2.0]]))
<tf.Tensor: shape=(1, 2), dtype=float32, numpy=array([[8.219393, 4.511119]], dtype=float32)>

Основное отличие здесь заключается в том, что входная форма указывается заранее как часть функционального процесса построения. Аргумент input_shape в этом случае не обязательно указывать полностью; вы можете оставить некоторые размеры как None .

Сохранение моделей Keras

Модели Keras могут быть отмечены контрольными точками, и это будет выглядеть так же, как tf.Module .

Модели Keras также можно сохранить с помощью tf.saved_model.save() , поскольку они являются модулями. Однако модели Keras имеют удобные методы и другие функции:

my_sequential_model.save("exname_of_file")
INFO:tensorflow:Assets written to: exname_of_file/assets

Так же легко их можно загрузить обратно:

reconstructed_model = tf.keras.models.load_model("exname_of_file")
WARNING:tensorflow:No training configuration found in save file, so the model was *not* compiled. Compile it manually.

SavedModels также сохраняет состояния метрик, потерь и оптимизатора.

Эту реконструированную модель можно использовать, и она будет давать тот же результат при вызове тех же данных:

reconstructed_model(tf.constant([[2.0, 2.0, 2.0]]))
<tf.Tensor: shape=(1, 2), dtype=float32, numpy=array([[5.5604653, 3.3511646]], dtype=float32)>

Нам нужно больше узнать о сохранении и сериализации моделей Keras, включая предоставление методов конфигурации для пользовательских слоев для поддержки функций. Ознакомьтесь с руководством по сохранению и сериализации .

Что дальше

Если вы хотите узнать более подробную информацию о Keras, вы можете следовать существующим руководствам по Keras здесь .

Еще одним примером высокоуровневого API, построенного на tf.module , является Sonnet от DeepMind, о котором рассказывается на их сайте .