class Generator(tf.keras.Model):
"""GANomaly Generator Model
Args:
input_shape (tuple): shape of one input datum (without batch size)
latent_size (int, optional): Size of the decoder input or of the latent space. Defaults to 100.
n_filters (int, optional): Filter count of the initial convolution layer. Defaults to 64.
n_extra_layers (int, optional): Count of additional layers. Defaults to 0.
"""
def __init__(self, input_shape, latent_size=100, n_filters=64, n_extra_layers=0, **kwargs):
kwargs['name'] = type(self).__name__
super().__init__(**kwargs)
self.encoder_i = Encoder(input_shape, latent_size, n_filters, n_extra_layers, name='encoder_i').model
self.decoder = Decoder(input_shape, latent_size, n_filters, n_extra_layers, name='decoder').model
self.encoder_o = Encoder(input_shape, latent_size, n_filters, n_extra_layers, name='encoder_o').model
def summary(self, **kwargs):
print_model(self, print_fn=kwargs.get('print_fn') or print)
super().summary(**kwargs)
self.encoder_i.summary(**kwargs)
self.decoder.summary(**kwargs)
self.encoder_o.summary(**kwargs)
def call(self, x, training=False):
latent_i = self.encoder_i(x, training)
fake = self.decoder(latent_i, training)
latent_o = self.encoder_o(fake, training)
return fake, latent_i, latent_o
def test_step(self, data):
# test_step(): https://github.com/tensorflow/tensorflow/blob/v2.3.0/tensorflow/python/keras/engine/training.py#L1148-L1180
# evaluate(): https://github.com/tensorflow/tensorflow/blob/v2.3.0/tensorflow/python/keras/engine/training.py#L1243-L1394
# fit(): https://github.com/tensorflow/tensorflow/blob/v2.3.0/tensorflow/python/keras/engine/training.py#L824-L1146
x, y, _ = tf.keras.utils.unpack_x_y_sample_weight(data)
# x.shape: (batchsize, width, height, depth)
# y.shape: (batchsize, 1) on numpy array or (batchsize,) on tf.data.Dataset
_, latent_i, latent_o = self(x, training=False)
# letent_x.shape: (batchsize, 1, 1, latent_size)
error = tf.keras.backend.mean(tf.keras.backend.square(latent_i - latent_o), axis=-1)
# error.shape: (batchsize, 1, 1, 1)
return {
"losses": tf.reshape(error, (-1, 1)),
"labels": tf.reshape(y, (-1, 1))
}
def predict_step(self, data):
# https://github.com/tensorflow/tensorflow/blob/v2.3.0/tensorflow/python/keras/engine/training.py#L1396
x, _, _ = tf.keras.utils.unpack_x_y_sample_weight(data)
# x.shape: (batchsize, width, height, depth)
_, latent_i, latent_o = self(x, training=False)
# letent_x.shape: (batchsize, 1, 1, latent_size)
error = tf.keras.backend.mean(tf.keras.backend.square(latent_i - latent_o), axis=-1)
# error.shape: (batchsize, 1, 1, 1)
return tf.reshape(error, (-1, 1))
class CAE(tf.keras.Model):
def __init__(self,
input_shape,
latent_size=100,
n_filters=64,
n_extra_layers=0,
**kwargs):
kwargs['name'] = type(self).__name__
super().__init__(**kwargs)
# Use the DCGAN encoder/decoder models
self.net_enc = Encoder(input_shape,
latent_size,
n_filters,
n_extra_layers,
name='encoder').model
self.net_dec = Decoder(input_shape,
latent_size,
n_filters,
n_extra_layers,
name='decoder').model
def summary(self, **kwargs):
print_model(self)
super().summary(**kwargs)
self.net_enc.summary(**kwargs)
self.net_dec.summary(**kwargs)
def load_weights(self, path):
if not (os.path.isfile(os.path.join(path, 'encoder.index'))
and os.path.isfile(os.path.join(path, 'decoder.index'))):
warning(
'No valid pre-trained network weights in: "{}"'.format(path))
return
self.net_enc.load_weights(os.path.join(path, 'encoder'))
self.net_dec.load_weights(os.path.join(path, 'decoder'))
info('Loaded pre-trained network weights from: "{}"'.format(path))
def save_weights(self, path):
self.net_enc.save_weights(os.path.join(path, 'encoder'))
self.net_dec.save_weights(os.path.join(path, 'decoder'))
info('Saved pre-trained network weights to: "{}"'.format(path))
def call(self, x, training=False):
encoded_image = self.net_enc(x, training=training)
decoded_image = self.net_dec(encoded_image, training=training)
return decoded_image
def train_step(self, data):
x, _, sample_weight = tf.keras.utils.unpack_x_y_sample_weight(data)
return super().train_step((x, x, sample_weight))
def test_step(self, data):
x, y, _ = tf.keras.utils.unpack_x_y_sample_weight(data)
x_pred = self(x, training=False)
#loss = self.compiled_loss(
# x,
# decoded_image,
# sample_weight=sample_weight,
# regularization_losses=self.losses,
# )
# we need a loss value per image which isn't provided by compiled_loss
# assume we always have a shape of (batch_size, width, height, depth)
losses = tf.keras.backend.mean(tf.keras.backend.square(x - x_pred),
axis=[1, 2, 3])
return {
"losses": tf.reshape(losses, (-1, 1)),
"labels": tf.reshape(y, (-1, 1))
}
def predict_step(self, data):
x, _, _ = tf.keras.utils.unpack_x_y_sample_weight(data)
x_pred = self(x, training=False)
losses = tf.keras.backend.mean(tf.keras.backend.square(x - x_pred),
axis=[1, 2, 3])
return tf.reshape(losses, (-1, 1))