Как правильно денормализовать ответ M xNet Neural Net?

Question

Я тренирую последовательный NN для регрессии. Входные данные должны быть нормализованы, а выходные данные должны быть денормализованы. Нормализация, в этом случае, означает масштабирование и перевод каждого ввода данных и обучение координат ожидаемого значения в 0..1, а денормализация означает преобразование выходных данных NN в исходные диапазоны.

Моя цель - иметь возможность сохранить нормализация и денормализация как операции слоя, когда NN сохраняется в файл вместо сохранения преобразований нормализации / денормализации в отдельном файле, что я сейчас и делаю.

До сих пор я подходил к проблеме через пользовательские слои. Например:

class NormDense( nn.Dense ):
   def __init__( self, factors, **kwargs ):
      super( NormDense, self ).__init__( **kwargs )
      with self.name_scope():
         self.scales = self.params.get( 'scales',
                                        shape = kwargs[ 'in_units' ],
                                        dtype = kwargs[ 'dtype' ],
                                        init = mx.init.Constant( factors[ 0 ].tolist() ),
                                        differentiable = False )
         self.shifts = self.params.get( 'shifts',
                                        shape = kwargs[ 'in_units' ],
                                        dtype = kwargs[ 'dtype' ],
                                        init = mx.init.Constant( factors[ 1 ].tolist() ),
                                        differentiable = False )

   def hybrid_forward( self, F, x, scales, shifts, *args, **kwargs ):
      xnormed = F.broadcast_add( F.broadcast_mul( x, scales ), shifts )
      return super( NormDense, self ).hybrid_forward( F, xnormed, *args, **kwargs )

class DenormDense( nn.Dense ):
   def __init__( self, factors, **kwargs ):
      super( DenormDense, self ).__init__( **kwargs )
      with self.name_scope():
         self.scales = self.params.get( 'scales',
                                        shape = kwargs[ 'units' ],
                                        dtype = kwargs[ 'dtype' ],
                                        init = mx.init.Constant( factors[ 0 ].tolist() ),
                                        differentiable = False )
         self.shifts = self.params.get( 'shifts',
                                        shape = kwargs[ 'units' ],
                                        dtype = kwargs[ 'dtype' ],
                                        init = mx.init.Constant( factors[ 1 ].tolist() ),
                                        differentiable = False )

   def hybrid_forward( self, F, x, scales, shifts, *args, **kwargs ):
      normed = super( DenormDense, self ).hybrid_forward( F, x, *args, **kwargs )
      return F.broadcast_add( F.broadcast_mul( normed, scales ), shifts )

def get_MunsellNet( n1, n2, n3, n4, NormDenorm ):
   Norm = numpy.transpose( NormDenorm[ 0:3 ] )
   Denorm = numpy.transpose( NormDenorm[ 3:5 ] )
   net = nn.HybridSequential()
   net.add( NormDense( Norm, units = n1, in_units = 3, activation = 'sigmoid', dtype = 'float64' ),
            nn.Dense( n2, activation = 'sigmoid', dtype = 'float64' ),
            nn.Dense( n3, activation = 'sigmoid', dtype = 'float64' ),
            nn.Dense( n4, activation = 'sigmoid', dtype = 'float64' ),
            DenormDense( Denorm, units = 2, dtype = 'float64' ) )
   net.hybridize()
   net.initialize( mx.init.Uniform(), ctx = ctx )
   return net

Я протестировал оба слоя, и они правильно рассчитали прямой проход. Однако, когда я обучаю эту сеть, она не учится, то есть и rmse, и l2loss вообще не прогрессируют. Если я заменю DenormDense на nn.Dense, NN обучается, как и ожидалось.

Возможен ли такой подход к нормализации / денормализации в M xNet?

Если да, чего мне не хватает заставить его работать?

Если нет, то есть ли другой способ добиться этого?