2021-06-02

# tensorflow自定义操作并附加自定义导数

## 说明

tensorflow最初给的例程中，虽然给了使用C++代码创建操作，并附加相应导数的能力，但是对于定义一个已有的复杂操作，并修改其导数的方法还不甚明晰。

本文结合@function.Defun() 和 @tf.RegisterGradient()以及g.gradient_override_map()来实现相关功能

## 示例

import tensorflow.compat.v1 as tf 
from tensorflow.python.framework import function

tf.compat.v1.disable_eager_execution() 

@function.Defun(tf.float32)
def customOp(x):
    return tf.identity(x * x * 10 - 2)

@tf.RegisterGradient("CustomOpGrad")
def _clip_grad(unused_op, grad):
  return grad * 0.1
 
x = tf.Variable([3.0], dtype=tf.float32)
 
g = tf.get_default_graph()
with g.gradient_override_map({customOp.name: "CustomOpGrad"}):
  y = customOp(x)

dydx = tf.gradients(y, x)

print("operation :", customOp.name)
print("x :", x)
print("y :", y)
print("dydx :", dydx)
 
with tf.Session() as sess:
  sess.run(tf.global_variables_initializer())
  print("dydx:", sess.run(dydx)[0])

## 结果

operation : customOp_yk92PMBQ5p0
x : <tf.Variable 'Variable:0' shape=(1,) dtype=float32>
y : Tensor("customOp_yk92PMBQ5p0:0", dtype=float32)
dydx : [<tf.Tensor 'gradients/customOp_yk92PMBQ5p0_grad/mul:0' shape=(1,) dtype=float32>]
dydx: [0.1]

## 说明

1. @function.Defun 注册customOp这个op.

需要注意的是，其name属性并不是叫customOp

2. @tf.RegisterGradient 注册CustomOpGrad这个op的求梯度操作

需要注意的是，其name属性就叫CustomOpGrad

3. g.gradient_override_map 将 customOp的求梯度操作替换成了CustomOpGrad

我看有些文章，给的示例都是tf.identity的

# 一般的示例
with g.gradient_override_map({"Identity": "CustomOpGrad"}):
  y = tf.identity(x, name="Identity")

# 对此示例不成立
# 只会影响tf.identity
# 不会影响tf.square
with g.gradient_override_map({"Identity": "CustomOpGrad"}):
  y = tf.identity(tf.square(x), name="Identity")

所以g.gradient_override_map这里填入的是定义的操作名，而不是操作之后的张量名

y就是一个张量，名字可以配置成Identity，但是我们想要的功能并没有实现

## 参考

这个可以按关键字自己百度，我是散的拼成+试错的