mxnet中自定义损失函数，并自动求导

mxnet中如果想自定义一个损失函数，并且利用autograd模块求导，需要重写该损失函数的前向和后向函数，下面以mxnet中一个例子来说明

class sigmoid(mx.autograd.Function):
    def forward(self, x):
        y = 1 / (1 + mx.nd.exp(-x))
        self.save_for_backward(y)
        return y

    def backward(self, dy):
        # backward takes as many inputs as forward's return value,
        # and returns as many NDArrays as forward's arguments.
        y, = self.saved_tensors
        return dy * y * (1-y)

 

以上是一个sigmoid函数的实例，首先计算出前向表达式 y = 1 / (1 + mx.nd.exp(-x))，由于反向传播需要用到，保存为self.saved_tensors，反向传播时，dy为正向传播时输出的导数，返回一个同正向传播输入具有相同shape的值。

func = sigmoid()
x = mx.nd.random.uniform(shape=(10,))
x.attach_grad()

with mx.autograd.record():
    m = func(x)
    m.backward()

-------------------------------分割线-----------------------------------------------

说一说用这种方式自定义loss的一个坑，由于训练过程中每次迭代都要新建一个Function的实例，而Function中又保存了saved_tensors等中间张量，计算的异步特性导致重新申请新的显存的速度超过显存回收的速度，因此训练时显存一直往上涨！最后的解决方案，每次迭代后手动释放自定义loss中saved_tensors等中间张量，并添加同步nd.waitall()