38 lines
1.3 KiB
Python
38 lines
1.3 KiB
Python
import tensorflow as tf
|
|
from models.base import AEEncoder, Decoder, AEEncoderNew, DecoderNew
|
|
|
|
|
|
class AutoEncoder(tf.keras.Model):
|
|
|
|
def __init__(self, config):
|
|
super(AutoEncoder, self).__init__()
|
|
|
|
self.encoder : AEEncoder = AEEncoder(config['latent_dim'], config['encoder_depths'], config['encoder_dims'], config['bottleneck_size'], config['kernel_size'])
|
|
self.decoder : Decoder = Decoder(config['decoder_input'], config['decoder_depths'], config['decoder_dims'], kernel_size=config['kernel_size'], is_label=config['is_label'], upscale=config['upscale_output'])
|
|
|
|
|
|
def call(self, inputs, training):
|
|
x = inputs
|
|
|
|
x = self.encoder(x, training=training)
|
|
x = self.decoder(x, training=training)
|
|
|
|
return x
|
|
|
|
|
|
class AutoEncoderNew(tf.keras.Model):
|
|
|
|
def __init__(self, config):
|
|
super(AutoEncoderNew, self).__init__()
|
|
|
|
self.encoder : AEEncoderNew = AEEncoderNew(config['encoder_depths'], config['encoder_dims'], config['kernel_size'], config['activation'])
|
|
self.decoder : DecoderNew = DecoderNew(config['decoder_depths'], config['decoder_dims'], kernel_size=config['kernel_size'], is_label=config['is_label'])
|
|
|
|
|
|
def call(self, inputs, training):
|
|
x = inputs
|
|
|
|
x = self.encoder(x, training=training)
|
|
x = self.decoder(x, training=training)
|
|
|
|
return x |