Tensorflow CuDNN RNN vs 그냥 RNN 비교

Cudnn은 GPU only로 GPU를 이용할 때 학습 속도가 더 빠르다는 장점이 있다.

 

www.tensorflow.org/api_docs/python/tf/keras/layers/GRU

tf.keras.layers.GRU  |  TensorFlow Core v2.2.0

위 page의 API 가이드대로 아래 조건을 만족하면 Cudnn Rnn 결과와 그냥 Rnn 결과가 정확히 일치한다.

현재는 이 조건을 만족할 경우 GPU가 잡히면 자동으로 Cudnn으로 학습을 진행하며,

TF 1.x버전에서 CuDNN을 import해서 사용할 수도 있는데, 이때는 weights의 configuration이 바뀌므로 이부분을 인지하고 있으면 좋다.

 

Cudnn의 단점은

1. mask 기능을 지원하지 않음 

2. weight구조가 keras api 와 달라서 inference(C++로 inference할 예정이기때문)를 하기 위해서는 weight reshape을 거쳐줘야 함.

 

단점이라기보다는 귀찮은 점..이라고 할 수 있다.

 

나는 실시간 inference가 필요한 프로젝트를 진행중이었기 때문에 mask기능이 필요했다.

 

Mask기능이 필요한 이유는 Batch로 학습시킬때, 학습데이터의 길이가 모두 일치해야하기 때문이다. 보통 0으로 padding을 해준다.

 

하지만 Padding된 길이만큼 Inference시 속도가 늦어지기 때문에, Mask기능을 사용하여 원래 Seqeunce길이만큼에 대해서만 학습하도록 하면 Padding없이도 실시간 Inference를 할 수 있다.

 

또한, Mask하지 않으면 Bidirectional 구조에서 forward , backward layer 끼리 결합할 때 문제가 생기는데

 

예를들어 input이 [1, 2, 3, 0, 0] 라고 가정하면

실제로 원하는 것은

forward :                                  -> [1, 2, 3, 0, 0]

backward : [3, 2, 1, 0, 0] ->reverse -> [1, 2, 3, 0, 0]

이런 순서로 계산되는 것이지만

 

Mask하지 않는다면

forward :                                   -> [1, 2, 3, 0, 0]

backward : [0, 0, 3, 2, 1] ->reverse ->  [1, 2, 3, 0, 0]

이런순서대로 output이 계산되어버린다.

 

즉, backward layer에서 3->2->1 과 0->0->3->2->1 은 결과가 다르기 때문에, 원하는 구조라고 할 수가 없게된다.

 

그래서 직접 Bidirectional GRU를 구현해보았다.  (Colab ver.)

 

%tensorflow_version 2.x
import tensorflow as tf
from tensorflow.keras import layers, Sequential
import numpy as np
import tensorflow.keras.backend as K
from tensorflow.compat.v1.keras.layers import CuDNNLSTM, CuDNNGRU,GRU
from tensorflow.python.ops import array_ops

tf.random.set_seed(21)
K.clear_session()

x = tf.constant([[1, 2, 3, 0, 0],
                 [1, 2, 3, 4, 0],
                 [1, 2, 0, 0, 0],
                 [1, 0, 0, 0, 0]],dtype=tf.int32)

embedding_layer = layers.Embedding(6,2)
embedding_x = embedding_layer(x)

#get reversed input sequence
tmp_index=tf.where(tf.equal(0,x),x=tf.range(tf.shape(x)[1])*tf.ones_like(x),y=(tf.shape(x)[1])*tf.ones_like(x))
x_length = tf.add(tf.reduce_min(tmp_index,axis=1),-1)
rx = tf.reverse_sequence(x, x_length, seq_axis=1, batch_axis=0)
embedding_rx = embedding_layer(rx)

# forward layer + reversed backward layer
f_cudnn_gru_cell = CuDNNGRU(2,return_sequences=True) #forward layer
b_cudnn_gru_cell = CuDNNGRU(2,return_sequences=True) # backward layer
xf = f_cudnn_gru_cell(embedding_x) # input
xb = b_cudnn_gru_cell(embedding_rx) # reversed_input
rxb = tf.reverse_sequence(txb, x_length, seq_axis=1, batch_axis=0)

# merge mode 'concat'
output = tf.concat([xf,rxb],2)

 

Input sequence와 backward layer를 뒤집어주는 것만 잘 신경쓰면 어렵지 않게 구현 가능하다.

이렇게 모델을구현 한 후 loss를 구할때 0으로 padding된 부분을 제외하고 loss를 계산하도록 하면 masking기능 없이도 같은 학습효과를 낼 수 있다.

 

 

CudnnRnn으로 학습한 모델을 CPU(standard Keras)로 inference하기 위해서는 weight Shape을 변형해주어야하는데

https://stackoverflow.com/questions/58807467/tensorflow-keras-cudnngru-to-gru-conversion

TensorFlow Keras CuDNNGRU to GRU conversion

위에 소스코드까지 친절하게 정리되어있으니, 그대로 갖다쓰면 된다.

 

Weight from CuDNNGRU

[array([[ 0.05315638,  0.48983377,  0.21304148, -0.61902946, -0.5883764 , 0.8355501 ],
        [ 0.6174968 , -0.06659609, -0.7425384 , -0.09426457,  0.267847  , -0.10589522]], dtype=float32),
 array([[-0.6191499 , -0.00290711, -0.3616762 ,  0.498306  , -0.34991384,  0.33924466],
        [ 0.33154872,  0.03050954,  0.03058927,  0.340796  ,  0.5393269 ,   0.69368124]], dtype=float32),
 array([0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.], dtype=float32)]

 

Weight  from tf.keras.layers.GRU

[array([[ 0.05315638,  0.6174968 ,  0.21304148, -0.7425384 , -0.5883764 ,  0.267847  ],
         [ 0.48983377, -0.06659609, -0.61902946, -0.09426457,  0.8355501, -0.10589522]], dtype=float32),
 array([[-0.6191499 ,  0.33154872, -0.3616762 ,  0.03058927, -0.34991384,  0.5393269 ],
        [-0.00290711,  0.03050954,  0.498306  ,  0.340796  ,  0.33924466, 0.69368124]], dtype=float32),
 array([[0., 0., 0., 0., 0., 0.],  [0., 0., 0., 0., 0., 0.]], dtype=float32)]

 

 이렇게 shape만 좀 다르고 weight는 같다.

 

 

실제로 Inference했을때 float유효자리수까지 정확히 일치함을 확인했다.