//+------------------------------------------------------------------+
//|                                              NeuronBatchNorm.mqh |
//|                                  Copyright 2021, MetaQuotes Ltd. |
//|                                             https://www.mql5.com |
//+------------------------------------------------------------------+
#property copyright "Copyright 2021, MetaQuotes Ltd."
#property link      "https://www.mql5.com"
//+------------------------------------------------------------------+
//| Подключаем библиотеки                                            |
//+------------------------------------------------------------------+
#include "neuronbase.mqh"
//+------------------------------------------------------------------+
//| Class CNeuronBatchNorm                                           |
//| Назначение: Класс пакетной нормализации                          |
//+------------------------------------------------------------------+
class CNeuronBatchNorm    :  public CNeuronBase
  {
protected:
   CBufferDouble     *m_cBatchOptions;
   uint              m_iBatchSize;       // Размер пакета

public:
                     CNeuronBatchNorm(void);
                    ~CNeuronBatchNorm(void);
   //---
   virtual bool      Init(CLayerDescription *description);
   virtual bool      FeedForward(CNeuronBase *prevLayer);
   virtual bool      CalcHiddenGradient(CNeuronBase *prevLayer);
   virtual bool      CalcDeltaWeights(CNeuronBase *prevLayer);
   //--- methods for working with files
   virtual bool      Save(const int file_handle);
   virtual bool      Load(const int file_handle);
   //--- method of identifying the object
   virtual int       Type(void)        const { return(defNeuronBatchNorm); }
  };
//+------------------------------------------------------------------+
//| Конструктор класса                                               |
//+------------------------------------------------------------------+
CNeuronBatchNorm::CNeuronBatchNorm(void)  :  m_iBatchSize(1)
  {
   m_cBatchOptions = new CBufferDouble;
  }
//+------------------------------------------------------------------+
//| Деструктор класса                                                |
//+------------------------------------------------------------------+
CNeuronBatchNorm::~CNeuronBatchNorm(void)
  {
   if(m_cBatchOptions)
      delete m_cBatchOptions;
  }
//+------------------------------------------------------------------+
//| Метод инициализации класса                                       |
//+------------------------------------------------------------------+
bool CNeuronBatchNorm::Init(CLayerDescription *description)
  {
   if(CheckPointer(description) == POINTER_INVALID ||
      description.window != description.count)
      return false;
   description.window = 1;
   if(!CNeuronBase::Init(description))
      return false;
//--- Инициализируем буфер обучаемых параметров
   if(!m_cWeights.m_mMatrix.Fill(0))
      return false;
//--- Инициализируем буфер параметров нормализации
   if(CheckPointer(m_cBatchOptions) == POINTER_INVALID)
     {
      m_cBatchOptions = new CBufferDouble();
      if(CheckPointer(m_cBatchOptions) == POINTER_INVALID)
         return false;
     }
   if(!m_cBatchOptions.BufferInit(description.count, 3, 0))
      return false;
   m_iBatchSize = description.batch;
//---
   return true;
  }
