update pyramid

This commit is contained in:
2022-10-16 14:57:36 +00:00
parent a6834caf80
commit 07c9fb57df
3 changed files with 10 additions and 10 deletions
+3 -3
View File
@@ -10,7 +10,7 @@ from custom_utils.custom_utils import load_model
from models.base import AEEncoder
from models.cvae_model import CVAE
from models.pyramid_level import PyramidLevel
from models.upscaler import UpscalerBasic, UpscalerBasicNew
from models.upscaler import UpscalerBasic, UpscalerBasicConv, UpscalerBasicNew
class Pyramid(tf.keras.Model):
@@ -22,7 +22,7 @@ class Pyramid(tf.keras.Model):
if config['pyramid']['static_model'] != '':
load_model(self.static_enc, config['pyramid']['static_model'] + '_enc')
self.body_enc : AEEncoder = AEEncoder(config['cond_ae']['latent_dim'], config['cond_ae']['encoder_depths'], config['cond_ae']['encoder_dims'], config['cond_ae']['bottleneck_size'], config['cond_ae']['kernel_size'])
self.body_enc : AEEncoder = AEEncoder(config['cond_ae_body']['latent_dim'], config['cond_ae_body']['encoder_depths'], config['cond_ae_body']['encoder_dims'], config['cond_ae_body']['bottleneck_size'], config['cond_ae_body']['kernel_size'])
if config['pyramid']['body_model'] != '':
load_model(self.body_enc, config['pyramid']['body_model'] + '_enc')
@@ -45,7 +45,7 @@ class Pyramid(tf.keras.Model):
self.upscale_enhancer = None
if config['pyramid']['final_upscale_enhancer'] != '':
self.upscale_enhancer = UpscalerBasicNew(low_res=config['pyramid']['resolutions'][-1])
self.upscale_enhancer = UpscalerBasicConv(low_res=config['pyramid']['resolutions'][-1])
load_model(self.upscale_enhancer, config['pyramid']['final_upscale_enhancer'])
self.upscale_enhancer.trainable = False