run abl_all
This commit is contained in:
+3
-3
@@ -111,18 +111,18 @@ class Decoder(tf.keras.Model):
|
||||
|
||||
class VariationalEncoderNew(tf.keras.Model):
|
||||
|
||||
def __init__(self, depths, dims, kernel_size=7):
|
||||
def __init__(self, depths, dims, kernel_size=7, latent_dim=64):
|
||||
super(VariationalEncoderNew, self).__init__()
|
||||
|
||||
self.encoder = ConvNeXt(depths=depths, dims=dims, include_top=False, kernel_size=kernel_size)
|
||||
|
||||
self.mean_branch = tf.keras.Sequential(
|
||||
[Block(dim=64, drop_path=0, kernel_size=kernel_size,
|
||||
[Block(dim=latent_dim, drop_path=0, kernel_size=kernel_size,
|
||||
layer_scale_init_value=1e-6, prefix=f'bottleneck{j}') for j in range(5)]
|
||||
)
|
||||
|
||||
self.var_branch = tf.keras.Sequential(
|
||||
[Block(dim=64, drop_path=0, kernel_size=kernel_size,
|
||||
[Block(dim=latent_dim, drop_path=0, kernel_size=kernel_size,
|
||||
layer_scale_init_value=1e-6, prefix=f'bottleneck{j}') for j in range(5)]
|
||||
)
|
||||
|
||||
|
||||
@@ -179,10 +179,10 @@ class CVAENew(tf.keras.Model):
|
||||
if config['use_curvature']:
|
||||
self.curv_enc.trainable = False
|
||||
|
||||
self.encoder = VariationalEncoderNew(config['cvae']['encoder_depths'], config['cvae']['encoder_dims'], kernel_size=config['cvae']['kernel_size'])
|
||||
self.encoder = VariationalEncoderNew(config['cvae']['encoder_depths'], config['cvae']['encoder_dims'], kernel_size=config['cvae']['kernel_size'], latent_dim=config['cvae']['latent_dim'])
|
||||
self.decoder = DecoderNew(config['cvae']['decoder_depths'], config['cvae']['decoder_dims'], config['cvae']['kernel_size'])
|
||||
|
||||
self.reshape = tf.keras.layers.Reshape((8, 4, 64))
|
||||
self.reshape = tf.keras.layers.Reshape((8, 4, config['cvae']['latent_dim']))
|
||||
|
||||
if config['avg_garment_path'] != '':
|
||||
self.avg_garment = np.load(config['avg_garment_path'])
|
||||
|
||||
Reference in New Issue
Block a user