Fix architecture name in network diagram script (#86)

* fix architecture name * add boxes to the diagram * deleting some extra redundant lines --------- Co-authored-by: Eduardo Hirata-Miyasaki <[email protected]>
mehta-lab · Jun 18, 2024 · c56b9d0 · c56b9d0
1 parent ae3f667
commit c56b9d0
Showing 1 changed file with 7 additions and 12 deletions.
diff --git a/viscy/scripts/network_diagram.py b/viscy/scripts/network_diagram.py
@@ -46,7 +46,7 @@
 # %%
 # 3D->2D
 model = VSUNet(
-    architecture="UNext2",
+    architecture="UNeXt2",
     model_config={
         "in_channels": 2,
         "out_channels": 3,
@@ -61,17 +61,18 @@
 model_graph = draw_graph(
     model,
     model.example_input_array,
-    graph_name="UNext2",
+    graph_name="UNeXt2",
     roll=True,
     depth=3,
+    expand_nested=True,
 )
 
 model_graph.visual_graph
 
 # %%
 # 3D->3D
 model = VSUNet(
-    architecture="UNext2",
+    architecture="UNeXt2",
     model_config={
         "in_channels": 1,
         "out_channels": 2,
@@ -85,9 +86,10 @@
 model_graph = draw_graph(
     model,
     model.example_input_array,
-    graph_name="UNext2",
+    graph_name="UNeXt2",
     roll=True,
     depth=3,
+    expand_nested=True,
 )
 
 model_graph.visual_graph
@@ -118,15 +120,8 @@
     graph_name="VSCyto2D",
     roll=True,
     depth=3,
+    expand_nested=True,
 )
 
 fcmae = model_graph.visual_graph
 fcmae
-
-# %%
-
-model_graph.visual_graph.render(
-    format="svg",
-)
-
-# %%