//+------------------------------------------------------------------+
//| Метод прямого прохода                                            |
//+------------------------------------------------------------------+
bool CNeuronBatchNorm::FeedForward(CNeuronBase *prevLayer)
  {
//--- Блок контролей
   if(!prevLayer || !prevLayer.GetOutputs() || !m_cOutputs || !m_cBatchOptions || !m_cWeights || !m_cActivation)
      return false;
//--- Проверка размера пакета нормализации
   if(m_iBatchSize <= 1)
     {
      m_cOutputs.m_mMatrix = prevLayer.GetOutputs().m_mMatrix;
      if(m_cOpenCL && !m_cOutputs.BufferWrite())
         return false;
      if(!m_cActivation.Activation(m_cOutputs))
         return false;
      return true;
     }
//--- Разветвление алгоритма по вычислительному устройству
   if(!m_cOpenCL)
     {
      VECTOR inputs = prevLayer.GetOutputs().m_mMatrix.Row(0);
      VECTOR mean = (m_cBatchOptions.m_mMatrix.Col(0) * ((double)m_iBatchSize - 1.0) + inputs) / (double)m_iBatchSize;
      VECTOR delt = inputs - mean;
      VECTOR variance = (m_cBatchOptions.m_mMatrix.Col(1) * ((double)m_iBatchSize - 1.0) + delt * delt) / (double)m_iBatchSize;
      VECTOR std = variance;
      for(uint r = 0; r < std.Size(); r++)
         std[r] = (std[r] > 0 ? sqrt(std[r]) : 1e-8);
      VECTOR nx = delt / std;
      VECTOR res = m_cWeights.m_mMatrix.Col(0) * nx + m_cWeights.m_mMatrix.Col(1);
      if(!m_cOutputs.m_mMatrix.Row(res, 0) ||
         !m_cBatchOptions.m_mMatrix.Col(mean, 0) ||
         !m_cBatchOptions.m_mMatrix.Col(variance, 1) ||
         !m_cBatchOptions.m_mMatrix.Col(nx, 2))
         return false;
     }
   else  // Блок OpenCL
     {
      //--- Создание буферов данных
      CBufferDouble *inputs = prevLayer.GetOutputs();
      if(inputs.GetIndex() < 0)
         return false;
      if(m_cBatchOptions.GetIndex() < 0)
         return false;
      if(m_cWeights.GetIndex() < 0)
         return false;
      if(m_cOutputs.GetIndex() < 0)
         return false;
      //--- Передача параметров кернелу
      if(!m_cOpenCL.SetArgumentBuffer(def_k_BatchNormFeedForward, def_bnff_inputs, inputs.GetIndex()))
         return false;
      if(!m_cOpenCL.SetArgumentBuffer(def_k_BatchNormFeedForward, def_bnff_weights, m_cWeights.GetIndex()))
         return false;
      if(!m_cOpenCL.SetArgumentBuffer(def_k_BatchNormFeedForward, def_bnff_options, m_cBatchOptions.GetIndex()))
         return false;
      if(!m_cOpenCL.SetArgumentBuffer(def_k_BatchNormFeedForward, def_bnff_outputs, m_cOutputs.GetIndex()))
         return false;
      if(!m_cOpenCL.SetArgument(def_k_BatchNormFeedForward, def_bnff_total, m_cOutputs.Total()))
         return false;
      if(!m_cOpenCL.SetArgument(def_k_BatchNormFeedForward, def_bnff_batch, m_iBatchSize))
         return false;
      //--- Постановка в очередь выполнения
      uint off_set[] = {0};
      uint NDRange[] = { (int)(m_cOutputs.Total() + 3) / 4 };
      if(!m_cOpenCL.Execute(def_k_BatchNormFeedForward, 1, off_set, NDRange))
         return false;
      //--- Получение результатов
      if(!m_cOutputs.BufferRead())
         return false;
     }
//---
   if(!m_cActivation.Activation(m_cOutputs))
      return false;
//---
   return true;
  }
