luni/cvpr
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

run abl_all

Browse Source
This commit is contained in:
luni
2023-02-24 20:31:04 +00:00
parent 0ce28a2e18
commit 784fe3eb45
5 changed files with 54 additions and 44 deletions
Download Patch File Download Diff File Expand all files Collapse all files
models/base.py
+3 -3
View File
@@ -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)]
)