[tensorflow] GradientTape

(tensorflow v2.4.0)

 

2021/01/12 - [ML & DL/tensorflow] - [tensorflow] Custom Training Loops (tf.GradientTape)

[tensorflow] Custom Training Loops (tf.GradientTape)

이전 게시글에서 Custom Training Loop를 직접 구현해서 모델을 학습시켜보았습니다. 이때, gradient를 계산하기 위해서 with tf.GradientTape() 를 사용하여 미분을 자동으로 계산하였습니다. 

이처럼 Tensorflow에서는 GradientTape를 통해서 주어진 입력 변수에 대한 연산의 미분값을 자동으로 계산할 수 있습니다. tf.GradientTape는 context 안에서 실행된 모든 연산(Ops)를 tape에 기록하고, Reverse mode differentiation(후진 방식 자동 미분)을 사용하여 기록된 연산의 gradient를 계산합니다.

 

이번에는 tf.GradientTape에 대해서 간단히 알아보겠습니다.

우선 GradientTape는 파이썬 with as 키워드를 통해서 사용할 수 있습니다. 그리고 code들이 with block 내에서 수행됩니다. 변수 tape는 GradientTape type의 객체가 되고, 이 변수는 나중에 gradient를 계산하는데 사용됩니다.

여기서 GradientTape의 파라미터로 persistent가 사용되었는데, 이 파라미터의 의미는 잠시 후에 설명하도록 하겠습니다.

 

GradientTape를 통한 미분값 계산은 다음과 같은 과정을 통해서 이루어집니다.

먼저, 예측값을 계산하고, 실제값과 예측값의 loss를 계산합니다.

그리고, gradient 메소드를 통해서 미분값을 계산합니다. 

tape.gradient(reg_loss, w)는 w에 대한 loss의 미분값이며, 동일하게 tape.gradient(reg_loss, b)는 b에 대한 loss의 미분값입니다.

마지막으로 이렇게 구한 gradient를 가지고 parameter w, b를 업데이트하면 한 번의 학습의 완료됩니다.

 

실제 tensorflow 모델을 사용하면 아래와 같이 사용할 수 있습니다.

 

기본 예제

import tensorflow as tf

# Define a 2x2 array of 1's
x = tf.ones((2,2))

with tf.GradientTape() as t:
    # Record the actions performed on tensor x with `watch`
    t.watch(x) 

    # Define y as the sum of the elements in x
    y =  tf.reduce_sum(x)

    # Let z be the square of y
    z = tf.square(y) 

# Get the derivative of z wrt the original input tensor x
dz_dx = t.gradient(z, x)

# Print our result
print(dz_dx)

주어진 x를 통해서 z가 계산되고 있습니다. 

 

그리고, Gradient tape는 한 번 사용된 후에는 기본적으로 만료(expire)됩니다. 무슨 말인지 살펴보도록 하죠.

x = tf.constant(3.0)

# Notice that persistent is False by default
with tf.GradientTape() as t:
    t.watch(x)
    
    # y = x^2
    y = x * x
    
    # z = y^2
    z = y * y

# Compute dz/dx. 4 * x^3 at x = 3 --> 108.0
dz_dx = t.gradient(z, x)
print(dz_dx)

위 계산과정을 통해서 x에 대한 z의 미분값은 t.gradient(z, x)를 통해서 구할 수 있습니다.

여기서 t.gradient(z, x)를 한번 더 호출하면 어떻게 될까요?

# If you try to compute dy/dx after the gradient tape has expired:
try:
    dy_dx = t.gradient(y, x)  # 6.0
    print(dy_dx)
except RuntimeError as e:
    print("The error message you get is:")
    print(e)

바로 에러가 발생합니다. 이처럼 Gradient Tape는 기본적으로 한 번 사용되면 만료가 되어 더 이상 사용할 수 없는 상태가 됩니다.

 

연산기록을 계속 유지하여 Gradient Tape를 사용하기 위해서는 아까 위에서 언급한 persistent 파라미터를 True로 설정해주어야 합니다. 이 파라미터를 True로 설정해주면, 사용되더라도 만료되지 않고 계속 Gradient Tape가 유지됩니다.

x = tf.constant(3.0)

# Set persistent=True so that you can reuse the tape
with tf.GradientTape(persistent=True) as t:
    t.watch(x)
    
    # y = x^2
    y = x * x
    
    # z = y^2
    z = y * y