//+------------------------------------------------------------------+
//| Метод распределение градиента через скрытый слой                 |
//+------------------------------------------------------------------+
bool CNeuronBatchNorm::CalcHiddenGradient(CNeuronBase *prevLayer)
  {
//--- Блок контролей
   if(!prevLayer || !prevLayer.GetOutputs() || !prevLayer.GetGradients() || !m_cActivation || !m_cBatchOptions || !m_cWeights)
      return false;
//--- Корректировка градиента ошибки на производну функции активации
   if(!m_cActivation.Derivative(m_cOutputs, m_cGradients))
      return false;
//--- Проверка размера пакета нормализации
   if(m_iBatchSize <= 1)
     {
      prevLayer.GetGradients().m_mMatrix = m_cGradients.m_mMatrix;
      if(m_cOpenCL && !prevLayer.GetGradients().BufferWrite())
         return false;
      return true;
     }
//--- Разветвление алгоритма по вычислительному устройству
   if(!m_cOpenCL)
     {
      VECTOR inputs = prevLayer.GetOutputs().m_mMatrix.Row(0);
      CBufferDouble *inputs_grad = prevLayer.GetGradients();
      uint total = m_cOutputs.Total();
      for(uint i = 0; i < total; i++)
        {
         double gnx = m_cGradients.m_mMatrix[1, i] * m_cWeights.m_mMatrix[i, 0];
         double temp = 1 / MathSqrt(m_cBatchOptions.m_mMatrix[i, 1] + 1e-8);
         double gvar = (inputs[i] - m_cBatchOptions.m_mMatrix[i, 0]) / (-2 * pow(m_cBatchOptions.m_mMatrix[i, 1] + 1.0e-8, 3.0 / 2.0)) * gnx;
         double gmu = (-temp) * gnx - gvar * 2 * (inputs[i] - m_cBatchOptions.m_mMatrix[i, 0]) / (double)m_iBatchSize;
         double gx = temp * gnx + gmu / (double)m_iBatchSize + gvar * 2 * (inputs[i] - m_cBatchOptions.m_mMatrix[i, 0]) / (double)m_iBatchSize;
         if(!inputs_grad.m_mMatrix.Flat(i, gx))
            return false;
        }
     }
   else  // Блок OpenCL
     {
      //--- Создание буферов данных
      CBufferDouble *inputs = prevLayer.GetOutputs();
      CBufferDouble *inputs_grad = prevLayer.GetGradients();
      if(inputs.GetIndex() < 0 && !inputs.BufferCreate(m_cOpenCL))
         return false;
      if(m_cBatchOptions.GetIndex() < 0 && !m_cBatchOptions.BufferCreate(m_cOpenCL))
         return false;
      if(m_cWeights.GetIndex() < 0 && !m_cWeights.BufferCreate(m_cOpenCL))
         return false;
      if(m_cOutputs.GetIndex() < 0 && !m_cOutputs.BufferCreate(m_cOpenCL))
         return false;
      if(m_cGradients.GetIndex() < 0 && !m_cGradients.BufferCreate(m_cOpenCL))
         return false;
      if(inputs_grad.GetIndex() < 0 && !inputs_grad.BufferCreate(m_cOpenCL))
         return false;
      //--- Передача параметров кернелу
      if(!m_cOpenCL.SetArgumentBuffer(def_k_BatchNormCalcHiddenGradient, def_bnhgr_inputs, inputs.GetIndex()))
         return false;
      if(!m_cOpenCL.SetArgumentBuffer(def_k_BatchNormCalcHiddenGradient, def_bnhgr_weights, m_cWeights.GetIndex()))
         return false;
      if(!m_cOpenCL.SetArgumentBuffer(def_k_BatchNormCalcHiddenGradient, def_bnhgr_options, m_cBatchOptions.GetIndex()))
         return false;
      if(!m_cOpenCL.SetArgumentBuffer(def_k_BatchNormCalcHiddenGradient, def_bnhgr_gradient, m_cGradients.GetIndex()))
         return false;
      if(!m_cOpenCL.SetArgumentBuffer(def_k_BatchNormCalcHiddenGradient, def_bnhgr_gradient_inputs, inputs_grad.GetIndex()))
         return false;
      if(!m_cOpenCL.SetArgument(def_k_BatchNormCalcHiddenGradient, def_bnhgr_total, m_cOutputs.Total()))
         return false;
      if(!m_cOpenCL.SetArgument(def_k_BatchNormCalcHiddenGradient, def_bnhgr_batch, m_iBatchSize))
         return false;
      //--- Постановка в очередь выполнения
      int off_set[] = {0};
      int NDRange[] = { (int)(m_cOutputs.Total() + 3) / 4 };
      if(!m_cOpenCL.Execute(def_k_BatchNormCalcHiddenGradient, 1, off_set, NDRange))
         return false;
      //--- Получение результатов
      if(!inputs_grad.BufferRead())
         return false;
      //--- Очистка памяти контекста OpenCL
      inputs.BufferFree();
      m_cWeights.BufferFree();
      m_cBatchOptions.BufferFree();
      m_cGradients.BufferFree();
     }
//---
   return true;
  }
