I want to develop a DCGAN with a resolution of 1024x1024. For this, I need to use multiple GPUs, otherwise it might take too much time. I refer to the introduction in https://www.tensorflow.org/guide/distributed_training documentation
At the top of the script I used
strategy = tf.distribute.MirroredStrategy()
Then inside the DCGAN I used
with strategy.scope():
The error I get is:
ValueError:Trying to create optimizer slot variable under the scope for tf.distribute.Strategy, which is different from the scope used for the original variable. Make sure the slot variables are created under the same strategy scope. This may happen if you're restoring from a checkpoint outside the scope.
The following is my code:
strategy = tf.distribute.MirroredStrategy()
dataset = keras.preprocessing.image_dataset_from_directory(
"test2", label_mode=None, image_size=(1024, 1024), batch_size=4)
dataset = dataset.map(lambda x: x / 255.0)
discriminator = keras.Sequential(
[
keras.Input(shape=(1024, 1024, 3)),
layers.Conv2D(8, kernel_size=4, strides=2, padding="same"),
layers.LeakyReLU(alpha=0.2),
layers.Conv2D(8, kernel_size=4, strides=2, padding="same"),
layers.LeakyReLU(alpha=0.2),
layers.Conv2D(16, kernel_size=4, strides=2, padding="same"),
layers.LeakyReLU(alpha=0.2),
layers.Conv2D(16, kernel_size=4, strides=2, padding="same"),
layers.LeakyReLU(alpha=0.2),
layers.Conv2D(32, kernel_size=4, strides=2, padding="same"),
layers.LeakyReLU(alpha=0.2),
layers.Conv2D(32, kernel_size=4, strides=2, padding="same"),
layers.LeakyReLU(alpha=0.2),
layers.Flatten(),
layers.Dropout(0.2),
layers.Dense(1, activation="sigmoid"),
],
name="discriminator",
)
discriminator.summary()
latent_dim = 1024
generator = keras.Sequential(
[
keras.Input(shape=(latent_dim,)),
layers.Dense(16 * 16 * 32),
layers.Reshape((16, 16, 32)),
layers.Conv2DTranspose(32, kernel_size=4, strides=2, padding="same"),
layers.LeakyReLU(alpha=0.2),
layers.Conv2DTranspose(32, kernel_size=4, strides=2, padding="same"),
layers.LeakyReLU(alpha=0.2),
layers.Conv2DTranspose(32, kernel_size=4, strides=2, padding="same"),
layers.LeakyReLU(alpha=0.2),
layers.Conv2DTranspose(32, kernel_size=4, strides=2, padding="same"),
layers.LeakyReLU(alpha=0.2),
layers.Conv2DTranspose(32, kernel_size=4, strides=2, padding="same"),
layers.LeakyReLU(alpha=0.2),
layers.Conv2DTranspose(32, kernel_size=4, strides=2, padding="same"),
layers.LeakyReLU(alpha=0.2),
layers.Conv2D(3, kernel_size=5, padding="same", activation="sigmoid"),
],
name="generator",
)
generator.summary()
class GAN(keras.Model):
def __init__(self, strategy, discriminator, generator, latent_dim):
super(GAN, self).__init__()
self.discriminator = discriminator
self.generator = generator
self.latent_dim = latent_dim
self.global_batchsize = 32
self.strategy = strategy
self.batchsize_per_replica = int(self.global_batchsize/self.strategy.num_replicas_in_sync)
def loss_fn(self, labels, predictions):
loss_fn = tf.keras.losses.BinaryCrossentropy(from_logits=True,\
reduction=tf.keras.losses.Reduction.NONE)
return loss_fn(labels, predictions)
def compile(self, d_optimizer, g_optimizer):
super(GAN, self).compile()
self.d_optimizer = d_optimizer
self.g_optimizer = g_optimizer
self.d_loss_metric = keras.metrics.Mean(name="d_loss")
self.g_loss_metric = keras.metrics.Mean(name="g_loss")
def metrics(self):
return [self.d_loss_metric, self.g_loss_metric]
def disc_loss(self, real_output, fake_output):
real_loss = self.loss_fn(tf.ones_like(real_output), real_output)
fake_loss = self.loss_fn(tf.zeros_like(fake_output), fake_output)
total_loss = real_loss + fake_loss
total_loss = total_loss/self.global_batchsize
return total_loss
def gen_loss(self, fake_output):
gen_loss = self.loss_fn(tf.ones_like(fake_output), fake_output)
gen_loss = gen_loss / self.global_batchsize
return gen_loss
def distribute_trainstep(self, dist_dataset):
per_replica_g_losses, per_replica_d_losses = self.strategy.experimental_run_v2(self.train_step,dist_dataset)
total_g_loss = self.strategy.reduce(tf.distribute.ReduceOp.SUM, per_replica_g_losses,axis=0)
total_d_loss = self.strategy.reduce(tf.distribute.ReduceOp.SUM, per_replica_d_losses, axis=0)
return total_g_loss, total_d_loss
def train_step(self, real_images):
batch_size = tf.shape(real_images)[0]
random_latent_vectors = tf.random.normal(shape=(batch_size, self.latent_dim))
generated_images = self.generator(random_latent_vectors)
combined_images = tf.concat([generated_images, real_images], axis=0)
labels = tf.concat(
[tf.ones((batch_size, 1)), tf.zeros((batch_size, 1))], axis=0
)
labels += 0.05 * tf.random.uniform(tf.shape(labels))
noise = tf.random.normal(shape=[tf.shape(real_images)[0], self.latent_dim])
with tf.GradientTape() as gen_tape, tf.GradientTape() as disc_tape:
generated_imgs = self.generator(noise, training=True)
real_output = self.discriminator(real_images, training=True)
fake_output = self.discriminator(generated_imgs, training=True)
d_loss = self.disc_loss(real_output, fake_output)
g_loss = self.gen_loss(fake_output)
G_grads = gen_tape.gradient(g_loss, self.generator.trainable_variables)
D_grads = disc_tape.gradient(d_loss, self.discriminator.trainable_variables)
self.g_optimizer.apply_gradients(zip(G_grads, self.generator.trainable_variables))
self.d_optimizer.apply_gradients(zip(D_grads, self.discriminator.trainable_variables))
with tf.GradientTape() as gen_tape:
generated_imgs = self.generator_model(noise, training=True)
fake_output = self.discriminator(generated_imgs, training=True)
g_loss = self.gen_loss(fake_output)
G_grads = gen_tape.gradient(g_loss, self.generator_model.trainable_variables)
self.g_optimizer.apply_gradients(zip(G_grads, self.generator.trainable_variables))
return g_loss, d_loss
class GANMonitor(keras.callbacks.Callback):
def __init__(self, num_img=6, latent_dim=32):
self.num_img = num_img
self.latent_dim = latent_dim
def on_epoch_end(self, epoch, logs=None):
random_latent_vectors = tf.random.normal(shape=(self.num_img, self.latent_dim))
generated_images = self.model.generator(random_latent_vectors)
generated_images *= 255
generated_images.numpy()
for i in range(self.num_img):
img = keras.preprocessing.image.array_to_img(generated_images[i])
if epoch %50 ==0:
img.save("./1024/generated_img_%03d_%d.png" % (epoch, i))
epochs = 5000
with strategy.scope():
gan = GAN(strategy, discriminator=discriminator, generator=generator, latent_dim=latent_dim)
gan.compile(
d_optimizer=keras.optimizers.Adam(learning_rate=0.0001),
g_optimizer=keras.optimizers.Adam(learning_rate=0.0001),
)
gan.fit(
dataset, epochs=epochs, callbacks=[GANMonitor(num_img=60, latent_dim=latent_dim)]
)
The error is the following
Epoch 1/5000
/home/kuo/.local/lib/python3.8/site-packages/tensorflow/python/util/dispatch.py:1082: UserWarning: "`binary_crossentropy` received `from_logits=True`, but the `output` argument was produced by a sigmoid or softmax activation and thus does not represent logits. Was this intended?"
return dispatch_target(*args, **kwargs)
Traceback (most recent call last):
File "1024.py", line 253, in <module>
gan.fit(
File "/home/kuo/.local/lib/python3.8/site-packages/keras/utils/traceback_utils.py", line 67, in error_handler
raise e.with_traceback(filtered_tb) from None
File "/home/kuo/.local/lib/python3.8/site-packages/tensorflow/python/framework/func_graph.py", line 1147, in autograph_handler
raise e.ag_error_metadata.to_exception(e)
ValueError: in user code:
File "/home/kuo/.local/lib/python3.8/site-packages/keras/engine/training.py", line 1021, in train_function *
return step_function(self, iterator)
File "/home/kuo/.local/lib/python3.8/site-packages/keras/engine/training.py", line 1010, in step_function **
outputs = model.distribute_strategy.run(run_step, args=(data,))
File "/usr/local/lib/python3.8/dist-packages/six.py", line 703, in reraise
raise value
File "/home/kuo/.local/lib/python3.8/site-packages/keras/engine/training.py", line 1000, in run_step **
outputs = model.train_step(data)
File "1024.py", line 179, in train_step
self.g_optimizer.apply_gradients(zip(G_grads, self.generator.trainable_variables))
File "/home/kuo/.local/lib/python3.8/site-packages/keras/optimizer_v2/optimizer_v2.py", line 639, in apply_gradients
self._create_all_weights(var_list)
File "/home/kuo/.local/lib/python3.8/site-packages/keras/optimizer_v2/optimizer_v2.py", line 825, in _create_all_weights
self._create_slots(var_list)
File "/home/kuo/.local/lib/python3.8/site-packages/keras/optimizer_v2/adam.py", line 117, in _create_slots
self.add_slot(var, 'm')
File "/home/kuo/.local/lib/python3.8/site-packages/keras/optimizer_v2/optimizer_v2.py", line 902, in add_slot
raise ValueError(
ValueError: Trying to create optimizer slot variable under the scope for tf.distribute.Strategy (<tensorflow.python.distribute.mirrored_strategy.MirroredStrategy object at 0x7f72f39c0430>), which is different from the scope used for the original variable (<tf.Variable 'dense_1/kernel:0' shape=(1024, 8192) dtype=float32, numpy=
array([[-0.00106893, 0.01506512, -0.01771315, ..., -0.01528796,
-0.02354955, -0.0135217 ],
[-0.01760183, -0.02044552, 0.00945723, ..., -0.02140231,
0.01164402, 0.01851213],
[ 0.00233763, -0.0196434 , 0.01152603, ..., -0.02139488,
0.0125667 , 0.0251492 ],
...,
[ 0.00782686, 0.00941393, 0.00423452, ..., -0.0052203 ,
-0.02194414, -0.0167138 ],
[ 0.02420759, -0.02258933, 0.01125678, ..., -0.00626962,
0.00758442, 0.0015665 ],
[-0.00925244, -0.02154037, -0.0209455 , ..., -0.01146874,
0.00285936, 0.01914702]], dtype=float32)>). Make sure the slot variables are created under the same strategy scope. This may happen if you're restoring from a checkpoint outside the scope.
Using the Keras functional API instead of the sequential API to specify the network architectures fixed this for me. See https://keras.io/guides/functional_api.