# Compute dz/dx. 4 * x^3 at x = 3 --> 108.0
dz_dx = t.gradient(z, x)
print(dz_dx)

# You can still compute dy/dx because of the persistent flag.
dy_dx = t.gradient(y, x)  # 6.0
print(dy_dx)

주의해야할 점은 더 이상 tape 변수를 사용하지 않는다면 Del 키워드를 통해서 제거해주어야 합니다. 

만약 연산이 매우 큰 경우에는 메모리를 많이 차지하여 overflow가 발생할 수 있습니다.

# Drop the reference to the tape
del t  

 

고차 미분(higher order derivative)

Nested Gradient tapes를 통해서 고차 미분도 가능합니다.

x = tf.Variable(1.0)

with tf.GradientTape() as tape_2:
    with tf.GradientTape() as tape_1:
        y = x * x * x
    
    # The first gradient calculation should occur at leaset
    # within the outer with block
    dy_dx = tape_1.gradient(y, x)
d2y_dx2 = tape_2.gradient(dy_dx, x)

print(dy_dx)
print(d2y_dx2)

적어도 첫번째 미분 계산은 outer with block 안에 존재해야합니다. 물론 첫번째 미분 계산이 inner with block에 존재해도 무관합니다.

x = tf.Variable(1.0)

with tf.GradientTape() as tape_2:
    with tf.GradientTape() as tape_1:
        y = x * x * x
    
        # The first gradient calculation can also be within the inner with block
        dy_dx = tape_1.gradient(y, x)
d2y_dx2 = tape_2.gradient(dy_dx, x)

print(dy_dx)
print(d2y_dx2)

 

만약 첫번째 미분 계산이 outer with block 외부에서 수행되면 어떻게 될까요?

x = tf.Variable(1.0)

with tf.GradientTape() as tape_2:
    with tf.GradientTape() as tape_1:
        y = x * x * x

# The first gradient call is outside the outer with block
# so the tape will expire after this
dy_dx = tape_1.gradient(y, x)

# The tape is now expired and the gradient output will be `None`
d2y_dx2 = tape_2.gradient(dy_dx, x)

print(dy_dx)
print(d2y_dx2)

두번째 미분 계산이 유지되지 않고, 만료되는 상황이 발생하여서 2차미분은 수행되지 않습니다.

 

이 경우에 persistent 파라미터를 True로 설정해주면, 미분 계산의 위치에 상관없이 지속적으로 미분 계산이 가능할 것처럼 보이지만, 아래의 경우에도 두번째 미분계산은 None이 됩니다.

x = tf.Variable(1.0)

# Setting persistent=True still won't work
with tf.GradientTape(persistent=True) as tape_2:
    # Setting persistent=True still won't work
    with tf.GradientTape(persistent=True) as tape_1:
        y = x * x * x

# The first gradient call is outside the outer with block
# so the tape will expire after this
dy_dx = tape_1.gradient(y, x)

# the output will be `None`
d2y_dx2 = tape_2.gradient(dy_dx, x)

print(dy_dx)
print(d2y_dx2)

두번째 미분 계산이 정상적으로 수행되기 위해서는 아래와 같은 방법으로 코드를 수행해야 합니다.

1.

x = tf.Variable(1.0)

with tf.GradientTape() as tape_2:
    with tf.GradientTape() as tape_1:
        y = x * x * x

        dy_dx = tape_1.gradient(y, x)
        
        # this is acceptable
        d2y_dx2 = tape_2.gradient(dy_dx, x)

print(dy_dx)
print(d2y_dx2)

2. 

x = tf.Variable(1.0)

with tf.GradientTape() as tape_2:
    with tf.GradientTape() as tape_1:
        y = x * x * x

        dy_dx = tape_1.gradient(y, x)
        
    # this is also acceptable
    d2y_dx2 = tape_2.gradient(dy_dx, x)

print(dy_dx)
print(d2y_dx2)

3.

x = tf.Variable(1.0)

with tf.GradientTape() as tape_2:
    with tf.GradientTape() as tape_1:
        y = x * x * x

        dy_dx = tape_1.gradient(y, x)
        
# this is also acceptable
d2y_dx2 = tape_2.gradient(dy_dx, x)

print(dy_dx)
print(d2y_dx2)

위 3가지 방법을 통해서 정상적으로 두번째 미분 계산을 수행할 수 있습니다.

 

 

 

 

- 참조

Coursera - Custom and Distributed Training with Tensorflow : Week 1