//+------------------------------------------------------------------+
//| Метод распределения градиента до уровня матрицы весов            |
//+------------------------------------------------------------------+
bool CNeuronBatchNorm::CalcDeltaWeights(CNeuronBase *prevLayer)
  {
//--- Блок контролей
   if(!m_cBatchOptions || !m_cGradients || !m_cDeltaWeights)
      return false;
//--- Проверка размера пакета нормализации
   if(m_iBatchSize <= 1)
      return true;
//--- Разветвление алгоритма по вычислительному устройству
   if(!m_cOpenCL)
     {
      VECTOR grad = m_cGradients.m_mMatrix.Row(0);
      VECTOR delta = m_cBatchOptions.m_mMatrix.Col(2) * grad + m_cDeltaWeights.m_mMatrix.Col(0);
      if(!m_cDeltaWeights.m_mMatrix.Col(delta, 0))
         return false;
      if(!m_cDeltaWeights.m_mMatrix.Col(grad + m_cDeltaWeights.m_mMatrix.Col(1), 1))
         return false;
     }
   else
     {
      //--- Создание буферов данных
      if(m_cBatchOptions.GetIndex() < 0 && !m_cBatchOptions.BufferCreate(m_cOpenCL))
         return false;
      if(m_cGradients.GetIndex() < 0 && !m_cGradients.BufferCreate(m_cOpenCL))
         return false;
      if(m_cDeltaWeights.GetIndex() < 0 && !m_cDeltaWeights.BufferCreate(m_cOpenCL))
         return false;
      //--- Передача параметров кернелу
      if(!m_cOpenCL.SetArgumentBuffer(def_k_BatchNormCalcDeltaWeights, def_bndelt_delta_weights, m_cDeltaWeights.GetIndex()))
         return false;
      if(!m_cOpenCL.SetArgumentBuffer(def_k_BatchNormCalcDeltaWeights, def_bndelt_options, m_cBatchOptions.GetIndex()))
         return false;
      if(!m_cOpenCL.SetArgumentBuffer(def_k_BatchNormCalcDeltaWeights, def_bndelt_gradient, m_cGradients.GetIndex()))
         return false;
      //--- Постановка в очередь выполнения
      int off_set[] = {0};
      int NDRange[] = {(int)m_cOutputs.Total()};
      if(!m_cOpenCL.Execute(def_k_BatchNormCalcDeltaWeights, 1, off_set, NDRange))
         return false;
      //--- Получение результатов
      if(!m_cDeltaWeights.BufferRead())
         return false;
      //--- Очистка памяти контекста
      m_cWeights.BufferFree();
      m_cBatchOptions.BufferFree();
      m_cDeltaWeights.BufferFree();
     }
//---
   return true;
  }
//+------------------------------------------------------------------+
//| Метод сохранения объектов класса в файл                          |
//+------------------------------------------------------------------+
bool CNeuronBatchNorm::Save(const int file_handle)
  {
//--- Блок контролей
   if(!m_cBatchOptions)
      return false;
//--- Вызываем метод родителького класса
   if(!CNeuronBase::Save(file_handle))
      return false;
//--- Сохраняем размер пакета нормализации
   if(FileWriteInteger(file_handle, m_iBatchSize) <= 0)
      return false;
//--- Сохранение параметров нормализации
   if(!m_cBatchOptions.Save(file_handle))
      return false;
//---
   return true;
  }
//+------------------------------------------------------------------+
//| Метод востановления класса из данных в файле                     |
//+------------------------------------------------------------------+
bool CNeuronBatchNorm::Load(const int file_handle)
  {
//--- Вызываем метод родителького класса
   if(!CNeuronBase::Load(file_handle))
      return false;
   m_iBatchSize = FileReadInteger(file_handle);
//--- Инициализируем динамический массив параметров оптимизации
   if(!m_cBatchOptions)
      if(!(m_cBatchOptions = new CBufferDouble()))
         return false;
   if(!m_cBatchOptions.Load(file_handle))
      return false;
//---
   return true;
  }
//+------------------------------------------------------------------+