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))