Como fazer op personalizada de saída ponderada softmax no mxnet?
Quero substituir mx.symbol.SoftmaxOutput pela versão ponderada (atribua peso diferente em relação à frequência do rótulo em todo o conjunto de dados)
A função original funciona bem como abaixo:
cls_prob = mx.symbol.SoftmaxOutput(data=data,
label=label,
multi_output=True,
normalization='valid',
use_ignore=True,
ignore_label=-1,
name='cls_prob')
O código atual que escrevi como abaixo. O código pode ser executado sem erros, mas a perda explode rapidamente para nan. Estou lidando com um problema de detecção, a perda RCNNL1 rapidamente se torna nan quando uso meu código como CustomOp. Outra coisa é que preciso ignorar o rótulo -1 e não tenho certeza de como fazê-lo corretamente. Qualquer ajuda será muito 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)