Tensorflow 2.0 custom gradient 만드는 법( vector output, multiple out의 경우)

 텐서플로우에서 이미 정의되어 있지않거나, 기존의 함수를 새로운 함수로 대체하고 싶을 경우, 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