I'm trying to implement a simple GAN model by following the official Tensorflow DCGAN tutorial. Posting the entire code here would be overkill, but for all practical purposes my code can be considered the same as the reference implementation. I did make changes to the layer architecture in both the discriminator and generator, but all of those modifications are ultimately irrelevant to my problem, with one key exception - the generator output layer.
Instead of outputting an image directly, like most generators do, I need mine to output some 1D feature vector instead. This feature vector is essentially a list of independent variables, which get converted to images by some featuresToImages() function in a deterministic and non-reversible fashion, and the resulting images can in turn be fed to the discriminator. To be clear, by "deterministic and non-reversible" I mean that the same feature vector will always generate an identical image, but there's no way of mapping a preexisting image back to this feature space. So, the output of my generator is the same shape/size as my feature vector. Which brings us to the crux of my problem, at least if the last couple of days of googling and getting nowhere have taught me anything...
I tested the discriminator, generator and featuresToImages() function independently, and they all work as expected. The problem crops up when I get to the training loop, or rather training step. The relevant code is:
@tf.function
def train_step(images):
noise = tf.random.normal([BATCH_SIZE, noise_dim])
with tf.GradientTape() as gen_tape, tf.GradientTape() as disc_tape:
generated_features = generator(noise, training=True)
generated_images = featuresToImages(generated_features) # Problem line - necessary but non-differentiable set of operations
real_output = discriminator(images, training=True)
fake_output = discriminator(generated_images, training=True)
gen_loss = generator_loss(fake_output)
disc_loss = discriminator_loss(real_output, fake_output)
gradients_of_generator = gen_tape.gradient(gen_loss, generator.trainable_variables)
gradients_of_discriminator = disc_tape.gradient(disc_loss, discriminator.trainable_variables)
generator_optimizer.apply_gradients(zip(gradients_of_generator, generator.trainable_variables))
discriminator_optimizer.apply_gradients(zip(gradients_of_discriminator, discriminator.trainable_variables))
My generator output isn't images, so I can't pass it directly to the discriminator like in the reference example.
If I convert the generator output to images via my featuresToImages() function, that's a non-differentiable operation, so I lose the generator gradients - gradients_of_generator returns a list of None values, so generator_optimizer.apply_gradients throws an exception, and the generator can't be trained.
I can't change the generator output, as that would defeat the entire point of the project.
I can't change the discriminator input, since I also need to pass it real images (my training dataset), which can't be reverse-mapped to a feature vector.
Is there anything I can do to make this concept work, or is it just a lost cause? Maybe a third "mapping" network that would learn how a specific feature vector correlates to a specific image output, ie. it would learn to represent the featuresToImages() function, be differentiable, and preserve the gradients between generator and discriminator? Though there's hopefully an obvious (to non-noobs like me) and cleaner solution/workaround that doesn't require an entire new subnet to implicitly learn an already explicitly defined function...