-
Notifications
You must be signed in to change notification settings - Fork 0
/
train.py
154 lines (124 loc) · 5.35 KB
/
train.py
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
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
import matplotlib.pyplot as plt
import numpy as np
import tensorflow as tf
import datetime
import argparse
from model import CVAE
train_size = 60000
test_size = 10000
# num_examples_to_generate = 16
current_time = datetime.datetime.now().strftime("%Y%m%d-%H%M%S")
train_log_dir = 'logs/gradient_tape/' + current_time + '/train'
test_log_dir = 'logs/gradient_tape/' + current_time + '/test'
train_summary_writer = tf.summary.create_file_writer(train_log_dir)
test_summary_writer = tf.summary.create_file_writer(test_log_dir)
def parse_args():
desc = 'Tensorflow implementation of CVAE'
parser = argparse.ArgumentParser(description=desc)
parser.add_argument('--EPOCHS', type=int, default=10, help='The number of epochs to run')
parser.add_argument('--batch_size', type=int, default=64, help='The size of batch')
parser.add_argument('--latent_dim', type=int, default=2, help='Dimension of latent space')
parser.add_argument('--inter_dim', type=int, default=128, help='Dimension of dense layer')
return check_args(parser.parse_args())
def check_args(args):
# --epoch
try:
assert args.EPOCHS >= 1
except:
print('number of epochs must be larger than or equal to one')
# --batch_size
try:
assert args.batch_size >= 1
except:
print('batch size must be larger than or equal to one')
# --latent_dim
try:
assert args.latent_dim >= 1
except:
print('dimension of latent space must be larger than or equal to one')
# --inter_dim
try:
assert args.inter_dim >= 1
except:
print('dimension of latent space must be larger than or equal to one')
return args
# model each pixel with a Bernoulli distribution
def preprocess_images(images):
images = images.reshape((images.shape[0], 28, 28, 1)) / 255.
return images.astype('float32')
def log_normal_pdf(sample, mean, logvar, raxis=1):
log2pi = tf.math.log(2. * np.pi)
return tf.reduce_sum(
-.5 * ((sample - mean) ** 2. * tf.exp(-logvar) + logvar + log2pi),
axis=raxis)
def compute_loss(model, x, y):
mean, logvar = model.encode(x,y)
z = model.reparameterize(mean, logvar)
x_logit = model.decode(z,y)
cross_ent = tf.nn.sigmoid_cross_entropy_with_logits(logits=x_logit, labels=x)
logpx_z = -tf.reduce_sum(cross_ent, axis=[1, 2, 3])
logpz = log_normal_pdf(z, 0., 0.)
logqz_x = log_normal_pdf(z, mean, logvar)
return -tf.reduce_mean(logpx_z + logpz - logqz_x)
train_loss = tf.keras.metrics.Mean('train_loss', dtype=tf.float32)
test_loss = tf.keras.metrics.Mean('test_loss', dtype=tf.float32)
@tf.function
def train_step(model, x, y, optimizer):
with tf.GradientTape() as tape:
loss = compute_loss(model, x, y)
gradients = tape.gradient(loss, model.trainable_variables)
optimizer.apply_gradients(zip(gradients, model.trainable_variables))
train_loss(loss)
@tf.function
def test_step(model, x_test, y_test):
test_loss(compute_loss(model, x_test, y_test))
def generate_and_save_images(model, epoch, test_sample_x, test_sample_y):
mean, logvar = model.encode(test_sample_x, test_sample_y)
z = model.reparameterize(mean, logvar)
predictions = model.sample(z, test_sample_y)
fig = plt.figure(figsize=(4, 4))
for i in range(predictions.shape[0]):
plt.subplot(4, 4, i + 1)
plt.imshow(predictions[i, :, :, 0], cmap='gray')
plt.axis('off')
plt.savefig('images/image_at_epoch_{:04d}.png'.format(epoch))
# plt.show()
if __name__ == '__main__':
# parse arguments
args = parse_args()
if args is None:
exit()
(train_images, train_y), (test_images, test_y) = tf.keras.datasets.mnist.load_data()
train_y = tf.one_hot(train_y, depth=10)
test_y = tf.one_hot(test_y, depth=10)
train_images = preprocess_images(train_images)
test_images = preprocess_images(test_images)
train_dataset = (tf.data.Dataset.from_tensor_slices((train_images, train_y))
.shuffle(train_size).batch(args.batch_size))
test_dataset = (tf.data.Dataset.from_tensor_slices((test_images, test_y))
.shuffle(test_size).batch(args.batch_size))
optimizer = tf.keras.optimizers.Adam(1e-4)
# # test
# model = CVAE(latent_dim, inter_dim)
# for test_batch in test_dataset.take(1):
# test_sample_x = test_batch[0][0:num_examples_to_generate, :, :, :]
# test_sample_y = test_batch[1][0:num_examples_to_generate]
# generate_and_save_images(model, 0, test_sample_x, test_sample_y)
# reset
model = CVAE(args.latent_dim, args.inter_dim)
for epoch in range(args.EPOCHS):
for train in train_dataset:
train_step(model, train[0], train[1], optimizer)
with train_summary_writer.as_default():
tf.summary.scalar('loss', train_loss.result(), step=epoch)
for test in test_dataset:
test_step(model, test[0], test[1])
with test_summary_writer.as_default():
tf.summary.scalar('loss', test_loss.result(), step=epoch)
template = 'Epoch {}, Loss: {}, Test Loss: {}'
print(template.format(epoch + 1, train_loss.result(), test_loss.result()))
# Reset metrics every epoch
train_loss.reset_states()
test_loss.reset_states()
# generate_and_save_images(model, epoch, test_sample_x, test_sample_y)
model.save_weights('saved_model/my_checkpoint')