tf2-AutoEncoder/mnist_autoencoder.py at master · ZeroAct/tf2-AutoEncoder · GitHub

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
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
### 필요한 모듈 임포트
from tensorflow.keras.layers import Conv2D, Conv2DTranspose, Dense, Flatten, Dropout, BatchNormalization, Reshape, LeakyReLU
from tensorflow.keras.models import Model, Sequential
from tensorflow.keras.callbacks import ModelCheckpoint, EarlyStopping

from tensorflow.keras import Input

import numpy as np

import tensorflow as tf
import os

### 데이터 로드
mnist = tf.keras.datasets.mnist
(x_train, y_train), (x_valid, y_valid) = mnist.load_data()

### 데이터 프로세싱
x_train = x_train.reshape(-1, 28, 28, 1)
x_train = x_train / 127.5 - 1
x_valid = x_valid.reshape(-1, 28, 28, 1)
x_valid = x_valid / 127.5 - 1

### Encoder 정의
encoder_input = Input(shape=(28,28,1))

x = Conv2D(32, 3, strides=1, padding='same')(encoder_input)
x = BatchNormalization()(x)
x = LeakyReLU()(x)

x = Conv2D(64, 3, strides=2, padding='same')(x)
x = BatchNormalization()(x)
x = LeakyReLU()(x)

x = Conv2D(64, 3, strides=2, padding='same')(x)
x = BatchNormalization()(x)
x = LeakyReLU()(x)

x = Conv2D(64, 3, padding='same')(x)
x = BatchNormalization()(x)
x = LeakyReLU()(x)

x = Flatten()(x)

encoder_output = Dense(3)(x)

encoder = Model(encoder_input, encoder_output)

### Decoder 정의
decoder_input = Input(shape=(3,))

x = Dense(7*7*64)(decoder_input)
x = Reshape((7, 7, 64))(x)

x = Conv2DTranspose(64, 3, strides=1, padding='same')(x)
x = BatchNormalization()(x)
x = LeakyReLU()(x)

x = Conv2DTranspose(64, 3, strides=2, padding='same')(x)
x = BatchNormalization()(x)
x = LeakyReLU()(x)

x = Conv2DTranspose(64, 3, strides=2, padding='same')(x)
x = BatchNormalization()(x)
x = LeakyReLU()(x)

x = Conv2DTranspose(32, 3, strides=1, padding='same')(x)
x = BatchNormalization()(x)
x = LeakyReLU()(x)

decoder_output = Conv2DTranspose(1, 3, strides=1, padding='same', activation='tanh')(x)

decoder = Model(decoder_input, decoder_output)

### Hyperparameter 정의
LEARNING_RATE = 0.0005
BATCH_SIZE = 32

### AutoEncoder 정의
encoder_in = Input(shape=(28,28,1))
x = encoder(encoder_in)
decoder_out = decoder(x)

auto_encoder = Model(encoder_in, decoder_out)

### 학습
auto_encoder.compile(optimizer=tf.keras.optimizers.Adam(LEARNING_RATE),
                     loss=tf.keras.losses.MeanSquaredError())

model_dir = './mnist_plain/'
os.mkdir(model_dir)
model_path = model_dir + '{epoch}__loss_{loss}.h5'

checkpoint_callback = ModelCheckpoint(model_path,
                                      save_best_only=True,
                                      save_weights_only=False,
                                      monitor='loss',
                                      verbose=1)
earlystop_callback = EarlyStopping(monitor='loss',
                                   patience=2)

auto_encoder.fit(x_train, x_train,
                 batch_size=BATCH_SIZE,
                 epochs=100,
                 callbacks=[checkpoint_callback, earlystop_callback])

import matplotlib.pyplot as plt

decoded_images = auto_encoder.predict(x_train[:15])

fig, axes = plt.subplots(3, 5)
fig.set_size_inches(12, 6)
for i in range(15):
    t = x_train[i]
    t = (t-t.min())/(t.max()-t.min())*255
    axes[i//5, i%5].imshow(t.astype(np.uint8))
    axes[i//5, i%5].axis('off')
plt.tight_layout()
plt.title('Original Images')
plt.show()

fig, axes = plt.subplots(3, 5)
fig.set_size_inches(12, 6)
for i in range(15):
    t = decoded_images[i]
    t = (t-t.min())/(t.max()-t.min())*255
    axes[i//5, i%5].imshow(t.astype(np.uint8))
    axes[i//5, i%5].axis('off')
plt.tight_layout()
plt.title('Auto Encoder Images')
plt.show()