trying to get resnet working

aaa_helpers.txt

@@ -1,6 +1,7 @@
 nohup python cvae.py --name="lhunor_weights_test_6" > nohup/iccv/weights_test_6.out 2>&1 &
 nohup python evaluate_ae.py --name iccv_weights_test_6 > nohup/iccv/evaluate_iccv_weights_test_6.out 2>&1 &
 nohup python calculate_error.py --name iccv_weights_test_6 > nohup/iccv/error_iccv_weights_test_6.out 2>&1 &
+nohup python auto_encoder_static.py --name="lhunor_ae_static_iccv_resnet_1" > nohup/iccv/ae_static_iccv_resnet_1.out 2>&1 &
 
 
 one final convolutional layer for smoothing, on the final coordinates

auto_encoder_static.py

@@ -37,7 +37,7 @@ from models.cond_encoder import AutoEncoderNew as AutoEncoder
 
 
 def init_wandb():
-    run = wandb.init(project='ae', name=config['name'], notes=config['notes'], dir='/home/workspace/local-data/temp', tags=config['tags'])
+    run = wandb.init(project='ae', name=config['name'], notes=config['notes'], dir='/tmp', tags=config['tags'])
     artifact = wandb.Artifact('code', type='code')
     artifact.add_file('auto_encoder_static.py')
     artifact.add_file('losses.py')
@@ -55,7 +55,7 @@ uv2mesh = UV2Mesh(res=256, is_skirt=config['is_skirt'], is_smpl_orig=config['is_
 
 
 
-# @tf.function
+@tf.function
 def train_step(batch):
     with tf.GradientTape() as tape:
         pred = model([batch['uv_static'], batch['uv_mask']], training=True)
@@ -139,10 +139,7 @@ def train():
         if step == 0:
             to_log = { **to_log, **{ 'model/summary': wandb.Html(get_summary_html(model)) } }
 
-        # if step % 1000 == 0:    
-        #     wandb.log(to_log, step=step, commit=True)
-        # else:
-            # wandb.log(to_log, step=step, commit=False)
+
         wandb.log(to_log, step=step)
 
         if step % 10000 == 0:
@@ -188,8 +185,20 @@ def test():
     print()
 
 
+class MyLRSchedule(tf.keras.optimizers.schedules.LearningRateSchedule):
+
+  def __init__(self, initial_learning_rate):
+    self.initial_learning_rate = initial_learning_rate
+
+  def __call__(self, step):
+    if step < 1000:
+        return self.initial_learning_rate / 10.0
+    else:
+        return self.initial_learning_rate
+  
+
 if __name__ == '__main__':
-    # init_wandb()
+    init_wandb()
     parser = argparse.ArgumentParser(description='')
     parser.add_argument('--name', type=str)
     args = vars(parser.parse_args())

configs/iccv/config_ae_static.py

@@ -23,14 +23,14 @@ config['cond_ae']['is_label'] = False
 config['cond_ae']['upscale_output'] = False
 config['cond_ae']['latent_dim'] = config['latent_dim']
 config['cond_ae']['kernel_size'] = 7
-config['cond_ae']['bottleneck_size'] = 4 # 8*4
-config['cond_ae']['layer_params'] = [3, 4, 5, 4, 3]
+config['cond_ae']['bottleneck_size'] = 8*4 # 8*4
+config['cond_ae']['layer_params'] = [3, 4, 23, 4, 3]
 # config['cond_ae']['layer_params'] = [3, 4, 23, 3, 3]
 config['cond_ae']['activation'] = 'linear'
 
 
 config['gpu'] = '5'
-config['lr'] = 0.0005
+config['lr'] = 0.0001
 config['name'] = 'static_resnet_1'
 config['notes'] = 'resnet101 based, one more block'
 config['tags'] = ['iccv', 'resnet']

models/base.py

@@ -116,7 +116,7 @@ class VariationalEncoderNew(tf.keras.Model):
     def __init__(self, features):
         super(VariationalEncoderNew, self).__init__()
 
-        self.encoder = resnet_101(features)
+        self.encoder = resnet_18(features)
 
         self.mean_branch = tf.keras.Sequential([
             tf.keras.layers.Conv2D(filters=features,
@@ -166,7 +166,7 @@ class AEEncoderNew(tf.keras.Model):
     def __init__(self, layer_params, features, activation='tanh'):
         super(AEEncoderNew, self).__init__()
 
-        self.encoder = resnet_101(layer_params, features)
+        self.encoder = resnet_18(layer_params, features)
         self.act = tf.keras.layers.Activation(activation)
 
 
@@ -193,9 +193,9 @@ class DecoderNew(tf.keras.Model):
         self.is_label = is_label
 
         if not is_label:
-            self.decoder = resnet_101(layer_params, features, reversed=True, is_cvae=is_cvae)
+            self.decoder = resnet_18(layer_params, features, reversed=True, is_cvae=is_cvae)
         else:
-            self.decoder = resnet_101(layer_params, features, reversed=True, labels=6, is_cvae=is_cvae)
+            self.decoder = resnet_18(layer_params, features, reversed=True, labels=6, is_cvae=is_cvae)
             
 
     def call(self, inputs, training):