I am having difficulties to implement a distributed strategy on my TensorFlow v2 model and I am looking for some help.
The main difference with the documentation is that my model overrides keras train_step function.
I got the distribute strategy to work with simple model or with custom training loop but I would like to avoid the latest as I prefer using model.fit which let me use keras callbacks out of the box.
There is a link to a collab model with custom train_step similar to my model https://colab.research.google.com/github/keras-team/keras-io/blob/master/examples/generative/ipynb/dcgan_overriding_train_step.ipynb
To reproduce the error, I modified the init of the model to get the network initialization and compile steps inside the strategy scope
class GAN(keras.Model):
def __init__(self, latent_dim):
super(GAN, self).__init__()
self.latent_dim = latent_dim
strategy = tf.distribute.MirroredStrategy()
with strategy.scope():
self.discriminator = keras.Sequential(
[
keras.Input(shape=(28, 28, 1)),
layers.Conv2D(64, (3, 3), strides=(2, 2), padding="same"),
layers.LeakyReLU(alpha=0.2),
layers.Conv2D(128, (3, 3), strides=(2, 2), padding="same"),
layers.LeakyReLU(alpha=0.2),
layers.GlobalMaxPooling2D(),
layers.Dense(1),
],
name="discriminator",
)
self.generator = keras.Sequential(
[
keras.Input(shape=(self.latent_dim,)),
layers.Dense(7 * 7 * 128),
layers.LeakyReLU(alpha=0.2),
layers.Reshape((7, 7, 128)),
layers.Conv2DTranspose(128, (4, 4), strides=(2, 2), padding="same"),
layers.LeakyReLU(alpha=0.2),
layers.Conv2DTranspose(128, (4, 4), strides=(2, 2), padding="same"),
layers.LeakyReLU(alpha=0.2),
layers.Conv2D(1, (7, 7), padding="same", activation="sigmoid"),
],
name="generator",
)
self.compile(
d_optimizer=keras.optimizers.Adam(learning_rate=0.0003),
g_optimizer=keras.optimizers.Adam(learning_rate=0.0003),
loss_fn=keras.losses.BinaryCrossentropy(from_logits=True),
)
def compile(self, d_optimizer, g_optimizer, loss_fn):
super(GAN, self).compile()
self.d_optimizer = d_optimizer
self.g_optimizer = g_optimizer
self.loss_fn = loss_fn
def train_step(self, real_images):
if isinstance(real_images, tuple):
print("debug")
real_images = real_images[0]
# Sample random points in the latent space
batch_size = tf.shape(real_images)[0]
random_latent_vectors = tf.random.normal(shape=(batch_size, self.latent_dim))
# Decode them to fake images
generated_images = self.generator(random_latent_vectors)
# Combine them with real images
combined_images = tf.concat([generated_images, real_images], axis=0)
# Assemble labels discriminating real from fake images
labels = tf.concat(
[tf.ones((batch_size, 1)), tf.zeros((batch_size, 1))], axis=0
)
# Add random noise to the labels - important trick!
labels += 0.05 * tf.random.uniform(tf.shape(labels))
# Train the discriminator
with tf.GradientTape() as tape:
predictions = self.discriminator(combined_images)
d_loss = self.loss_fn(labels, predictions)
grads = tape.gradient(d_loss, self.discriminator.trainable_weights)
self.d_optimizer.apply_gradients(
zip(grads, self.discriminator.trainable_weights)
)
# Sample random points in the latent space
random_latent_vectors = tf.random.normal(shape=(batch_size, self.latent_dim))
# Assemble labels that say "all real images"
misleading_labels = tf.zeros((batch_size, 1))
# Train the generator (note that we should *not* update the weights
# of the discriminator)!
with tf.GradientTape() as tape:
predictions = self.discriminator(self.generator(random_latent_vectors))
g_loss = self.loss_fn(misleading_labels, predictions)
grads = tape.gradient(g_loss, self.generator.trainable_weights)
self.g_optimizer.apply_gradients(zip(grads, self.generator.trainable_weights))
return {"d_loss": d_loss, "g_loss": g_loss}
Then run the code with:
epochs = 30
gan = GAN(latent_dim=latent_dim)
gan.fit(
dataset, epochs=epochs, callbacks=[GANMonitor(num_img=3, latent_dim=latent_dim)]
)
The full error message :
RuntimeError: in user code:
/usr/local/lib/python3.6/dist-packages/tensorflow/python/keras/engine/training.py:806 train_function *
return step_function(self, iterator)
/usr/local/lib/python3.6/dist-packages/tensorflow/python/keras/engine/training.py:796 step_function **
outputs = model.distribute_strategy.run(run_step, args=(data,))
/usr/local/lib/python3.6/dist-packages/tensorflow/python/distribute/distribute_lib.py:1211 run
return self._extended.call_for_each_replica(fn, args=args, kwargs=kwargs)
/usr/local/lib/python3.6/dist-packages/tensorflow/python/distribute/distribute_lib.py:2585 call_for_each_replica
return self._call_for_each_replica(fn, args, kwargs)
/usr/local/lib/python3.6/dist-packages/tensorflow/python/distribute/distribute_lib.py:2945 _call_for_each_replica
return fn(*args, **kwargs)
/usr/local/lib/python3.6/dist-packages/tensorflow/python/keras/engine/training.py:789 run_step **
outputs = model.train_step(data)
<ipython-input-12-54a272890fa5>:78 train_step
zip(grads, self.discriminator.trainable_weights)
/usr/local/lib/python3.6/dist-packages/tensorflow/python/keras/optimizer_v2/optimizer_v2.py:519 apply_gradients
self._create_all_weights(var_list)
/usr/local/lib/python3.6/dist-packages/tensorflow/python/keras/optimizer_v2/optimizer_v2.py:702 _create_all_weights
_ = self.iterations
/usr/local/lib/python3.6/dist-packages/tensorflow/python/keras/optimizer_v2/optimizer_v2.py:709 __getattribute__
return super(OptimizerV2, self).__getattribute__(name)
/usr/local/lib/python3.6/dist-packages/tensorflow/python/keras/optimizer_v2/optimizer_v2.py:836 iterations
aggregation=tf_variables.VariableAggregation.ONLY_FIRST_REPLICA)
/usr/local/lib/python3.6/dist-packages/tensorflow/python/keras/optimizer_v2/optimizer_v2.py:1040 add_weight
aggregation=aggregation)
/usr/local/lib/python3.6/dist-packages/tensorflow/python/training/tracking/base.py:750 _add_variable_with_custom_getter
**kwargs_for_getter)
/usr/local/lib/python3.6/dist-packages/tensorflow/python/keras/engine/base_layer_utils.py:145 make_variable
shape=variable_shape if variable_shape else None)
/usr/local/lib/python3.6/dist-packages/tensorflow/python/ops/variables.py:260 __call__
return cls._variable_v1_call(*args, **kwargs)
/usr/local/lib/python3.6/dist-packages/tensorflow/python/ops/variables.py:221 _variable_v1_call
shape=shape)
/usr/local/lib/python3.6/dist-packages/tensorflow/python/ops/variables.py:67 getter
return captured_getter(captured_previous, **kwargs)
/usr/local/lib/python3.6/dist-packages/tensorflow/python/distribute/distribute_lib.py:2024 creator_with_resource_vars
created = self._create_variable(next_creator, **kwargs)
/usr/local/lib/python3.6/dist-packages/tensorflow/python/distribute/mirrored_strategy.py:450 _create_variable
values.MirroredVariable, values.SyncOnReadVariable, **kwargs)
/usr/local/lib/python3.6/dist-packages/tensorflow/python/distribute/distribute_utils.py:291 create_mirrored_variable
value_list = real_mirrored_creator(**kwargs)
/usr/local/lib/python3.6/dist-packages/tensorflow/python/distribute/mirrored_strategy.py:443 _real_mirrored_creator
v = next_creator(**kwargs)
/usr/local/lib/python3.6/dist-packages/tensorflow/python/ops/variables.py:67 getter
return captured_getter(captured_previous, **kwargs)
/usr/local/lib/python3.6/dist-packages/tensorflow/python/distribute/distribute_lib.py:2856 creator
_require_strategy_scope_strategy(strategy)
/usr/local/lib/python3.6/dist-packages/tensorflow/python/distribute/distribute_lib.py:344 _require_strategy_scope_strategy
_wrong_strategy_scope(strategy, context)
/usr/local/lib/python3.6/dist-packages/tensorflow/python/distribute/distribute_lib.py:322 _wrong_strategy_scope
(context.strategy, strategy))
RuntimeError: Mixing different tf.distribute.Strategy objects: <tensorflow.python.distribute.mirrored_strategy.MirroredStrategy object at 0x7ff32ba58f60> is not <tensorflow.python.distribute.distribute_lib._DefaultDistributionStrategy object at 0x7ff3c0ca0320>