Tensorflow로 작성된 코드를 Keras로 바꾸기
앞에서 설명한 논문의 코드는 Tensorflow로 구현된 것이었습니다.
하지만 배치단위로 처리하는 코드였기 때문에 실시간으로 한 개씩 처리에는 적합하지 않았고,
배치마다 그래프를 다시 읽어야 하는 방식으로 구현되어 있었습니다.
그래서 Tensorflow로 된 코드를 Keras로 바꾸고, 그래프는 실행시 한번만 읽을 수 있도록 로직을 수정하기로 했습니다!
1. 코드 바꾸기
레이어 부분의 코드를 바꾸는 일은 어렵지 않습니다.
tensorflow에서 지원하는 layer들은 대부분 Keras에도 있기 때문에 함수명과 인자만 잘 맞춰주면 됩니다.
예를 들어 아래 tensorflow 코드는
layer01 = tf.layers.conv2d(images, 64, 3,
padding="same",
activation=tf.nn.relu,
data_format=self._data_format,
name="conv1/conv1_1")
layer02 = tf.layers.conv2d(layer01, 64, 3,
padding="same",
activation=tf.nn.relu,
data_format=self._data_format,
name="conv1/conv1_2")
layer03 = tf.layers.max_pooling2d(layer02, 2, 2,
data_format=self._data_format)
Keras로 이렇게 변환하면 됩니다. (Keras의 layer 작성 방식 중 Functional API를 사용했습니다.)
from keras.layers import Conv2D
from keras.layers.pooling import MaxPooling2D
x = Conv2D(filters=64, kernel_size=3,
strides=(1, 1), padding='same',
activation='relu', name='conv1/conv1_1',
data_format=self._data_format)(img_input)
x = Conv2D(filters=64, kernel_size=3,
strides=(1, 1), padding='same',
activation='relu', name='conv1/conv1_2',
data_format=self._data_format)(x)
x = MaxPooling2D(pool_size=(2, 2), strides=(2, 2),
name="pool1", padding='valid',
data_format=self._data_format)(x)
여기서 data_format 인자는 GPU 사용시 쓰이는 인자 입니다.
기본적으로 tensor는 (num, w, h, channel) 순서로 구성되어 있고, 그렇기 때문에 data_format의 디폴트값은 data_format='channel_last' 입니다.data_format='channel_first' 라고 하면 tensor의 순서가 (channel, num, w, h)로 변경되고, GPU 연산을 위해서는 반드시 명시해 주어야 하는 인자 입니다.
문제는 tensorflow는 layer 중간에 tensor를 처리하는 함수를 넣을 수 있지만, Keras는 layer를 쌓아가는 형태이기 때문에 별도의 처리함수를 사용할 수 없다는 점이었습니다.
예를들어
branch5 = tf.reduce_mean(features,
axis=self._dims_axis,
keepdims=True)
branch5 = tf.layers.conv2d(branch5, 256, 1,
padding="valid",
activation=tf.nn.relu,
data_format=self._data_format,
name="aspp/conv1_5")
tensorflow에서는 tensor처리 후 layer를 사용할 수 있지만, Keras에서는 사용할 수 없습니다.
Keras에서 이런 방식을 사용하려면 첫번째 branch5도 layer 형태여야 합니다.
Keras에서는 Lambda 라는 layer를 지원하고 있습니다.
Lambda를 이용하면 tensor처리한 것을 layer형태로 만들어 다음 layer에 적용할 수 있습니다.
from keras import backend as K
from keras.layers import Lambda, Conv2D
branch5 = Lambda(lambda inputs: K.mean(inputs,
axis=self._dims_axis,
keepdims=True))(encoder_output)
branch5 = Conv2D(filters=256, kernel_size=1,
strides=(1, 1), padding='valid',
activation='relu', name='aspp/conv1_5',
data_format=self._data_format)(branch5)
이렇게 하면 tensorflow에서 사용하는 여러 처리 함수들도 쉽게 불러와서 쓸 수 있습니다.
#tensorflow
stack = tf.image.resize_bilinear(stack, (shape[self._dims_axis[0]] * factor,
shape[self._dims_axis[1]] * factor))
#Keras
stack = Lambda(lambda inputs: tf.image.resize_bilinear(stack,
(shape[self._dims_axis[0]]*factor, shape[self._dims_axis[1]]*factor))
2. Weight 파일 바꾸기
사용한 딥러닝 모델의 깃허브는 .pb 형태의 weight 파일을 제공하고 있습니다.
하지만 이 파일을 Keras의 load_model 로 읽으면 에러가 나는데, load_model은 .h5 형태의 파일만 지원하기 때문입니다. (공식 문서 참고)
즉, Keras를 사용하기 위해서는 .h5 형태의 weight 파일이 필요합니다. 가장 좋은 방법은 1번에서 새로 작성한 모델 구조로 다시 학습을 시켜 .h5 파일을 만드는 것인데, 시간도 자원도 부족합니다,,, 기브미GPU,,,
그래서 최대한 .pb 파일을 사용해보기로 했습니다!
전체적인 과정은 다음과 같습니다.
- Keras로 작성한 모델을 저장 후 불러옴
- Keras로 작성한 모델 구조 (초기화 시 랜덤으로 들어간 weight 값을 포함)에서 변수명과 weight 값을 가져옴
.pb파일을 읽어 그래프 구조를 파싱- 그래프의 노드 이름과 Keras 모델의 레이어 이름을 짝으로 비교해 weight를 저장
- 1번에서 저장한 모델 구조에 3번에서 얻은 weight를 넣음
Keras로 작성한 모델을 저장 후 불러옵니다.
import tensorflow as tf
model.save('model_structure_gpu.h5')
model_structure = load_model('model_structure_gpu.h5', custom_objects={'tf': tf})
.pb 파일의 그래프 노드들의 이름과 비교하기 위해 Keras 모델의 레이어 이름을 저장합니다.
또한 그래프 노드에서 weight값을 가져와 저장하기 위해 랜덤으로 초기화된 weight값들도 저장합니다.
target_model_variable = model_structure.weights # 모델의 레이어 이름
target_model_weights = model_structure.get_weights() # 모델의 weight들
.pb파일을 읽어 그래프들을 노드단위로 파싱합니다.
import tensorflow as tf
from tensorflow.python.platform import gfile
GRAPH_PB_PATH = './model_salicon_gpu.pb'
with tf.Session() as sess:
with gfile.FastGFile(GRAPH_PB_PATH,'rb') as f:
graph_def = tf.GraphDef()
graph_def.ParseFromString(f.read())
sess.graph.as_default()
tf.import_graph_def(graph_def, name='')
graph_nodes=[n for n in graph_def.node]
그리고 노드 중 속성이 상수인 노드 (= weight) 만 불러옵니다.
wts = [n for n in graph_nodes if n.op=='Const'] # weight만 불러옴
weight 노드들의 이름과 Keras 모델의 레이어 이름들을 비교하면서, 이름이 같으면 초기화된 weight 값이 있던 부분에 weight 노드의 weight 값을 넣어줍니다.
from tensorflow.python.framework import tensor_util
# 노드명이 일치하면 weight를 바꿔줌
for n in wts:
for idx, var in enumerate(target_model_variable):
if str(n.name) == str(var.name)[:-2]:
target_model_weights[idx] = tensor_util.MakeNdarray(n.attr['value'].tensor)
[:-2] 로 슬라이싱 해 비교한 이유는 여기서 .pb 파일의 노드 이름은 aspp/conv1_2/kernel 형식이고, Keras 레이어 이름은 aspp/conv1_2/kernel:0 형식이기 때문입니다.
이제 다했습니다! 마지막으로 set_weights 함수를 이용해 방금 만든 weight를 처음 모델에 적용해주면 됩니다!
model_structure.set_weights(target_model_weights)
이제 테스트를 위해 모델을 저장하고 불러오기만 하면 됩니다!
import tensorflow as tf
model_structure.save('model_weights_gpu.h5')
model_gpu = load_model('model_structure_gpu.h5', custom_objects={'tf': tf})
model_gpu.load_weights('model_weights_gpu.h5')
여기서 custom_objects 옵션으로 Lambda 레이어에서 사용한 라이브러리들을 작성해야만 정상적으로 모델을 불러올 수 있습니다.
테스트 결과 .pb 파일로 얻은 결과물과 동일한 결과물을 얻을 수 있었습니다!
약간 미련한 짓으로 보일수도 있겠지만,,, 다시 학습할 시간과 자원을 아꼈습니다,,! 시간은 금이니까요~ 다들 아껴쓰시길 바라며 포스팅 마무리하겠습니다!