MXNet-获取中间层输出及相应的梯度

最近为了得到CNN中每一层激活值对应的梯度（自动计算的梯度只有每一层对应的权值和偏置，MX又没有pytorch那样的backward_hook那样的神器），考虑用autograd.中的一些函数来实现.

• autograd.attach_grad()只能对最后一个attach的变量进行求梯度（不懂怎么用，目前没见过对多个变量用attach_grad）
  想到两种方法：
  1. register_forward_hook，虽然mx有一个这样的函数，确没有任何官方用例，google: mxnet register_forward_hook出来的全是pytorch的问答，真的气。另外，gluon.nn.Conv2D还没有register_forward_hook这个接口。这条路是死了
  2. 参考 这个，自己写一个conv层。但是这里又碰到另一个坑，不知道怎么用另一模型中的参数初始化新的模型。后来用了以下代码来初始化：

  @mx.init.register
  class myInitializer(mx.init.Initializer):
      def __init__(self, weight,bias):
          super(myInitializer,self).__init__()
          self.weight = weight[0]
          self.bias = bias[0]
      def _init_weight(self, _, arr):
          arr[:] = self.weight
      def _init_bias(self, _, arr):
          arr[:] = self.bias
  for ii, layer in enumerate(model.features):
      if isinstance(layer, nn.Conv2D):
          fsize = layer.weight.shape[2]
          new_layer = gradcam.Conv2D(layer._channels,(fsize,fsize))
          new_layer.initialize(init=myInitializer(layer.weight._data,
                                                        layer.bias._data), ctx=ctx)
          self.features.add(new_layer)
      elif ii<31: self.features.add(layer) 
  

  • 然后在需要获取梯度的地方用gradcam中的get_conv_out_grad方法就行了，不过这个方法一次只能获取一个梯度。
  • 需要注意的是，在官方文档上有说明如果在autograd.record(train_mode=False)，那么backward(train_mode=False)，否则梯度将不被定义
  3. 用自定义identity operation