log model architecture

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
 
@@ -152,11 +152,13 @@ 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}
+        to_log = { **train_loss, **valid_loss }
         if step % 5000 == 0 or (step < 5000 and step % 500 == 0):
-            imgs_train = plot_everything(batch, 'train')
-            imgs_valid = plot_everything(data_valid, 'valid')
-            to_log = { **to_log, **imgs_train, **imgs_valid }
+            train_figs = plot_everything(batch, 'train')
+            valid_figs = plot_everything(data_valid, '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)