Sign in Subscribe
tensorflow

TF 1에서 TF 2로 porting하기

장동주

• 11 min read

인턴 기간 동안 진행했던 CycleGAN code porting을 진행하면서 배웠던 버전 간의 차이점과 tensorflow 2를 사용하면서 썼던 api들에 대해 설명하려고 합니다. 기존의 코드는 tensorflow 1.8 버전으로 작성되었고 프로젝트를 진행했던 시점에서 최신 버전인 2.2버전으로 수정하였습니다.

코드를 수정하면서 기능적으로 달라진 부분은 다음과 같습니다.

  • Red CNN, WGAN VGG, Cyclegan, Cycle Identity 4종류의 모델 중 하나를 선택하는 방식에서 Cycle Identity 모델만 사용하도록 수정하였습니다.
  • 데이터 형태에 따라 unpaired 또는 paired를 선택했던 방식에서 unpaired한 데이터만을 다루는 방식으로 수정하였습니다.
  • end_epoch과 decay_epoch을 이용하여 learning rate를 조절하는 방식을 사용하지 않고 Adam optimizer를 그대로 사용하도록 수정하였습니다.

이를 제외하고는 기존과 동일한 작업을 수행하도록 수정하였습니다.

Data Input Pipeline

모델을 학습하는 과정에서 데이터를 적절하게 전처리하고 모델에 입력할 수 있게 하는 작업은 어려운 작업입니다. 데이터의 크기가 커질수록 데이터 입력 파이프라인을 잘 만드는 것이 중요해집니다.

tensorflow에서는 데이터를 미리 전처리하여 TFRecord 형태로 변환하고 학습을 할 때 tfrecorddataset을 사용하는 방법을 가장 권장하는 것 같습니다. 데이터의 크기가 크고 실시간으로 쌓일수록 더 알맞는 방법입니다. 하지만 이번에는 기존의 코드처럼 전처리를 포함하도록 했습니다.

기존의 코드에서는 tf.FIFOqueue를 이용하여 병렬적으로 작동하는 파이프라인을 구현했습니다.

  1. glob을 통해 지정한 형식의 파일명을 모두 찾습니다.
  2. 찾은 파일명을 이용해 데이터를 모두 read합니다.
  3. 모델을 학습이 진행되는 과정에서 데이터를 전처리하고 tf.FIFOqueue에 enqueue합니다.
  4. tf.FIFOqueue에서 batch size만큼 dequeue해서 모델에 입력합니다.

하지만 이 방식은 두 가지 문제가 있습니다.

  • 데이터 전체를 한꺼번에 메모리에 할당하기 때문에 데이터가 memory의 크기 보다 큰 경우는 대처할 수 없습니다.
  • 전체 데이터를 load 한 뒤 학습을 시작하기 때문에 모델에  입력 데이터가 들어가기까지 시간이 오래걸립니다.

tf.data

tf.data API는 기존의 코드에서 사용한 tf.FIFOqueue와 방식이 다르기 때문에 위의 두 가지 문제점을 해결할 수 있습니다.

tensorflow 2에서는 tf.data API를 이용하여 data input pipeline을 만드는 것을 권장합니다. tf.data API를 통해 dataset 객체를 만들어 손쉽게 전처리 과정과 파이프라인을 구성할 수 있습니다.

제가 이해한 수준에서 tf.data는 저장소와 메모리 사이에 통로를 놓아준다고 볼 수 있습니다. 전체 데이터를 메모리에 올려놓고 시작하는 것이 아닌 통로를 통해 데이터를 저장소에서 하나씩 가져오면서 먼저 가져온 데이터부터 바로 사용할 수 있는 방식이라 할 수 있습니다. 그래서 tf.data.Dataset을 출력해보면 하나의 Dataset 객체로 표시가 될 뿐이고 for문 같은 iterator를 이용해야 실제 데이터를 확인할 수 있습니다.

AWS sagemaker에서 학습을 하는 경우에도 tf.data API를 사용하여 pipe line을 구성해두면 'Pipe' input mode를 오류없이 이용할 수 있어 학습을 빠르게 진행할 수 있습니다.

def dcm_read(path):
    path = path.numpy().decode('utf-8')
    dcm_file = pydicom.dcmread(path)
    img = get_pixel_hu(dcm_file)
    return img


def read_function_dcm(fn):
    out = tf.py_function(dcm_read, [fn], tf.int16)
    return out
    
def get_image_dataset(patent_no_list):
    path_pattern_list = [os.path.join(self.data_path, patent_no, '*.' + self.extension) for patent_no in patent_no_list]
    p_path = tf.data.Dataset.list_files(path_pattern_list)

    p = p_path.map(read_function_dcm, num_parallel_calls=tf.data.experimental.AUTOTUNE)

    p = p.map(normalize, num_parallel_calls=tf.data.experimental.AUTOTUNE)

    return p
dcm 파일을 읽어 dataset을 만드는 과정 예시
  • tf.data.list_file : 지정한 형식의 이름을 가진 파일의 이름을 읽어오는 파이프라인을 만듭니다.
  • tf.data.Dataset.map : element 별로 지정한 작업을 진행할 수 있습니다.

