I’ve been following along with a video series from 2020 that details how to build a variational autoencoder inside of Python, and I’ve had to reformat the code in some places to get it to work with the modern version of Tensorflow. I’ve run into an issue where I try to use K.function on a dense layer (called self.mu), but I keep getting an error saying that I can’t pass a Keras tensor into a Tensorflow function. For example, here is some of the problematic code:

def _calculate_kl_loss(self, y_target, y_predicted):
    kl_loss = -0.5 * K.sum(1 + self.log_variance - K.square(self.mu) -
                           K.exp(self.log_variance), axis=1)
    return kl_loss

K.square throws up the error; I initially was able to solve or bypass it by wrapping the function in a layer, which is what Tensorflow itself suggests, but I ran into the same error at this point in the code:

@tf.function
def train(self, x_train, batch_size, num_epochs):
    self.model.fit(x_train,
                   x_train,
                   batch_size=batch_size,
                   epochs=num_epochs,
                   shuffle=True)

When I tried wrapping this in a layer, nothing happened. Here’s what I tried:

class TrainingLayer(Layer):
    def __init__(self, model, x_train, batch_size, num_epochs):
        super(TrainingLayer, self).__init__()
        self.model = model
        self.x_train = x_train
        self.batch_size = batch_size
        self.num_epochs = num_epochs

    def call(self, inputs):
        self.model.fit(self.x_train, self.x_train,
                       batch_size=self.batch_size,
                       epochs=self.num_epochs,
                       shuffle=True)

In my VAE class, I have:

@tf.function
def train(self, x_train, batch_size, num_epochs):
    TrainingLayer(model=self.model, x_train=x_train, batch_size=batch_size, num_epochs=num_epochs)

When I do this, the code runs, but the model is not trained at all. Does anybody know what’s going on and how I can fix this?

Here is a link to the github repo with the original code that was generated in the video series: https://github.com/musikalkemist/generating-sound-with-neural-networks/blob/main/11%20Implementing%20VAE/code/autoencoder.py

The author seems unresponsive.