diff --git a/.gitignore b/.gitignore index 451ce6f..b4e25e5 100644 --- a/.gitignore +++ b/.gitignore @@ -1,3 +1,6 @@ -__pycache__ +__pycache__/ +dist/ +build/ +torchsummaryX.egg-info/ .DS_Store *.swp diff --git a/README.md b/README.md index 0d63695..003a792 100644 --- a/README.md +++ b/README.md @@ -2,6 +2,8 @@ Improved visualization tool of [torchsummary](https://github.com/sksq96/pytorch-summary). Here, it visualizes kernel size, output shape, # params, and Mult-Adds. Also the torchsummaryX can handle RNN, Recursive NN, or model with multiple inputs. ## Usage +`pip install torchsummaryX` and + ```python from torchsummaryX import summary summary(your_model, torch.zeros((1, 3, 224, 224))) diff --git a/setup.py b/setup.py new file mode 100644 index 0000000..42c69ff --- /dev/null +++ b/setup.py @@ -0,0 +1,11 @@ +from setuptools import setup, find_packages + +setup( + name = "torchsummaryX", + version = "1.0.0", + description = "Improved visualization tool of torchsummary.", + author = "Namhyuk Ahn", + author_email = "nmhkahn@gmail.com", + url = "https://github.com/nmhkahn/torchsummaryX", + packages =["torchsummaryX"], +) diff --git a/torchsummaryX/__init__.py b/torchsummaryX/__init__.py new file mode 100644 index 0000000..ca59865 --- /dev/null +++ b/torchsummaryX/__init__.py @@ -0,0 +1 @@ +from .torchsummaryX import summary