log model architecture

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