Implement LeNet Style Visualization (#38)

* Add pyproject.toml and fix ruff issues * open file using Path * fix ruff issues * Add lenet style visualization * add docs for lenet style * add usage example for lenet style * fix usage example for LeNet style * Add LeNet style
willyfh · Apr 10, 2024 · 71dab1b · 71dab1b
1 parent 09f26c6
commit 71dab1b
Show file tree

Hide file tree

Showing 9 changed files with 410 additions and 5 deletions.
diff --git a/README.md b/README.md
@@ -5,13 +5,13 @@
 
 </div>
 
-**VisualTorch** aims to help visualize Torch-based neural network architectures. It currently supports generating layered-style and graph-style architectures for PyTorch Sequential and Custom models. This tool is adapted from [visualkeras](https://github.com/paulgavrikov/visualkeras), [pytorchviz](https://github.com/szagoruyko/pytorchviz), and [pytorch-summary](https://github.com/sksq96/pytorch-summary).
+**VisualTorch** aims to help visualize Torch-based neural network architectures. It currently supports generating layered-style, graph-style, and LeNet-style architectures for PyTorch Sequential and Custom models. This tool is adapted from [visualkeras](https://github.com/paulgavrikov/visualkeras), [pytorchviz](https://github.com/szagoruyko/pytorchviz), and [pytorch-summary](https://github.com/sksq96/pytorch-summary).
 
 **Note:** VisualTorch may not yet support complex models, but contributions are welcome!
 
 <div align="center">
 
-![VisualTorch Examples](https://github.com/willyfh/visualtorch/assets/5786636/7e2c35ea-d34d-4b92-b414-285bccb8576a)
+![VisualTorch Examples](https://github.com/willyfh/visualtorch/assets/5786636/398c3356-4de0-446b-a30b-d8ebe532d2c2)
 
 </div>
 

diff --git a/docs/examples/lenet_style/README.md b/docs/examples/lenet_style/README.md
@@ -0,0 +1,5 @@
+============
+LeNet Style
+============
+
+These examples cover LeNet style visualization
diff --git a/docs/examples/lenet_style/plot_basic_sequential_lenet_style.py b/docs/examples/lenet_style/plot_basic_sequential_lenet_style.py
@@ -0,0 +1,26 @@
+"""LeNet Style Basic Sequential
+=======================================
+
+Visualization of basic custom model
+"""  # noqa: D205
+
+import matplotlib.pyplot as plt
+import visualtorch
+from torch import nn
+
+# Example of a simple CNN model using nn.Sequential
+model = nn.Sequential(
+    nn.Conv2d(3, 8, kernel_size=3, padding=1),
+    nn.MaxPool2d(2, 2),
+    nn.Conv2d(8, 16, kernel_size=3, padding=1),
+    nn.MaxPool2d(2, 2),
+)
+
+input_shape = (1, 3, 128, 128)
+
+img = visualtorch.lenet_view(model, input_shape=input_shape)
+
+plt.axis("off")
+plt.tight_layout()
+plt.imshow(img)
+plt.show()
diff --git a/docs/source/conf.py b/docs/source/conf.py
@@ -91,7 +91,7 @@
     "examples_dirs": "../examples",  # path to your example scripts
     "gallery_dirs": "usage_examples",  # path to where to save gallery generated output
     "min_reported_time": 10,
-    "subsection_order": ExplicitOrder(["../examples/layered", "../examples/graph"]),
+    "subsection_order": ExplicitOrder(["../examples/layered", "../examples/graph", "../examples/lenet_style"]),
     "within_subsection_order": ExampleTitleSortKey,
 }
 

diff --git a/docs/source/index.md b/docs/source/index.md
@@ -1,6 +1,6 @@
 # VisualTorch Documentation
 
-VisualTorch aims to help visualize Torch-based neural network architectures. It currently supports generating layered-style and graph-style architectures for PyTorch Sequential and Custom models.
+VisualTorch aims to help visualize Torch-based neural network architectures. It currently supports generating layered-style, graph-style, and LeNet-style architectures for PyTorch Sequential and Custom models.
 
 ![VisualTorch](_static/images/banners/visualizations-preview.png)
 
@@ -64,6 +64,7 @@ markdown/get_started/installation
 
 usage_examples/layered/index
 usage_examples/graph/index
+usage_examples/lenet_style/index
 ```
 
 ```{toctree}
@@ -72,6 +73,8 @@ usage_examples/graph/index
 
 markdown/api_references/layered
 markdown/api_references/graph
+markdown/api_references/lenet_style
+
 ```
 
 ```{toctree}

diff --git a/docs/source/markdown/api_references/lenet_style.md b/docs/source/markdown/api_references/lenet_style.md
@@ -0,0 +1,7 @@
+# LeNet Style View
+
+```{eval-rst}
+.. automodule:: visualtorch.lenet_style
+   :members:
+   :show-inheritance:
+```
diff --git a/visualtorch/__init__.py b/visualtorch/__init__.py
@@ -5,5 +5,6 @@
 
 from visualtorch.graph import graph_view
 from visualtorch.layered import layered_view
+from visualtorch.lenet_style import lenet_view
 
-__all__ = ["layered_view", "graph_view"]
+__all__ = ["layered_view", "graph_view", "lenet_view"]