Keras 모델 load시 Unknown layer, 외부 라이브러리 해결 방법

Keras를 사용하다보면, 간단하게 구현된 라이브러리를 사용할 때가 있다.

그래서 더 사랑을 받는게 아닐까. 너무 간결하고 간단한 라이브러리들이 많다.

최근에 사용했던 간단한 라이브러리 Transformer 모델의 Self-Attention이다.

이 Function을 사용하기 위해서는 따로 import를 해줘야하는데 아래와 같이 하면 된다.

from keras_self_attention import SeqSelfAttention

자 이제 CNN 몇개와 RNN 몇개를 쌓고, 그 아래에 Self-Attention을 쌓았다고 가정해보자.

이 모델을 Fine-tune 하기 위해 일정 learning rate로 돌려놓고, 어느정도 loss를 깎은 다음 일정 이상 loss가 떨어지지 않으면 early stopping을 하여 learning rate를 기존보다 0.1배 정도로 낮춰서 fine-tune 할 예정이다. 이를 위해 모델을 save해야한다.

그럼 모델을 load할 때 종속성 문제가 생길 수 있다. 위의 SeqSelfAttention 인식을 못 할 수 있는 것이다.

내가 안되었던 경우는 다음과 같다. 

from keras.models import load_model, Model
import tensorflow as tf
import os
import time
from keras_self_attention import SeqSelfAttention
global graph, model
graph = tf.get_default_graph()
### GPU Setting  : Not Using GPU
os.environ["CUDA_VISIBLE_DEVICES"]="-1"

# Load Model
model = load_model('/data/Auditory_Emotion_Recognition/model_attention/fianl_emotion_model.h5')

코드를 보면 단지 model을 load하는데, Unknown layer라고 뜬다. 아래의 사진은 그 현상을 얘기한다.

즉 SeqSelfAttention을 import 하였는데도 이러한 문제가 생긴다.

이를 해결 하기 위해 model load시 다음과 같은 줄을 선언해주면 쉽게 해결할 수 있다.

model = load_model('/data/Auditory_Emotion_Recognition/model_attention/fianl_emotion_model.h5',custom_objects={'SeqSelfAttention':SeqSelfAttention})

custom_objects={'name':import_library_name}

위와 같이 custom_objects 를 선언해주면 사용 가능하다.