fix cvae
This commit is contained in:
+22
-6
@@ -118,13 +118,29 @@ class VariationalEncoderNew(tf.keras.Model):
|
||||
|
||||
self.encoder = resnet_18(features)
|
||||
|
||||
self.mean_branch = make_basic_block_layer(filter_num=features,
|
||||
blocks=2,
|
||||
stride=1)
|
||||
self.mean_branch = tf.keras.Sequential([
|
||||
tf.keras.layers.Conv2D(filters=features,
|
||||
kernel_size=(3, 3),
|
||||
strides=1,
|
||||
padding="same"),
|
||||
tf.keras.layers.BatchNormalization(),
|
||||
tf.keras.layers.Conv2D(filters=features,
|
||||
kernel_size=(3, 3),
|
||||
strides=1,
|
||||
padding="same")
|
||||
])
|
||||
|
||||
self.var_branch = make_basic_block_layer(filter_num=features,
|
||||
blocks=2,
|
||||
stride=1)
|
||||
self.var_branch = tf.keras.Sequential([
|
||||
tf.keras.layers.Conv2D(filters=features,
|
||||
kernel_size=(3, 3),
|
||||
strides=1,
|
||||
padding="same"),
|
||||
tf.keras.layers.BatchNormalization(),
|
||||
tf.keras.layers.Conv2D(filters=features,
|
||||
kernel_size=(3, 3),
|
||||
strides=1,
|
||||
padding="same")
|
||||
])
|
||||
|
||||
self.flatten_mean = tf.keras.layers.Flatten()
|
||||
self.flatten_var = tf.keras.layers.Flatten()
|
||||
|
||||
Reference in New Issue
Block a user