5.Создание класса нового нейронного слоя

Давайте перейдем к практической части и посмотрим на реализацию нашего нейронного слоя многоголового внимания. Для его реализации мы создаем новый класс MHAttention наследником от базового класса всех нейронных слоев tf.keras.layers.Layer.

# Модель Multi-Head Self-Attention
class MHAttention(tf.keras.layers.Layer):

Вначале мы переопределим метод инициализации слоя __init__. В параметрах метода инициализации мы будем указывать две константы:

key_size — размер вектора описания одного элемента последовательности в теноре ключей Key;
heads — количество голов внимания.

В теле метода мы сохраним параметры в локальные переменные для будущего использования и сразу посчитаем размер конкатенированного выхода голов внимания в переменную m_iDimension.

Для вашего удобства я постарался максимально повторить наименование переменных из реализации MQL5.

Затем объявим внутренние объекты нашего нейронного слоя. Но обратите внимание, что в данном случае мы не указываем размер вектора одного элемента последовательности исходных данных. Это стало возможно благодаря использованию многомерных тензоров.

Библиотека TensorFlow работает с многомерными массивами или тензорами, представленными в виде объектов. Такой подход делает понимание модели более удобным и наглядным. Для возможности реализации в OpenCL мы были вынуждены использовать одномерные буферы данных, а для получения доступа к необходимому элементу рассчитывали смещение в одномерном буфере. Сейчас же, при использовании многомерных массивов, для доступа к элементу матрицы нам достаточно указать строку и столбец элемента. Это удобно и наглядно.

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

Но есть и вторая сторона медали. На момент инициализации класса мы еще не получили исходные данные. Соответственно, их размерность нам не известна. И пользователь не указал их в параметрах. Поэтому не все объекты мы можем создать в методе инициализации. Но не беда. Сделаем то, что можем.

В методе инициализации мы объявим объекты, создание которых возможно без понимания размерности исходных данных:

m_cQuerys — нейронный слой формирования конкатенированного тензора запросов Query;
m_cKeys — нейронный слой формирования конкатенированного тензора ключей Key;
m_cValues — нейронный слой формирования конкатенированного тензора значений Values;
m_cNormAttention — слой нормализации данных блока Multi-Head Self-Attention;
m_cNormOutput — слой нормализации результатов нейронного слоя.

  def __init__(self,key_size, heads, **kwargs):
    super(MHAttention, self).__init__(**kwargs)

    self.m_iHeads = heads
    self.m_iKeysSize = key_size
    self.m_iDimension=self.m_iHeads*self.m_iKeysSize;

    self.m_cQuerys = tf.keras.layers.Dense(self.m_iDimension)
    self.m_cKeys = tf.keras.layers.Dense(self.m_iDimension)
    self.m_cValues = tf.keras.layers.Dense(self.m_iDimension)
    self.m_cNormAttention=tf.keras.layers.LayerNormalization(epsilon=1e-6)
    self.m_cNormOutput=tf.keras.layers.LayerNormalization(epsilon=1e-6)

После создания метода инициализации перейдем к работе над методом build. Именно этот метод позволит нам инициализировать недостающие объекты. Данный метод запускается только один раз перед первым вызовом метода call и получает в параметрах размерность исходных данных. Благодаря этому мы можем инициализировать объекты, структуры и/или параметры, которые зависят от размера исходных данных.

В теле метода мы сохраняем последнее измерение тензора исходных данных в качестве размера вектора описания одного элемента последовательности исходных данных в локальную переменную m_iWindow. Потом создадим еще три внутренних нейронных слоя:

m_cW0 — полносвязный слой понижающей матрицы W0;
m_cFF1 — первый полносвязный слой блока Feed Forward;
m_cFF2 — второй полносвязный слой блока Feed Forward.

  def build(self, input_shape):
    self.m_iWindow=input_shape[-1]
    self.m_cW0 = tf.keras.layers.Dense(self.m_iWindow)
    self.m_cFF1=tf.keras.layers.Dense(4*self.m_iWindow,
                                      activation=tf.nn.swish)
    self.m_cFF2=tf.keras.layers.Dense(self.m_iWindow)

Вот мы и определили все внутренние объекты, необходимые для реализации алгоритма Multi-Head Self-Attention внутри нашего нового слоя. Но прежде чем приступить к реализации, давайте еще раз посмотрим каким образом мы можем записать алгоритм многоголового внимания средствами матричной математики. Ведь работая с многомерными тензорами, мы должны оперировать матричными операциями.

Первый шаг — определение тензоров Query, Key, Value. Для получения данных запросов нам необходимо умножить тензор исходных данных на соответствующую матрицу весовых коэффициентов. Эту операцию мы возлагаем на три внутренних нейронных слоя.

  def call(self, data):
    batch_size = tf.shape(data)[0]
    query = self.m_cQuerys(data)
    key = self.m_cKeys(data)
    value = self.m_cValues(data)

Второй шаг — определяем матрицу коэффициентов зависимости. По алгоритму Self-Attention сначала необходимо умножить тензор запросов на транспонированный тензор ключей.

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

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

