//+------------------------------------------------------------------+
//|                                              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(CheckPointer(m_cBatchOptions) != POINTER_INVALID)
      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.BufferInit(m_cWeights.Total(), 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(CheckPointer(prevLayer) == POINTER_INVALID ||
      CheckPointer(prevLayer.GetOutputs()) == POINTER_INVALID ||
      CheckPointer(m_cOutputs) == POINTER_INVALID ||
      CheckPointer(m_cBatchOptions) == POINTER_INVALID ||
      CheckPointer(m_cWeights) == POINTER_INVALID ||
      CheckPointer(m_cActivation) == POINTER_INVALID)
      return false;
//--- Проверка размера пакета нормализации
   if(m_iBatchSize <= 1)
     {
      if(!m_cOutputs.AssignArray(prevLayer.GetOutputs()))
         return false;
      if(!m_cActivation.Activation(m_cOutputs))
         return false;
      return true;
     }
//--- Разветвление алгоритма по вычислительному устройству
   if(CheckPointer(m_cOpenCL) == POINTER_INVALID)
     {
      int total = m_cOutputs.Total();
      CBufferDouble *inputs = prevLayer.GetOutputs();
      for(int i = 0; i < total; i++)
        {
         int shift_options = i * 3;
         int shift_weights = i * 2;
         double mean = m_cBatchOptions[shift_options] * ((double)m_iBatchSize - 1.0) + inputs[i];
         if(m_cBatchOptions[shift_options] != 0 && m_cBatchOptions[shift_options + 1] != 0)
            mean /= (double)m_iBatchSize;
         double delt = inputs[i] - mean;
         double variance = m_cBatchOptions[shift_options + 1] * ((double)m_iBatchSize - 1.0) + MathPow(delt, 2);
         if(m_cBatchOptions[shift_options + 1] > 0)
            variance /= (double)m_iBatchSize;
         double nx = delt / MathSqrt(variance + 1e-8);
         //---
         if(m_cWeights[shift_weights] == 0)
            if(!m_cWeights.Update(shift_weights, 1))
               return false;
         //---
         double res = m_cWeights[shift_weights] * nx + m_cWeights[shift_weights + 1];
         if(!m_cOutputs.Update(i, res))
            return false;
         if(!m_cBatchOptions.Update(shift_options, mean) ||
            !m_cBatchOptions.Update(shift_options + 1, variance) ||
            !m_cBatchOptions.Update(shift_options + 2, nx))
            return false;
        }
     }
   else  // Блок OpenCL
     {
      //--- Создание буферов данных
      CBufferDouble *inputs = prevLayer.GetOutputs();
      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_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;
      //--- Постановка в очередь выполнения
      int off_set[] = {0};
      int s = m_cOutputs.Total();
      int d = s % 4;
      s = (s - d) / 4 + (d > 0 ? 1 : 0);
      int NDRange[] = {s};
      if(!m_cOpenCL.Execute(def_k_BatchNormFeedForward, 1, off_set, NDRange))
         return false;
      //--- Получение результатов
      if(!m_cOutputs.BufferRead())
         return false;
      if(!m_cBatchOptions.BufferRead())
         return false;
      //--- Очистка памяти кернела
      inputs.BufferFree();
      m_cWeights.BufferFree();
      m_cBatchOptions.BufferFree();
      m_cOutputs.BufferFree();
     }
//---
   if(!m_cActivation.Activation(m_cOutputs))
      return false;
//---
   return true;
  }
//+------------------------------------------------------------------+
//| Метод распределение градиента через скрытый слой                 |
//+------------------------------------------------------------------+
bool CNeuronBatchNorm::CalcHiddenGradient(CNeuronBase *prevLayer)
  {
//--- Блок контролей
   if(CheckPointer(prevLayer) == POINTER_INVALID ||
      CheckPointer(prevLayer.GetOutputs()) == POINTER_INVALID ||
      CheckPointer(prevLayer.GetGradients()) == POINTER_INVALID ||
      CheckPointer(m_cActivation) == POINTER_INVALID ||
      CheckPointer(m_cBatchOptions) == POINTER_INVALID ||
      CheckPointer(m_cWeights) == POINTER_INVALID)
      return false;
//--- Корректировка градиента ошибки на производну функции активации
   if(!m_cActivation.Derivative(m_cOutputs, m_cGradients))
      return false;
//--- Проверка размера пакета нормализации
   if(m_iBatchSize <= 1)
      return prevLayer.GetGradients().AssignArray(m_cGradients);
//--- Разветвление алгоритма по вычислительному устройству
   if(CheckPointer(m_cOpenCL) == POINTER_INVALID)
     {
      CBufferDouble *inputs = prevLayer.GetOutputs();
      CBufferDouble *inputs_grad = prevLayer.GetGradients();
      int total = m_cOutputs.Total();
      for(int i = 0; i < total; i++)
        {
         int shift_options = i * 3;
         int shift_weights = i * 2;
         double inp = inputs[i];
         double gnx = m_cGradients[i] * m_cWeights[shift_weights];
         double temp = 1 / MathSqrt(m_cBatchOptions[shift_options + 1] + 1e-8);
         double gvar = (inp - m_cBatchOptions[shift_options]) / (-2 * pow(m_cBatchOptions[shift_options + 1] + 1.0e-8, 3.0 / 2.0)) * gnx;
         double gmu = (-temp) * gnx - gvar * 2 * (inp - m_cBatchOptions[shift_options]) / m_iBatchSize;
         double gx = temp * gnx + gmu / m_iBatchSize + gvar * 2 * (inp - m_cBatchOptions[shift_options]) / m_iBatchSize;
         if(!inputs_grad.Update(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 s = m_cOutputs.Total();
      int d = s % 4;
      s = (s - d) / 4 + (d > 0 ? 1 : 0);
      int NDRange[] = {s};
      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(CheckPointer(m_cBatchOptions) == POINTER_INVALID ||
      CheckPointer(m_cGradients) == POINTER_INVALID ||
      CheckPointer(m_cDeltaWeights) == POINTER_INVALID)
      return false;
//--- Проверка размера пакета нормализации
   if(m_iBatchSize <= 1)
      return true;
//--- Разветвление алгоритма по вычислительному устройству
   if(CheckPointer(m_cOpenCL) == POINTER_INVALID)
     {
      int total = m_cGradients.Total();
      for(int i = 0; i < total; i++)
        {
         int shift_weights = i * 2;
         int shift_options = i * 3;
         double grad = m_cGradients[i];
         double delta = m_cBatchOptions[shift_options + 2] * grad + m_cDeltaWeights[shift_weights];
         if(!m_cDeltaWeights.Update(shift_weights, delta))
            return false;
         if(!m_cDeltaWeights.Update(shift_weights + 1, grad + m_cDeltaWeights[shift_weights + 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[] = {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(CheckPointer(m_cBatchOptions) == POINTER_INVALID)
      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(CheckPointer(m_cBatchOptions) == POINTER_INVALID)
     {
      m_cBatchOptions = new CBufferDouble();
      if(CheckPointer(m_cBatchOptions) == POINTER_INVALID)
         return false;
     }
   if(!m_cBatchOptions.Load(file_handle))
      return false;
//---
   return true;
  }
//+------------------------------------------------------------------+