[Keras] VAEモデルを学習するとき、”you should implement a `call` method"のメッセージが出ます

Question

Kerasを用いて変分オートエンコーダ(Variable AutoEncoder: VAE)を作りたいと思います。Kerasのブログの内容を参考に、自分の画像データで学習したく、データの供給部分をImageDataGeneratorの部分を変更しました。

参考元のVAEコード

しかし、コード実行の時、最後行である学習の部分(fit)で、”you should implement a call method"というエラーが発生しています。

何が原因でエラーが発生しているか教えて頂ければと思います。

import tensorflow as tf
from tensorflow import keras
from tensorflow.keras import layers

# Creating a sampling layer
class Sampling(layers.Layer):
    '''Uses (z_mean, z_log_var) to sample z'''

    def call(self, inputs):
        z_mean, z_log_var = inputs
        batch = tf.shape(z_mean)[0] #tensorのshapeを返す
        dim = tf.shape(z_mean)[1]
        epsilon = tf.keras.backend.random_normal(shape=(batch, dim))
        return z_mean + tf.exp(0.5 * z_log_var) * epsilon



# Encoder Network

latent_dim = 2

encoder_iputs = keras.Input(shape=(28,28,1))

x = layers.Conv2D(32,3, activation='relu', strides=2, padding='same')(encoder_iputs)
x = layers.Conv2D(64,3, activation='relu', strides=2, padding='same')(x)
x = layers.Flatten()(x)
x = layers.Dense(16, activation='relu')(x)

z_mean = layers.Dense(latent_dim, name='z_mean')(x)
z_log_var = layers.Dense(latent_dim, name='z_log_var')(x)

z = Sampling()([z_mean, z_log_var])
encoder = keras.Model(encoder_iputs, [z_mean, z_log_var, z], name = 'encoder')
encoder.summary()


# Decoder Network

latent_inputs = keras.Input(shape=(latent_dim,))
x = layers.Dense(7*7*64, activation='relu')(latent_inputs)
x = layers.Reshape((7,7,64))(x)
x = layers.Conv2DTranspose(64, 3, activation='relu', strides=2, padding='same')(x)
x = layers.Conv2DTranspose(32, 3, activation='relu', strides=2, padding='same')(x)
decoder_outputs = layers.Conv2DTranspose(1,3,activation='sigmoid', padding='same')(x)
decoder = keras.Model(latent_inputs, decoder_outputs, name='decoder')
decoder.summary()

#Define the VAE as a Model with a custom train_step

class VAE(keras.Model):
    def __init__(self, encoder, decoder, **kwargs):
        super(VAE, self).__init__(**kwargs)
        self.encoder = encoder
        self.decoder = decoder

    def train_step(self,data):
        if isinstance(data, tuple):
            data = data[0]
        with tf.GradientTape() as tape:
            z_mean, z_log_var, z = encoder(data)
            reconstruction = decoder(z)
            reconstruction_loss = tf.reduce_mean(
                keras.losses.binary_crossentropy(data, reconstruction)
            )
            reconstruction_loss *= 28* 28
            kl_loss = 1 + z_log_var - tf.square(z_mean)-tf.exp(z_log_var)
            kl_loss = tf.reduce_mean(kl_loss)
            kl_loss *= -0.5
            total_loss = reconstruction_loss + kl_loss

        grads = tape.gradient(total_loss, self.trainable_weights)
        self.optimizer.apply_gradients(zip(grads, self.trainable_weights))
        return{
            "loss": total_loss,
            "reconstruction_loss": reconstruction_loss,
            "kl_loss": kl_loss,
        }


#Generator
from tensorflow.keras.preprocessing.image import ImageDataGenerator

base_dir = 'C:\DATA\DOGS'  # Surface
BATCH_SIZE = 256; EPOCHS = 200
train_datagen = ImageDataGenerator(rescale=1./255)
train_batches = train_datagen.flow_from_directory(base_dir,
        target_size=(28,28), shuffle=True, class_mode='input', batch_size=BATCH_SIZE,color_mode="grayscale")

data_batch, labels_batch = next(train_batches)

print('data batch shape', data_batch.shape)

#Train

vae = VAE(encoder, decoder)
vae.compile(optimizer=keras.optimizers.Adam())
vae.fit(x=train_batches, steps_per_epoch = train_batches.samples // BATCH_SIZE,  epochs=EPOCHS,verbose=1 )

エラーコード

Traceback (most recent call last):
  File "C:/Users/se/84.VAE_again/01.CAE-VAE-Test-Generato-Q&Ar.py", line 102, in <module>
    vae.fit(x=train_batches, steps_per_epoch = train_batches.samples // BATCH_SIZE,  epochs=EPOCHS,verbose=1 )
  File "C:\Users\se\Anaconda3\envs\kaneko-pip-exe-test\lib\site-packages\tensorflow\python\keras\engine\training.py", line 66, in _method_wrapper
    return method(self, *args, **kwargs)
  File "C:\Users\se\Anaconda3\envs\kaneko-pip-exe-test\lib\site-packages\tensorflow\python\keras\engine\training.py", line 815, in fit
    model=self)
  File "C:\Users\se\Anaconda3\envs\kaneko-pip-exe-test\lib\site-packages\tensorflow\python\keras\engine\data_adapter.py", line 1112, in __init__
    model=model)
  File "C:\Users\se\Anaconda3\envs\kaneko-pip-exe-test\lib\site-packages\tensorflow\python\keras\engine\data_adapter.py", line 908, in __init__
    **kwargs)
  File "C:\Users\se\Anaconda3\envs\kaneko-pip-exe-test\lib\site-packages\tensorflow\python\keras\engine\data_adapter.py", line 781, in __init__
    lambda x: model(x, training=False), args=(concrete_x,))
  File "C:\Users\se\Anaconda3\envs\kaneko-pip-exe-test\lib\site-packages\tensorflow\python\distribute\distribute_lib.py", line 951, in run
    return self._extended.call_for_each_replica(fn, args=args, kwargs=kwargs)
  File "C:\Users\se\Anaconda3\envs\kaneko-pip-exe-test\lib\site-packages\tensorflow\python\distribute\distribute_lib.py", line 2290, in call_for_each_replica
    return self._call_for_each_replica(fn, args, kwargs)
  File "C:\Users\se\Anaconda3\envs\kaneko-pip-exe-test\lib\site-packages\tensorflow\python\distribute\distribute_lib.py", line 2649, in _call_for_each_replica
    return fn(*args, **kwargs)
  File "C:\Users\se\Anaconda3\envs\kaneko-pip-exe-test\lib\site-packages\tensorflow\python\autograph\impl\api.py", line 282, in wrapper
    return func(*args, **kwargs)
  File "C:\Users\se\Anaconda3\envs\kaneko-pip-exe-test\lib\site-packages\tensorflow\python\keras\engine\data_adapter.py", line 781, in <lambda>
    lambda x: model(x, training=False), args=(concrete_x,))
  File "C:\Users\se\Anaconda3\envs\kaneko-pip-exe-test\lib\site-packages\tensorflow\python\keras\engine\base_layer.py", line 968, in __call__
    outputs = self.call(cast_inputs, *args, **kwargs)
  File "C:\Users\se\Anaconda3\envs\kaneko-pip-exe-test\lib\site-packages\tensorflow\python\keras\engine\network.py", line 714, in call
    raise NotImplementedError('When subclassing the `Model` class, you should'
NotImplementedError: When subclassing the `Model` class, you should implement a `call` method.