Описанную процедуру вынесем в отдельную функцию split_heads.

  def split_heads(self, x, batch_size):
    x = tf.reshape(x, (batch_size, -1,
                                self.m_iHeads,
                                self.m_iKeysSize))
    return tf.transpose(x, perm=[0, 2, 1, 3])

Внутри метода call преобразуем тензоры и перемножим их согласно алгоритму Self-Attention.

    query = self.split_heads(query, batch_size)
    key = self.split_heads(key, batch_size)
    value = self.split_heads(value, batch_size)
    score = tf.matmul(query, key, transpose_b=True)

Далее полученные коэффициенты зависимости нам необходимо разделить на квадратный корень из размерности вектора ключей и нормализовать функцией Softmax по последнему измерению тензора.

score = score / tf.math.sqrt(tf.cast(self.m_iKeysSize, tf.float32))
score = tf.nn.softmax(score, axis=-1)

Нам остается умножить нормализованные коэффициенты зависимости на тензор ключей Value.

attention = tf.matmul(score, value)

И в результате этой операции мы получим итог работы блока внимания для каждой головы внимания. Для продолжения алгоритма нам нужен конкатенированный тензор всех голов внимания. Следовательно, нам надо осуществить обратную процедуру трансформации тензора. Мы еще раз переставляем местами первое и второе измерения и меняем размерность тензора с трехмерного на двухмерное.

attention = tf.transpose(attention, perm=[0, 2, 1, 3])
attention = tf.reshape(attention,(batch_size, -1, self.m_iDimension))

После этого с помощью матрицы W0 мы приводим конкатенированный тензор результатов к размеру тензора исходных данных. Складываем два тензора и нормализуем полученный результат.

attention = self.m_cW0(attention)
attention=self.m_cNormAttention(data + attention)

На этом заканчивается первый блок алгоритма Multi-Head Self-Attention и далее следует два последовательных полносвязных слоя блока Feed Forward. Первый нейронный слой будет с функцией активации Swish, а второй — без функции активации.

output=self.m_cFF1(attention)
output=self.m_cFF2(output)

В заключение метода мы складываем тензоры результатов блоков Multi-Head Self-Attention и Feed Forward и нормализуем слой. Результат операций возвращаем в виде тензора.

output=self.m_cNormOutput(attention+output)
return output

Мы реализовали минимальный набор методов класса, достаточный для проверки его функциональных возможностей. Но в таком виде, мы не сможем сохранить модель с данным классом. А это, как вы понимаете, совсем нехорошо, ведь мы хотим построить и обучить модель с последующей возможностью промышленной эксплуатации. Поэтому, для нас возможность сохранения модели и последующее ее восстановление является одним из ключевых требований.

Прежде всего, для создания возможности сохранения нового объекта, коим является наш нейронный слой, необходимо добавить его в список пользовательских объектов и дать возможность сериализации объекта. Это позволяет сделать директива register_keras_serializable, которую мы добавим перед объявлением класса нашего нейронного слоя.

# Модель Multi-Head Self-Attention
@tf.keras.utils.register_keras_serializable(package="Custom", name='MHAttention')
class MHAttention(tf.keras.layers.Layer):

Но это не все. Нам еще нужно добавить метод get_config, который бы возвращал содержимое переменных для сохранения в файл. Обратите внимание, что среди переменных есть как указываемые пользователем при инициализации объекта класса, так и сохраняемые из размерности исходных данных. Ведь наши весовые коэффициенты настроены именно на эти размерности.

  def get_config(self):
    config={'key_size': self.m_iKeysSize,
            'heads': self.m_iHeads,
            'dimension': self.m_iDimension,
            'window': self.m_iWindow
            }
    base_config = super(MHAttention, self).get_config()
    return dict(list(base_config.items()) + list(config.items()))

За восстановление данных из списка конфигурации отвечает метод from_config. Здесь тоже есть свои нюансы. Дело в том, что в обычной логике в словаре конфигурации указываются параметры из метода инициализации класса. Но мы сохранили и данные, зависящие от размерности исходных данных. А их, как вы помните, нет в параметрах метода инициализации. В чистом виде мы получим ошибку о наличии неизвестных параметров. Поэтому в начале метода мы их удаляем из справочника конфигурации, но при этом сохраняем значения в локальные переменные. И только после этого восстанавливаем слой.

  @classmethod
  def from_config(cls, config):
    dimension=config.pop('dimension')
    window=config.pop('window')
    layer = cls(**config)
    layer._build_from_signature(dimension, window)
    return layer

После инициализации нашего нейронного слоя из справочника конфигурации нам необходимо передать в соответствующие переменные предварительно извлеченные нами значения о конфигурации исходных данных. Для выполнения этого функционала мы вызовем метод _build_from_signature, который нам предстоит также переопределить.

  def _build_from_signature(self, dimension, window):
    self.m_iDimension=dimension
    self.m_iWindow=window

На этом мы завершаем работу над классом нашего нейронного слоя и можем перейти к созданию модели для тестирования созданного нейронного слоя Multi-Head Self-Attention.

Построение Multi-Head Self-Attention в Python

5.2.4.2 Создание скрипта для тестирования технологии Multi-Head Self-Attention