Implement Softmax with Loss layer using python but keep getting high loss

[yunchi]

I'm trying to design my own loss layer to train FCN, and my first step is to implement the softmax with loss layer writing with python.

But I keep getting a high loss while training. Here's my code in SoftmaxLoss.py:

class SoftmaxLossLayer(caffe.Layer):

    """

    Compute the Softmax Loss in the same manner but consider soft labels

    as the ground truth

    """

    def setup(self, bottom, top):

        # check input pair

        if len(bottom) != 3:

            raise Exception("Need three inputs to compute distance.")

    def reshape(self, bottom, top):

        # check input dimensions match

        if bottom[0].num != bottom[1].num:

            raise Exception("Inputs must have the same dimension.")

        # difference is shape of inputs

        self.diff = np.zeros_like(bottom[0].data, dtype=np.float32)

        # loss output is scalar

        top[0].reshape(1)

    def forward(self, bottom, top):

        print 'bottom[0].num: {}'.format(bottom[0].num)

        scores = bottom[0].data

        scores -= np.max(scores)

        exp_scores = np.exp(scores)

        probs = exp_scores / np.sum(exp_scores, axis=1, keepdims=True)

        logprob = -np.log(probs)

        # ignore label

        for x in range(scores[0,0].shape[0]):

            for y in range(scores[0,0].shape[1]):

                if bottom[1].data[0,0,x,y] == 255:

                    logprob[:,:,x,y] = 0

        data_loss = np.sum(np.sum(logprob,axis=1))/bottom[0].num

        print 'data loss: {}'.format(data_loss)

        self.diff[...] = probs

        top[0].data[...] = data_loss

    def backward(self, top, propagate_down, bottom):

        delta = self.diff

        for i in range(2):

            if not propagate_down[i]:

                continue

            if i==0:

                delta[range(bottom[0].num), np.array(bottom[1].data,dtype=np.uint16)] -= 1

            bottom[i].diff[...] = delta/bottom[0].num

Could anyone tell me what's going on here?