텐서플로우에서 이미 정의되어 있지않거나, 기존의 함수를 새로운 함수로 대체하고 싶을 경우, back propagation을 위한 gradient를 새로 정의할 필요가 있을 때가 있다. 이런 경우, 텐서플로우에서 @tf.custom_gradient 데코레이터를 이용하여 custom gradient를 정의해 줄수 있다.
간단한 scalar 함수(y=f(x) 모양, 예를 들어 y = log(1+exp(x) ) 처럼)의 경우에는 https://www.tensorflow.org/api_docs/python/tf/custom_gradient 에 있는 예제를 보는 것으로 충분하다. 그러나, 결과값,output 이 array, vector, tensor 인 경우에 어떻게 custom gradient를 정의해야하는지는 예제를 찾기 어렵다. 나도 한참 코드를 만들고 고치면서 실험을 하고서야 방법을 알아낼 수 있었다.
그러나, 그렇게 고생해서 알아낸 것이 이미 위의 링크 페이지 맨 아래에 있는 설명을 잘 읽으면 되는 것이라는 것을 나중에야 알게 되었다.
원칙적으로 input이 vector이고 output이 vector인 경우 수학적인 의미의 gradient는 Jacobian이 된다.
예를 들어 input 이 (x1,x2) , output이 (y1,y2,y3) 인 경우, Jacobian은 dy1/dx1, dy1/dx2, dy2/dx1, dy2/dx2, dy3/dx1,dy3/dx2 를 모두 계산해야 한다.
그러나, 신경망에서 관심이 있는 gradient는 사실상 loss function 에 대한 변수의 미분이기 때문에, 실제로는 dy/dx 가 아니라 d(loss)/dx = d(loss)/dy * dy/dx 만 필요하다. 여기서 , d(loss)/dy 를 upstream 이라고 부르고, 그 shape는 output y 의 shape와 같다. custom_gradient로 정의하는 gradient는 바로 d(loss)/dx 이지, dy/dx 가 아니라는 것이다. (다만, Gradient tape를 이용하여 tape.gradient(y,x) 를 계산할 경우 y 가 scalar라면 upstream=1이라고 둘 수 있으므로 dy/dx 를 구하는 것과 마찬가지가 된다.)
gradient의 output인 d(loss)/dx 는 input x의 shape와 같다. 따라서, custom gradient에서 정의하는 함수는 upstream을 input 으로 받아서, input x의 shape와 같은 모양의 결과를 내도록 정의되어야한다. 즉, y 가 vector라면, Jacobian에서 y 의 vector 축으로 그 값들을 모두 합한 결과가 필요한 것이다.
위의 예의 경우에는 d(loss)/dx1 = d(loss)/dy1 * dy1/dx1+d(loss)/dy2 * dy2/dx1 +d(loss)/dy3 * dy3/dx1 로 구하고, 마찬가지로 d(loss)/dx2 를 계산하여 d(loss)/dx 를 return 하도록 custom gradient를 정의하면 된다.
한편 새로 정의하는 함수 자신이 variable을 가지고 있을 경우에는 , (y=f(x;w) 처럼), 변수에 대한 미분도 정의해 주어야 한다. 이 경우에도 d(loss)/dw 는 변수 w 의 shape와 같아지도록 custom gradient를 정의한다. (즉, d(loss)/dw= d(loss)/dy * dy/dw 에서 d(loss)/dy 의 shape는 y 의 shape와 같고, dy/dw의 Jacobian에서 redundant한 축에대해 더해주어 올바른 shape가 되도록 만들어야한다.)
다음은 Y = X.W +b 인 경우를 custom gradient를 이용하여 구현한 것이다. 여기서 X.shape는 (batch_size, input_size), W.shape 는 (unit_size, input_size), Y는 (batch_size, unit_size) 의 모양을 가지는 tensor 이다.
@tf.custom_gradient
def custom_op(x):
@tf.function
def _inner_function():
y = x @ w + b
return y # y shape = (batch_size,unit_size)
y = _inner_function()
def grads(upstream,variables):
# here upstream is a shape of (batch_size,unit_size)
assert variables[0] is w
dydx = upstream @ tf.transpose(w) # (batch_size,input_size)
dydw = tf.transpose(x) @ upstream # (input_size,unit_size)
dydb = tf.reduce_sum(upstream,axis=0) # (unit_size,)
return dydx, [dydw, dydb]
return y, grads
