forked from lsds/KungFu
-
Notifications
You must be signed in to change notification settings - Fork 0
/
tf2_mnist_gradient_tape.py
75 lines (60 loc) · 2.8 KB
/
tf2_mnist_gradient_tape.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
import argparse
import tensorflow as tf
from kungfu.python import current_cluster_size, current_rank
from kungfu.tensorflow.optimizers import (PairAveragingOptimizer,
SynchronousAveragingOptimizer,
SynchronousSGDOptimizer)
parser = argparse.ArgumentParser(description='KungFu mnist example.')
parser.add_argument('--kf-optimizer',
type=str,
default='sync-sgd',
help='available options: sync-sgd, async-sgd, sma')
args = parser.parse_args()
(mnist_images, mnist_labels), _ = \
tf.keras.datasets.mnist.load_data(path='mnist-%d.npz' % current_rank())
dataset = tf.data.Dataset.from_tensor_slices(
(tf.cast(mnist_images[..., tf.newaxis] / 255.0,
tf.float32), tf.cast(mnist_labels, tf.int64)))
dataset = dataset.repeat().shuffle(10000).batch(128)
mnist_model = tf.keras.Sequential([
tf.keras.layers.Conv2D(32, [3, 3], activation='relu'),
tf.keras.layers.Conv2D(64, [3, 3], activation='relu'),
tf.keras.layers.MaxPooling2D(pool_size=(2, 2)),
tf.keras.layers.Dropout(0.25),
tf.keras.layers.Flatten(),
tf.keras.layers.Dense(128, activation='relu'),
tf.keras.layers.Dropout(0.5),
tf.keras.layers.Dense(10, activation='softmax')
])
loss = tf.losses.SparseCategoricalCrossentropy()
# KungFu: adjust learning rate based on number of GPUs.
# opt = tf.keras.optimizers.SGD(0.001 * current_cluster_size())
opt = tf.compat.v1.train.AdamOptimizer(0.001 * current_cluster_size())
# KungFu: wrap tf.compat.v1.train.Optimizer.
if args.kf_optimizer == 'sync-sgd':
opt = SynchronousSGDOptimizer(opt)
elif args.kf_optimizer == 'async-sgd':
opt = PairAveragingOptimizer(opt)
elif args.kf_optimizer == 'sma':
opt = SynchronousAveragingOptimizer(opt)
else:
raise RuntimeError('Unknown KungFu optimizer')
@tf.function
def training_step(images, labels, first_batch):
with tf.GradientTape() as tape:
probs = mnist_model(images, training=True)
loss_value = loss(labels, probs)
grads = tape.gradient(loss_value, mnist_model.trainable_variables)
opt.apply_gradients(zip(grads, mnist_model.trainable_variables))
# KungFu: broadcast is done after the first gradient step to ensure optimizer initialization.
if first_batch:
from kungfu.tensorflow.initializer import broadcast_variables
broadcast_variables(mnist_model.variables)
broadcast_variables(opt.variables())
return loss_value
# KungFu: adjust number of steps based on number of GPUs.
for batch, (images, labels) in enumerate(
dataset.take(10000 // current_cluster_size())):
loss_value = training_step(images, labels, batch == 0)
if batch % 10 == 0 and current_rank() == 0:
print('Step #%d\tLoss: %.6f' % (batch, loss_value))