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)
-        else:
-            wandb.log({ **train_loss, **valid_loss }, step=step)
+            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)  
             
         if step % 10000 == 0:
             manager_enc.save(checkpoint_number=step)