started training pyramid

This commit is contained in:
2023-02-26 22:25:50 +00:00
parent a678724711
commit 64ecb900f5
8 changed files with 314 additions and 79 deletions
+3 -4
View File
@@ -102,13 +102,12 @@ class PyramidLevelNew(tf.keras.Model):
super(PyramidLevelNew, self).__init__()
self.size = config['resolution']
self.encoder = VariationalEncoderNew(config['encoder_depths'], config['encoder_dims'], kernel_size=config['kernel_size'])
self.encoder = VariationalEncoderNew(config['encoder_depths'], config['encoder_dims'], kernel_size=config['kernel_size'], latent_dim=config['latent_dim'])
self.decoder = DecoderNew(config['decoder_depths'], config['decoder_dims'], config['kernel_size'])
self.upscale_enhancer = None
if config['upscale_enhancer'] != '':
k_size = config['upscale_enhancer_kernel_size'] if 'upscale_enhancer_kernel_size' in config.keys() else ((3 if config['resolution'] <= 64 else 5))
self.upscale_enhancer = UpscalerBasicConv(config['resolution'], kernel_size=k_size)
self.upscale_enhancer = UpscalerBasicConv(config['resolution'], kernel_size=config['kernel_size'])
load_model(self.upscale_enhancer, config['upscale_enhancer'])
self.upscale_enhancer.trainable = not config['upscale_enhancer_freeze']
@@ -117,7 +116,7 @@ class PyramidLevelNew(tf.keras.Model):
else:
self.avg_garment = None
self.reshape = tf.keras.layers.Reshape((8, 4, 64))
self.reshape = tf.keras.layers.Reshape((8, 4, config['latent_dim']))
def enhance(self, low, mask, training=None):