add extra block, lower feature sizes

aaa_helpers.txt

@@ -1,6 +1,6 @@
-nohup python auto_encoder_static.py --name="lhunor_ae_static_r7_2" > nohup/review/ae_static_r7_2.out 2>&1 &
-nohup python auto_encoder_body.py --name="lhunor_ae_body_r7_2" > nohup/review/ae_body_r7_2.out 2>&1 &
-nohup python auto_encoder_curv.py --name="lhunor_ae_curv_r7_2" > nohup/review/ae_curv_r7_2.out 2>&1 &
+nohup python auto_encoder_static.py --name="lhunor_ae_static_r7_3" > nohup/review/ae_static_r7_3.out 2>&1 &
+nohup python auto_encoder_body.py --name="lhunor_ae_body_r7_3" > nohup/review/ae_body_r7_3.out 2>&1 &
+nohup python auto_encoder_curv.py --name="lhunor_ae_curv_r7_3" > nohup/review/ae_curv_r7_3.out 2>&1 &
 
 one final convolutional layer for smoothing, on the final coordinates
 

auto_encoder_body.py

@@ -137,7 +137,7 @@ def train():
         if step == 0:
             to_log = { **to_log, **{ 'model/summary': wandb.Html(get_summary_html(model)) } }
 
-        if step % 1000:    
+        if step % 1000 == 0:    
             wandb.log(to_log, step=step, commit=True)
         else:
             wandb.log(to_log, step=step, commit=False)

auto_encoder_curv.py

@@ -177,7 +177,10 @@ def train():
         if step == 0:
             to_log = { **to_log, **{ 'model/summary': wandb.Html(get_summary_html(model)) } }
  
-        wandb.log(to_log, step=step)  
+        if step % 1000 == 0:    
+            wandb.log(to_log, step=step, commit=True)
+        else:
+            wandb.log(to_log, step=step, commit=False)  
         # print(step)
 
         if step % 10000 == 0:

auto_encoder_static.py

@@ -42,6 +42,7 @@ def init_wandb():
     artifact.add_file('auto_encoder_static.py')
     artifact.add_file('losses.py')
     artifact.add_file('models/base.py')
+    artifact.add_file('models/resnet.py')
     artifact.add_file('models/cond_encoder.py')
     artifact.add_file('dataset.py')
     artifact.add_file('convnext/models/convnext_tf.py')
@@ -54,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)
@@ -138,10 +139,11 @@ def train():
         if step == 0:
             to_log = { **to_log, **{ 'model/summary': wandb.Html(get_summary_html(model)) } }
 
-        if step % 100:    
-            wandb.log(to_log, step=step, commit=True)
-        else:
-            wandb.log(to_log, step=step, commit=False)
+        # 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:
             manager_enc.save(checkpoint_number=step)
@@ -187,7 +189,7 @@ def test():
 
 
 if __name__ == '__main__':
-    # init_wandb()
+    init_wandb()
     # allow_memory_growth()
     parser = argparse.ArgumentParser(description='')
     parser.add_argument('--name', type=str)

configs/review/config_ae_static.py

@@ -31,10 +31,10 @@ config['cond_ae']['decoder_dims'] = config['cond_ae']['encoder_dims']
 config['cond_ae']['activation'] = 'linear'
 
 
-config['gpu'] = '0'
+config['gpu'] = '6'
 config['lr'] = 0.0005
-config['name'] = 'static_r7_2'
-config['notes'] = 'resnet18 based'
+config['name'] = 'static_r7_3'
+config['notes'] = 'resnet18 based, one more block'
 config['tags'] = ['review_7']
 
     

models/resnet.py

@@ -145,16 +145,20 @@ class ResNetTypeI(tf.keras.Model):
                                              blocks=layer_params[0],
                                              stride=1 if not reversed else 2,
                                              upsample=reversed)
-        self.layer2 = make_basic_block_layer(filter_num=128,
+        self.layer2 = make_basic_block_layer(filter_num=64,
                                              blocks=layer_params[1],
                                              stride=2,
                                              upsample=reversed)
-        self.layer3 = make_basic_block_layer(filter_num=256,
+        self.layer3 = make_basic_block_layer(filter_num=64,
                                              blocks=layer_params[2],
                                              stride=2,
                                              upsample=reversed)
-        self.layer4 = make_basic_block_layer(filter_num=features if not is_cvae else features*2,
+        self.layer4 = make_basic_block_layer(filter_num=128,
                                              blocks=layer_params[3],
+                                             stride=2,
+                                             upsample=reversed)
+        self.layer5 = make_basic_block_layer(filter_num=features if not is_cvae else features*2,
+                                             blocks=layer_params[4],
                                              stride=2 if not reversed else 1,
                                              upsample=reversed)
 
@@ -168,8 +172,10 @@ class ResNetTypeI(tf.keras.Model):
             x = self.layer2(x, training=training)
             x = self.layer3(x, training=training)
             x = self.layer4(x, training=training)
+            x = self.layer5(x, training=training)
         else:
-            x = self.layer4(inputs, training=training)
+            x = self.layer5(inputs, training=training)
+            x = self.layer4(x, training=training)
             x = self.layer3(x, training=training)
             x = self.layer2(x, training=training)
             x = self.layer1(x, training=training)
@@ -229,7 +235,7 @@ class ResNetTypeII(tf.keras.Model):
 
 
 def resnet_18(features, reversed=False, labels=None, is_cvae=False):
-    return ResNetTypeI(layer_params=[2, 2, 2, 2], features=features, reversed=reversed, labels=labels, is_cvae=is_cvae)
+    return ResNetTypeI(layer_params=[2, 2, 2, 2, 2], features=features, reversed=reversed, labels=labels, is_cvae=is_cvae)
 
 
 def resnet_34():