Quiero reemplazar mx.symbol.SoftmaxOutput con la versión ponderada (asigne un peso diferente con respecto a la frecuencia de la etiqueta en todo el conjunto de datos)

La función original funciona bien como a continuación:

cls_prob = mx.symbol.SoftmaxOutput(data=data,
                                   label=label,
                                   multi_output=True,
                                   normalization='valid',
                                   use_ignore=True, 
                                   ignore_label=-1,
                                   name='cls_prob')

El código actual que escribí como a continuación. El código puede ejecutarse sin errores, pero la pérdida explota rápidamente a nan. Estoy lidiando con un problema de detección, la pérdida de RCNNL1 se convierte rápidamente en nan cuando uso mi código como CustomOp. Otra cosa es que tengo que ignorar la etiqueta -1 y no estoy seguro de cómo hacerlo correctamente. Cualquier ayuda será apreciada.

import mxnet as mx
import numpy as np

class WeightedSoftmaxCrossEntropyLoss(mx.operator.CustomOp):
    def __init__(self, num_class):
        self.num_class = int(num_class)

    def forward(self, is_train, req, in_data, out_data, aux):

        data = in_data[0]
        label = in_data[1]
        pred = mx.nd.SoftmaxOutput(data, label, multi_output=True,
                               normalization='valid', use_ignore=True, ignore_label=-1,
                               name='rcnn_cls_prob')

        self.assign(out_data[0], req[0], pred)

    def backward(self, req, out_grad, in_data, out_data, in_grad, aux):
        cls_weight = np.array([
            0.002852781814876101, 
            0.30715984513157385, 
            1.0932468996115976, 
            1.1598757152765971, 
            0.20739109264009636, 
            1.1984256112776808, 
            0.18746186040248036, 
            2.9009928470737023, 
            0.92140970338602113, 
            1.200317380251021
        ])
        label = in_data[1]
        pred = out_data[0]
        label = label.asnumpy().astype('int32').reshape((-1))
        pred = pred.asnumpy().reshape((pred.shape[0], pred.shape[1], -1)).transpose((0, 2, 1))
        pred = pred.reshape((label.shape[0], -1))

        # Need to ignore label (how)
        out_inds = np.where(label == -1)[0]
        #label = label[keep_inds]
        one_hot = np.zeros((label.shape[0], self.num_class))
        one_hot[np.arange(label.shape[0]), label] = 1
        # gradient
        dx = pred - one_hot
        #dx[out_inds] = 0.0
        weighted_dx = cls_weight * dx / 4
        self.assign(in_grad[0], req[0], weighted_dx)

@mx.operator.register("weighted_softmax_ce_loss")
class WeightedSoftmaxCrossEntropyLossProp(mx.operator.CustomOpProp):
    def __init__(self, num_class):
        super(WeightedSoftmaxCrossEntropyLossProp, self).__init__(need_top_grad=False)
        self.num_class = num_class

    def list_arguments(self):
        return ['data', 'label']

    def list_outputs(self):
        return ['output']

    def infer_shape(self, in_shapes):
        data_shape = in_shapes[0]
        label_shape = (in_shapes[0][0],)
        output_shape = in_shapes[0]
        return [data_shape, label_shape], [output_shape], []

    def create_operator(self, ctx, in_shapes, in_dtypes):
        #  create and return the CustomOp class.
        `enter code here`return WeightedSoftmaxCrossEntropyLoss(self.num_class)