# Download and load MNIST dataset.
(x_train, y_train), (x_test, y_test) = tf.keras.datasets.mnist.load_data(
"MNIST-data-%d" % smp.rank()
)
x_train, x_test = x_train / 255.0, x_test / 255.0
# Add a channels dimension
x_train = x_train[..., tf.newaxis]
x_test = x_test[..., tf.newaxis]
# Rubik: Seed the shuffle with smp.dp_rank(), and drop_remainder
# in batching to make sure batch size is always divisible by number of microbatches
train_ds = (
tf.data.Dataset.from_tensor_slices((x_train, y_train))
.shuffle(10000, seed=smp.dp_rank())
.batch(256, drop_remainder=True)
)
test_ds = (
tf.data.Dataset.from_tensor_slices((x_test, y_test))
.shuffle(10000, seed=smp.dp_rank())
.batch(256, drop_remainder=True)
)
# Rubik: Define smp.DistributedModel the same way as Keras sub-classing API
class MyModel(smp.DistributedModel):
def __init__(self):
super(MyModel, self).__init__()
self.conv = Conv2D(32, 3, activation="relu")
self.flatten = Flatten()
else:
raise
# Download and load MNIST dataset.
(x_train, y_train), (x_test, y_test) = tf.keras.datasets.mnist.load_data(
"MNIST-data-%d" % smp.rank()
)
x_train, x_test = x_train / 255.0, x_test / 255.0
# Add a channels dimension
x_train = x_train[..., tf.newaxis]
x_test = x_test[..., tf.newaxis]
train_ds = (
tf.data.Dataset.from_tensor_slices((x_train, y_train))
.shuffle(10000, seed=123 + smp.dp_rank())
.batch(32)
)
class MyModel(smp.DistributedModel):
def __init__(self):
super(MyModel, self).__init__()
self.conv1 = Conv2D(
32, 3, activation="relu", kernel_initializer=tf.keras.initializers.GlorotNormal(seed=12)
)
self.conv0 = Conv2D(
32, 3, activation="relu", kernel_initializer=tf.keras.initializers.GlorotNormal(seed=12)
)
self.flatten = Flatten()
self.d1 = Dense(
raise
# Download and load MNIST dataset.
(x_train, y_train), (x_test, y_test) = tf.keras.datasets.mnist.load_data(
"MNIST-data-%d" % smp.rank())
x_train, x_test = x_train / 255.0, x_test / 255.0
# Add a channels dimension
x_train = x_train[..., tf.newaxis]
x_test = x_test[..., tf.newaxis]
# SMP: If needed, seed the shuffle with smp.dp_rank(), and drop_remainder
# in batching to make sure batch size is always divisible by number of microbatches
train_ds = (tf.data.Dataset.from_tensor_slices(
(x_train, y_train)).shuffle(10000,
seed=smp.dp_rank()).batch(256,
drop_remainder=True))
test_ds = (tf.data.Dataset.from_tensor_slices(
(x_test, y_test)).shuffle(10000,
seed=smp.dp_rank()).batch(256,
drop_remainder=True))
# SMP: Define smp.DistributedModel the same way as Keras sub-classing API
class MyModel(smp.DistributedModel):
def __init__(self):
super(MyModel, self).__init__()
self.conv = Conv2D(32, 3, activation="relu")
self.flatten = Flatten()
self.dense1 = Dense(128)
self.dense2 = Dense(10)
if e.errno == errno.EEXIST and os.path.isdir(cache_dir):
pass
else:
raise
# Download and load MNIST dataset.
(x_train, y_train), (x_test, y_test) = tf.keras.datasets.mnist.load_data(
"MNIST-data-%d" % smp.rank())
x_train, x_test = x_train / 255.0, x_test / 255.0
# Add a channels dimension
x_train = x_train[..., tf.newaxis]
x_test = x_test[..., tf.newaxis]
train_ds = (tf.data.Dataset.from_tensor_slices(
(x_train, y_train)).shuffle(10000, seed=123 + smp.dp_rank()).batch(32))
class MyModel(smp.DistributedModel):
def __init__(self):
super(MyModel, self).__init__()
self.conv1 = Conv2D(
32,
3,
activation="relu",
kernel_initializer=tf.keras.initializers.GlorotNormal(seed=12))
self.conv0 = Conv2D(
32,
3,
activation="relu",
kernel_initializer=tf.keras.initializers.GlorotNormal(seed=12))