list_file을 이용해 glob과 같은 기능을 할 수 있고, map을 통해 read 및 전처리를 수행합니다.

map 함수에서는 처리할 작업으로 python function을 직접 지정할 수 없고 wrapping 해줘야 합니다. 위의 코드에서 dicom 파일을 읽기 위해서 pydicom 라이브러리의 함수를 사용했고 tf.py_funtion을 통해 wrapping을 해주었습니다.‌

모델 구성

tensorflow 2에서는 모델을 구성할 때 tf.keras API를 사용하도록 권장합니다. Keras의 high level api를 통해 직관적으로 모델을 구성할 수 있습니다.

원래 Keras는 다양한 framework를 back-end로 지원했지만 앞으로는 tensorflow만 지원을 하는 방향으로 개발이 진행됩니다. 또한 tensorflow 라이브러리에 포함되어 모델 구성 부분을 tf.keras API로 대부분 처리할 수 있게 발전하고 있습니다. 이 글에서 tensorflow와 keras의 관계를 잘 설명하고 있습니다.

tf.keras로 model을 만드는 방식은 총 3가지가 있습니다.

  1. Sequential API
  2. functional API
  3. model class 상속

그 중에서 functional API를 이용하여 모델을 만들었습니다.

functional API를 사용하는 방식은 기존의 코드에서 graph 구조를 정의하는 부분과 유사합니다. 

import tensorflow as tf

def discriminator(image, options, reuse=False, name='discriminator'):
    def first_layer(input_, out_channels, ks=3, s=1, name='disc_conv_layer'):
    	with tf.variable_scope(name):
    		return lrelu(conv2d(input_, out_channels, ks=ks, s=s))
    def conv_layer(input_, out_channels, ks=3, s=1, name='disc_conv_layer'):
    	with tf.variable_scope(name):
    		return lrelu(batchnorm(conv2d(input_, out_channels, ks=ks, s=s)))
    def last_layer(input_, out_channels, ks=4, s=1, name='disc_conv_layer'):
    	with tf.variable_scope(name):
    		return tf.contrib.layers.fully_connected(conv2d(input_, out_channels, ks=ks, s=s), out_channels)

    with tf.variable_scope(name):
    	if reuse:
    		tf.get_variable_scope().reuse_variables()
    	else:
    		assert tf.get_variable_scope().reuse is False
    l1 = first_layer(image, options.df_dim, ks=4, s=2, name='disc_layer1')
    l2 = conv_layer(l1, options.df_dim*2, ks=4, s=2, name='disc_layer2')
    l3 = conv_layer(l2, options.df_dim*4, ks=4, s=2, name='disc_layer3')
    l4 = conv_layer(l3, options.df_dim*8, ks=4, s=1, name='disc_layer4')
    l5 = last_layer(l4, options.img_channel, ks=4, s=1, name='disc_layer5')
    return l5


def lrelu(x, leak=0.2, name="lrelu"):
    with tf.variable_scope(name):
        return tf.maximum(x, leak*x)

def batchnorm(input_, name="batch_norm"):
    with tf.variable_scope(name):
        return tf.layers.batch_normalization(input_, axis=3, epsilon=1e-5, \
            momentum=0.1, training=True, \
            gamma_initializer=tf.random_normal_initializer(1.0, 0.02))

def conv2d(batch_input, out_channels, ks=4, s=2, name="cov2d"):
    with tf.variable_scope(name):
        padded_input = tf.pad(batch_input, [[0, 0], [1, 1], [1, 1], [0, 0]], mode="CONSTANT")
        return tf.layers.conv2d(padded_input, out_channels, kernel_size=ks, \
            strides=s, padding="valid", \
            kernel_initializer=tf.random_normal_initializer(0, 0.02))
tf1
import tensorflow as tf
from tensorflow.keras import layers

def discriminator(image_shape, options, name='discriminator'):
    def first_layer(input_, out_channels, ks=3, s=1):
    	return lrelu(conv2d(input_, out_channels, ks=ks, s=s))

    def conv_layer(input_, out_channels, ks=3, s=1):
    	return lrelu(batchnorm(conv2d(input_, out_channels, ks=ks, s=s)))

    def last_layer(input_, out_channels, ks=4, s=1):
    	return layers.Dense(units=out_channels)(conv2d(input_, out_channels, ks=ks, s=s))

    inputs = tf.keras.Input(shape=image_shape)
    l1 = first_layer(inputs, options.df_dim, ks=4, s=2)
    l2 = conv_layer(l1, options.df_dim * 2, ks=4, s=2)
    l3 = conv_layer(l2, options.df_dim * 4, ks=4, s=2)
    l4 = conv_layer(l3, options.df_dim * 8, ks=4, s=1)
    l5 = last_layer(l4, options.img_channel, ks=4, s=1)

    model = tf.keras.Model(inputs=inputs, outputs=l5, name=name)
    return model

def lrelu(x, leak=0.2):
    return layers.LeakyReLU(alpha=leak)(x)


