pythontensorflowdcganmultiple-gpu

Trying to create optimizer slot variable under the scope for tf.distribute.Strategy, which is different from the scope used for the original variable

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.
Solution

  • 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.