본문 바로가기
Developer/Machine Learning

[Tensorflow] 학습 모델 저장하고 불러오는 방법 (Saver, Restore)

by Doony 2020. 1. 13.

저는 주로 Jupyter notebook으로 작업을 합니다. 때문에 창을 끄거나 컴퓨터를 끄는 등 종료를 시키게 되면, 한번 학습했던 모델이 다 사라지게 되는데요. 우리가 만든 모델을 실생활에 사용하기 위해서는 학습 모델을 저장하고, 또 불러와서 사용할 수 있어야합니다.
MNIST 데이터셋에 대해 다룬 포스팅을 기준으로, 학습 모델을 저장하고 불러오는 방법에 대해 알아보겠습니다.

학습 모델 저장하기

위 포스팅에서 모델 중간에 아래 코드가 있습니다. 텐서플로우의 Saver을 통해 모델을 저장할 수 있습니다.


1
2
3
4
5
6
7
# 모델 저장을 위한 부분
import os
save_file = './model_mnist.ckpt'
SAVER_DIR = "modelMNIST"
saver = tf.train.Saver()
checkpoint_path = os.path.join(SAVER_DIR, "modelMNIST")
ckpt = tf.train.get_checkpoint_state(SAVER_DIR)
cs

위 코드를 학습 이전에 실행합니다. 그 후, 학습을 할 때 아래와 같이 지정하고 싶은 주기별로 체크포인트를 생성합니다.


1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
for epoch in range(training_epochs):
    avg_cost = 0
    avg_cost_test = 0
    total_batch = int(mnist.train.num_examples / batch_size)
    
    for i in range(total_batch):
        batch_xs, batch_ys = mnist.train.next_batch(bach_size)
        feed_dict = {X: batch_xs, Y: batch_ys}
        summary, c, a, _ = sess.run([merged, cost, accuracy, optimizer], feed_dict = feed_dict)
        avg_cost += c / total_batch
        
    
    # Tensorboard에서 epoch별 스칼라값 확인하기 위함
    writer.add_summary(summary, global_step = epoch)
    
    # epoch별 모델 체크포인트 저장
    saver.save(sess, checkpoint_path, global_step = epoch)
 
cs

저는 epoch 마다 저장하게 해놨습니다.


학습된 모델 불러오기

모델을 불러오기 위해서는 학습 모델의 기본 구조는 가지고 있어야합니다.
일단 기존 세션과 관계 없는 새로운 노트를 만든 후, 아래와 같이 기본 골격은 동일하게 설정합니다.


1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
import tensorflow as tf
 
from tensorflow.examples.tutorials.mnist import input_data
mnist = input_data.read_data_sets("/tmp/data/", one_hot = True)
tf.set_random_seed(777)
tf.reset_dafault_graph()
 
learning_rate = 0.001
training_epochs = 10
batch_size = 100
 
= tf.placeholder(tf.float32, [None, 784])
X_img = tf.reshape(X, [-128281])
= tf.placeholder(tf.float32, [None, 10])
 
 
W1 = tf.Variable(tf.random_normal([33132], stddev = 0.01))
L1 = tf.nn.conv2d(X_img, W1, strides = [1111], padding = 'SAME')
L1 = tf.nn.relu(L1)
L1 = tf.nn.max_pool(L1, ksize = [1221], strides = [1221], padding = 'SAME')
 
 
W2 = tf.Variable(tf.random_normal([333264], stddev = 0.01))
L2 = tf.nn.conv2d(L1, W2, strides = [1111], padding = 'SAME')
L2 = tf.nn.relu(L2)
L2 = tf.nn.max_pool(L2, ksize = [1221], strides = [1221], padding = 'SAME')
L2_flat = tf.reshape(L2, [-17*7*64])
 
W3 = tf.Variable(tf.random_normal([7*7*6410], stddev = 0.01))
= tf.Variable(tf.random_normal([10]))
logits = tf.matmul(L2_flat, W3) + b
 
cost = tf.reduce_mean(tf.nn.softmax_cross_entropy_with_logits(logits = logits, labels = Y))
optimizer = tf.train.AdamOptimizer(learning_rate = learning_rate).minimize(cost)
correct_prediction = tf.equal(tf.argmax(logits, 1), tf.argmax(Y, 1))
accuracy = tf.reduce_mean(tf.cast(correct_prediction, tf.float32))
cs

모델에서 정의된 weight나 bias 등 계층 구조가 있어야 그 값을 불러올 수 있기 때문에 위와 같은 과정이 필요합니다. 그 후, 아래 코드를 통해 테스트 데이터셋을 검증해보겠습니다.


1
2
3
4
5
6
7
8
with tf.Session() as sess:
    sess.run(tf.glbal_variables_initializer())
    saver = tf.train.import_meta_graph('./model_mnist.ckpt.meta')
    saver.restore(sess, tf.train.latest_checkpoint('./'))
    
    test_accuracy = sess.run(accuracy, feed_dict={X: mnist.test.images, Y: mnist.test.labels})
    
    print(test_accuracy)
cs

약 98%의 정확도를 가지는 것으로 결과가 확인됩니다. 즉, 모델이 잘 불러와진 것을 알 수 있습니다.

댓글