-
Notifications
You must be signed in to change notification settings - Fork 10
/
train.py
executable file
·196 lines (161 loc) · 6.26 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
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
#!/usr/bin/python3
# -*- coding: utf-8 -*-
"""tensorflow/datasets
Automatically generated by Colaboratory.
Original file is located at
https://colab.research.google.com/github/tensorflow/datasets/blob/master/docs/keras_example.ipynb
# Training a neural network on MNIST with Keras
This simple example demonstrate how to plug TFDS into a Keras model.
Copyright 2020 The TensorFlow Datasets Authors, Licensed under the Apache License, Version 2.0
<table class="tfo-notebook-buttons" align="left">
<td>
<a target="_blank" href="https://www.tensorflow.org/datasets/keras_example"><img src="https://www.tensorflow.org/images/tf_logo_32px.png" />View on TensorFlow.org</a>
</td>
<td>
<a target="_blank" href="https://colab.research.google.com/github/tensorflow/datasets/blob/master/docs/keras_example.ipynb"><img src="https://www.tensorflow.org/images/colab_logo_32px.png" />Run in Google Colab</a>
</td>
<td>
<a target="_blank" href="https://github.com/tensorflow/datasets/blob/master/docs/keras_example.ipynb"><img src="https://www.tensorflow.org/images/GitHub-Mark-32px.png" />View source on GitHub</a>
</td>
</table>
"""
import tensorflow.compat.v2 as tf
import tensorflow_datasets as tfds
import sys
import re
tfds.disable_progress_bar()
tf.enable_v2_behavior()
"""## Step 1: Create your input pipeline
Build efficient input pipeline using advices from:
* [TFDS performance guide](https://www.tensorflow.org/datasets/performances)
* [tf.data performance guide](https://www.tensorflow.org/guide/data_performance#optimize_performance)
### Load MNIST
Load with the following arguments:
* `shuffle_files`: The MNIST data is only stored in a single file, but for larger datasets with multiple files on disk, it's good practice to shuffle them when training.
* `as_supervised`: Returns tuple `(img, label)` instead of dict `{'image': img, 'label': label}`
"""
dataset = 'fashion_mnist' if 'fashion' in sys.argv else 'mnist'
(ds_train, ds_test), ds_info = tfds.load(
dataset,
split=['train', 'test'],
shuffle_files=True,
as_supervised=True,
with_info=True,
)
"""### Build training pipeline
Apply the following transormations:
* `ds.map`: TFDS provide the images as tf.uint8, while the model expect tf.float32, so normalize images
* `ds.cache` As the dataset fit in memory, cache before shuffling for better performance.<br/>
__Note:__ Random transformations should be applied after caching
* `ds.shuffle`: For true randomness, set the shuffle buffer to the full dataset size.<br/>
__Note:__ For bigger datasets which do not fit in memory, a standard value is 1000 if your system allows it.
* `ds.batch`: Batch after shuffling to get unique batches at each epoch.
* `ds.prefetch`: Good practice to end the pipeline by prefetching [for performances](https://www.tensorflow.org/guide/data_performance#prefetching).
"""
def normalize_img(image, label):
"""Normalizes images: `uint8` -> `float32`."""
return tf.cast(image, tf.float32) / 255., label
ds_train = ds_train.map(
normalize_img, num_parallel_calls=tf.data.experimental.AUTOTUNE)
ds_train = ds_train.cache()
ds_train = ds_train.shuffle(ds_info.splits['train'].num_examples)
ds_train = ds_train.batch(128)
ds_train = ds_train.prefetch(tf.data.experimental.AUTOTUNE)
"""### Build evaluation pipeline
Testing pipeline is similar to the training pipeline, with small differences:
* No `ds.shuffle()` call
* Caching is done after batching (as batches can be the same between epoch)
"""
ds_test = ds_test.map(
normalize_img, num_parallel_calls=tf.data.experimental.AUTOTUNE)
ds_test = ds_test.batch(128)
ds_test = ds_test.cache()
ds_test = ds_test.prefetch(tf.data.experimental.AUTOTUNE)
"""## Step 2: Create and train the model
Plug the input pipeline into Keras.
"""
network = 'A'
n_epochs = 10
lr = 0.1
adam = False
amsgrad = False
if len(sys.argv) > 1:
network = sys.argv[1]
if len(sys.argv) > 2:
n_epochs = int(sys.argv[2])
if len(sys.argv) > 3:
adam = re.match('adam(.*)', sys.argv[3])
amsgrad = re.match('amsgrad(.*)', sys.argv[3])
if adam or amsgrad:
try:
lr = float((adam or amsgrad).group(1))
except:
lr = .001
if adam:
print('use Adam with lr', lr)
else:
print('use AMSGrad with lr', lr)
else:
lr = float(sys.argv[3])
if network == 'A':
layers = [
tf.keras.layers.Flatten(),
tf.keras.layers.Dense(128,activation='relu'),
tf.keras.layers.Dense(128,activation='relu'),
tf.keras.layers.Dense(10, activation='softmax')
]
elif network == 'D':
layers = [
tf.keras.layers.Conv2D(5, 5, 2, 'same', activation='relu'),
tf.keras.layers.Flatten(),
tf.keras.layers.Dense(100, activation='relu'),
tf.keras.layers.Dense(10, activation='softmax')
]
model = tf.keras.models.Sequential(layers)
elif network == 'B':
layers = [
tf.keras.layers.Conv2D(16, 5, 1, 'valid', activation='relu'),
tf.keras.layers.MaxPooling2D(2),
tf.keras.layers.Conv2D(16, 5, 1, 'valid', activation='relu'),
tf.keras.layers.MaxPooling2D(2),
tf.keras.layers.Flatten(),
tf.keras.layers.Dense(100, activation='relu'),
tf.keras.layers.Dense(10, activation='softmax')
]
elif network == 'C':
layers = [
tf.keras.layers.Conv2D(20, 5, 1, 'valid', activation='relu'),
tf.keras.layers.MaxPooling2D(2),
tf.keras.layers.Conv2D(50, 5, 1, 'valid', activation='relu'),
tf.keras.layers.MaxPooling2D(2),
tf.keras.layers.Flatten(),
tf.keras.layers.Dense(500, activation='relu'),
tf.keras.layers.Dense(10, activation='softmax')
]
if 'dropout' in sys.argv:
layers.insert(-2, tf.keras.layers.Dropout(0.5))
else:
raise Exception('unknown network: ' + network)
print (layers)
model = tf.keras.models.Sequential(layers)
if adam or amsgrad:
optim = tf.keras.optimizers.Adam(lr, amsgrad=amsgrad)
else:
optim = tf.keras.optimizers.SGD(momentum=0.9, learning_rate=lr,)
model.compile(
loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),
metrics=[tf.keras.metrics.SparseCategoricalAccuracy()],
optimizer=optim
)
for i in range(1):
history = model.fit(
ds_train,
epochs=n_epochs,
validation_data=ds_test,
)
out = open('log-' + '-'.join(sys.argv[1:]), 'w')
for i in range(n_epochs):
out.write(str(i) + ' ')
for x in history.history:
out.write('%s: %.4f ' % (x, history.history[x][i]))
out.write('\n')