def batchnorm(input_, name="batch_norm"):
    return layers.BatchNormalization(axis=3, epsilon=1e-5, momentum=0.1,
                                     gamma_initializer=tf.random_normal_initializer(1.0, 0.02))(input_, training=True)


def conv2d(batch_input, out_channels, ks=4, s=2):
    padded_input = tf.pad(batch_input, [[0, 0], [1, 1], [1, 1], [0, 0]], mode="CONSTANT")
    return layers.Conv2D(out_channels, kernel_size=ks, strides=s, padding="valid",
                         kernel_initializer=tf.random_normal_initializer(0, 0.02))(padded_input)
tf.keras functional api

API 이름을 바꿔주는 수준의 수정으로 tf.keras API를 사용하는 것이 가능했습니다.

이번에는 generator, discriminator 각각 2개씩 총 4개의 tf.keras.Model이 사용하도록 하였습니다.

학습 과정 설계

Graph mode

tensorflow 1에서는 Graph 모드를 사용했습니다. Graph 모드는 모델의 연산을 python interpretor와 분리 가능한 계산 그래프 형태로 만든 다음 학습을 진행합니다.

Graph 모드의 장점은 다음과 같습니다.

  • 연산 간의 관계를 파악하여 서로 독립적인 연산은 병렬 처리할 수 있습니다.
  • 연산의 중복을 막을 수 있습니다.
  • 계산 그래프를 이용하기 때문에 미분 값을 구하기 쉽습니다.
  • 여러 디바이스로 분산 처리하기 적합합니다.

다양한 이유로 graph 모드는 성능적으로는 뛰어나지만, 모델을 구성할 때 직관적이지 못해 어려움이 있습니다.

Eager mode

반면에 tensorflow 2에서는 Eager 모드가 기본입니다. Eager 모드는 python처럼 연산을 바로 수행하는 방식입니다.

eager 모드의 장점은 다음과 같습니다.

  • python 스타일로 코드를 작성할 수 있기 때문에 직관적입니다.
  • 디버깅이 쉽습니다.
  • python 제어문을 사용할 수 있어서 직관적으로 유연한 모델 구성을 할 수 있습니다.

이 글에서 두 가지 모드의 장단점을 잘 설명했습니다.

Eager 모드를 이용하기 때문에 python 코드처럼 학습 과정을 작성할 수 있습니다.

@tf.function
def train_step(patch_X, patch_Y, g_optim, d_optim, step):
	with tf.GradientTape() as tape:
            '''
            Forwarding
            (ex) G_X = self.generator_G(patch_X)
            '''
            ''' 
            Loss
            (ex) cycle_loss = md.cycle_loss(patch_X, F_GX, patch_Y, G_FY, args.L1_lambda)
            '''
	# get gradients values from tape
    # (ex) generator_g_gradients = tape.gradient(G_loss,self.generator_G.trainable_variables)
            g_optim.apply_gradients(zip(generator_g_gradients,                    self.generator_G.trainable_variables))

for epoch in range(args.epoch):
	for patch_X, patch_Y in tf.data.Dataset.zip((self.patch_X_set, self.patch_Y_set)):
		train_step(patch_X, patch_Y, g_optim, d_optim, self.ckpt.step)

tf.GradientTape

eager 모드는 계산 그래프의 구조를 미리 알 수 없기 때문에 미분 값을 계산하기 위해 tf.GradientTape()를 사용하여 연산 순서를 기록을 해야합니다. 이를 통해 gradient 함수로 원하는 weight에 해당하는 gradient를 계산할 수 있습니다. 그리고 opimizer의 apply_gradients 함수를 통해 weight 값을 변경할 수 있습니다.

@tf.fuction

eager 모드에서는 graph 모드의 성능적인 장점을 얻기 어렵습니다. 그래서 eager 모드에서는 Autograph 기능을 통해 성능에서의 단점을 극복하려고 합니다. tensorflow 2에서는 eager 모드로 작성된 코드에서도 tf.function decorator를 이용하면 모델을 graph 형태로 바꿀 수 있도록 지원합니다. 적절한 함수에 decorator만 붙여 사용하기 때문에 편리합니다.

train step

tensorflow 2에서는 estimator의 fit을 이용하여 학습하는 방법도 있지만 이번 경우에는 tf.keras.Model 객체 4개를 사용했기 때문에 custom training step을 사용했습니다. 하나의 학습 스텝에서 하는 작업을 train_step 함수로 만들어 줍니다.

train_step 함수에 @tf.funtion을 붙여 주었습니다. 한 가지 주의 할 점은 tf.function decorator가 붙은 함수에 trace하지 못하는 parameter가 들어오면 매 스탭마다 그래프를 새로 그리게 되어 성능이 저하될 수 있다는 warning이 발생합니다. 수정하는 과정에서 train_step 함수에서 python 변수를 이용하여 step 정보를 기록했을 때 warning message가 발생하였습니다. 이것을 해결하기 위해서 python 변수 대신 tf.variable을 이용하였습니다.

Sign up for more like this.

Enter your email
Subscribe