问题描述

在深度神经网络的训练过程中，经常需要自定义函数或一些变换
，如果tf.gradient无法对自定义的函数进行自动求导计算梯度，
那就需要对自定义的函数补充梯度计算函数。这里，我们以多输
出神经网络结合自定义损失函数的例子展示如何使用TF2.0 进行
梯度下降并优化参数

解决方法

• 定义神经网络

  定义神经网络结构，并且为多输出网络

  class MyNet(tf.keras.Model):
      def __init__(self):
          super(MyNet, self).__init__()
          self.deep_cnn=DeepCNN()
          self.avgpool = AvgPool()
          self.fc_1 = Dense(classes)
          self.fc_2 = Dense(classes)
   
      def call(self,x):
          x = self.deep_cnn(x)
          x = self.avgpool(x)
          output1 = self.fc_1(x)
          output2 = self.fc_2(x)
          
          return output1,output2

• 自定义损失函数&计算梯度函数

  损失函数对神经网络的输出做一定变换，并计算出损失值，此处应在损失
  函数内定义梯度计算函数，并返回梯度计算函数对象。损失函数需要使用
  @tf.custom_gradient装饰

  @tf.custom_gradient
  def myLossFn(output1,output2,labels):
      ...
      ...
      loss = f_loss(output1,output2,labels)  # 计算损失值
      
      def grad_fn(grad):  # 计算梯度函数
          ...
          ...
          grad1,grad2=f_grad(output1,output2,labels)  #只需要计算output1和output2 的梯度
          
          return grad_1,grad2,None  # 由于损失函数有三个参数，所以需要返回3个值
      
      return loss,grad_fn

• 训练设置

  在 forward propagation之后进行损失值计算，
  然后计算梯度,然后通过梯度更新参数

  # training setting
  model = MyNet()
  ...
  with tf.GradientTape() as tape:
      output1,output2=model(x)
      loss = myLossFn(output1,output2,labels)
  grad_list = model.trainable_variables
  grad_list += [output1,output2]
  grad = tape.gradient(loss,grad_list)
  optimazer.apply_gradients(zip(grad,grad_list)) # 如果output1和output2没有可优化的参数，可以使用model.trainable_variables替换grad_list
  ...