[TF 2.2] tf.keras.layers.experimental.preprocessing.TextVectorization 한계와 해결 방법

이전에 tf.keras 의  layer 자체에 preprocessing을 하는 기능을 넣어 텍스트를 쉽게 vectorization하는 것의 기능을 살펴보았다. (tf.keras.layers.experimental.preprocessing.TextVectorization)

 

해당기능은 tensorflow 2.2 버전의 experimental로 들어와 있는데 이 기능의 한계를 알아보겠다.

 

먼저 모델 layer에 preprocessing.TextVectorization을 넣을때,

해당 layer를 넣기 전 text vocab구성을 위해 adapt 라는 method를 호출해 fit을 한다.

이렇게 하면 vocab을 형성하여 vectorization을 한다.

 

그러나 해당 모델을 잘 사용하고 저장 후 다시 불러오면 어떻게 될까?

 

결론부터 말하면 loaded한 모델에 있는 TextVectorization layer은 train했을 당시의 vocab을 가지고 있지 않아 잘못된 vectorization을 수행하게 된다.

 

코드로 살펴보자

import numpy as np
import tensorflow as tf
from tensorflow.keras.layers.experimental.preprocessing import TextVectorization

vector_layer =TextVectorization(max_tokens=5,output_mode='int',output_sequence_length=5)

data = np.asarray(["tensorflow is so easy","I am the greatest genius","keras is beautiful"])
vector_layer.adapt(data) ## 데이터를 fit 하여 vacab을 만든다.

vector_layer.get_vocabulary()
# [b'is', b'the', b'tensorflow', b'so']

## vectorization model
inputs = tf.keras.layers.Input(shape=(1,),dtype=tf.string,name='text')
outputs=vector_layer(inputs)
model=tf.keras.Model(inputs,outputs)

model.predict(["tensorflow is so easy"])
# array([[4, 2, 5, 1, 0]])

## save model
MODEL_EXPORT_PATH = './textvectors_test/'
tf.saved_model.save(model, MODEL_EXPORT_PATH)

## load model
loaded_model = tf.keras.models.load_model(MODEL_EXPORT_PATH)

### loaded model은 train 했을 당시의 vocab을 가지고 있지 않는 것 확인
loaded_model.predict(["tensorflow is so easy"])
# array([[1, 1, 1, 1, 0]])

위 코드에서 처럼 같은 문장을 모델을 저장하기 전과 후로 vector화 되는 것을 비교 했다.

저장한 모델을 다시불러온 모델은 vocab이 없어 oov로 text를 모두 1로 vector화하는 것을 알 수 있다.

 

이러한 문제를 해결하기 위해서는!!!!

## load model
loaded_model = tf.keras.models.load_model(MODEL_EXPORT_PATH, custom_objects={"TextVectorization":TextVectorization})

loaded_model.predict(["tensorflow is so easy"])
=====
array([[4, 2, 5, 1, 0]])

이렇게 custom_objects를 지정해주면 모델을 로드할때도 vocab을 가지고 있다.

 

ai-platform predict에 적용할때는 별도의 custom_objects를 지정하지 않아도 아주 잘 예측이 된다.

## load model
loaded_model = tf.keras.models.load_model(MODEL_EXPORT_PATH, custom_objects={"TextVectorization":TextVectorization})

@tf.function(input_signature=[tf.TensorSpec([None], dtype=tf.string), tf.TensorSpec([None, 1], dtype=tf.string)])
def keyed_prediction(key, data):
    pred = loaded_model(data, training=False)
    return {
        'output': pred,
        'key': key
    }
    
KEYED_EXPORT_PATH = './keyed_textvectors_test/'
loaded_model.save(KEYED_EXPORT_PATH, signatures={'serving_default': keyed_prediction})

os.environ["MODEL_LOCATION"] = KEYED_EXPORT_PATH

!gcloud ai-platform models create text_vectorization_layer \
  --regions us-central1
  
# Create model version based on that SavedModel directory
!gcloud beta ai-platform versions create v2 \
    --model text_vectorization_layer \
    --runtime-version 2.1 \
    --python-version 3.7 \
    --framework tensorflow \
    --origin ${MODEL_LOCATION} \
    --staging-bucket gs://daehwan
    
    
with open("keyed_string_input.json", "w") as file:
    print('{"data": ["tensorflow is so easy"], "key": "id_1234"}', file=file)
    
!gcloud ai-platform predict --model text_vectorization_layer --json-instances keyed_string_input.json --version v2
    
===========================
KEY      OUTPUT
id_1234  [4, 2, 5, 1, 0]