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()