log model architecture

auto_encoder_body.py

@@ -21,7 +21,7 @@ import wandb
 import random
 import cv2
 import numpy as np
-from custom_utils.custom_utils import UV2Mesh, load_model, maskFacesBatch
+from custom_utils.custom_utils import UV2Mesh, get_summary_html, load_model, maskFacesBatch
 import plotly.graph_objects as go
 import plotly.express as ex
 
@@ -127,12 +127,15 @@ def train():
         data_valid = static_dataset_valid.expand_batch(batch_to_dict(data_valid, config['signature_dynamic']))
         valid_loss = valid_step(data_valid)
 
+        to_log = { **train_loss, **valid_loss }
         if step % 5000 == 0 or (step < 5000 and step % 500 == 0):
             train_figs = plot_everything(batch, 'train')
             valid_figs = plot_everything(data_valid, 'valid')
-            wandb.log({ **train_loss, **valid_loss, **train_figs, **valid_figs }, step=step)
+            to_log = { **to_log, **train_figs, **valid_figs }
-        else:
+        if step == 0:
-            wandb.log({ **train_loss, **valid_loss }, step=step)
+            to_log = { **to_log, **{ 'model/summary': wandb.Html(get_summary_html(model)) } }
+
+        wandb.log(to_log, step=step)  
             
         if step % 10000 == 0:
             manager_enc.save(checkpoint_number=step)

auto_encoder_curv.py

@@ -22,7 +22,7 @@ import random
 import numpy as np
 
 
-from custom_utils.custom_utils import UV2Mesh, UV2MeshLabel, maskFacesBatch
+from custom_utils.custom_utils import UV2Mesh, UV2MeshLabel, get_summary_html, maskFacesBatch
 import plotly.graph_objects as go
 import plotly.express as ex
 
@@ -154,9 +154,11 @@ def train():
 
         to_log = { **train_loss, **valid_loss }
         if step % 5000 == 0 or (step < 5000 and step % 500 == 0):
-            imgs_train = plot_everything(batch, 'train')
+            train_figs = plot_everything(batch, 'train')
-            imgs_valid = plot_everything(data_valid, 'valid')
+            valid_figs = plot_everything(data_valid, 'valid')
-            to_log = { **to_log, **imgs_train, **imgs_valid }
+            to_log = { **to_log, **train_figs, **valid_figs }
+        if step == 0:
+            to_log = { **to_log, **{ 'model/summary': wandb.Html(get_summary_html(model)) } }
 
         wandb.log(to_log, step=step)  
         print(step)

auto_encoder_static.py

@@ -20,7 +20,7 @@ import argparse
 import wandb
 import random
 import numpy as np
-from custom_utils.custom_utils import UV2Mesh, load_model, maskFacesBatch
+from custom_utils.custom_utils import UV2Mesh, get_summary_html, load_model, maskFacesBatch, get_summary_str
 import plotly.graph_objects as go
 import plotly.express as ex
 
@@ -133,12 +133,16 @@ def train():
         data_valid = static_dataset_valid.expand_batch(batch_to_dict(data_valid, config['signature_dynamic']))
         valid_loss = valid_step(data_valid)
 
+        to_log = dict({ **train_loss, **valid_loss })
         if step % 5000 == 0 or (step < 5000 and step % 500 == 0):
             train_figs = plot_everything(batch, 'train')
             valid_figs = plot_everything(data_valid, 'valid')
-            wandb.log({ **train_loss, **valid_loss, **train_figs, **valid_figs }, step=step)
+            to_log = { **to_log, **train_figs, **valid_figs }
-        else:
+        if step == 0:
-            wandb.log({ **train_loss, **valid_loss }, step=step)
+            to_log = { **to_log, **{ 'model/summary': wandb.Html(get_summary_html(model)) } }
+
+            
+        wandb.log(to_log, step=step)
             
         if step % 10000 == 0:
             manager_enc.save(checkpoint_number=step)

custom_utils/custom_utils.py

@@ -220,3 +220,17 @@ def load_model(model, checkpoint, disable_assert=False, assert_consumed_on=False
         ckpt.restore(manager.latest_checkpoint).assert_consumed()
     else:
         ckpt.restore(manager.latest_checkpoint).assert_existing_objects_matched()#.assert_consumed()
+
+
+def get_summary_str(model):
+    lines = []
+    model.summary(print_fn=lines.append, line_length=150,expand_nested=True,show_trainable=True)
+    # Add initial spaces to avoid markdown formatting in TensorBoard
+    return lines
+
+def get_summary_html(model):
+    lines = get_summary_str(model)
+    html = ''
+    for line in lines:
+        html += "<pre>" + line + "</pre>"
+    return html