Package: visualtorch
🔥 VisualTorch 🔥VisualTorchaims 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 fromvisualkeras,pytorchviz, andpytorch-summary.Note:VisualTorch may not yet support complex models, but contributions are welcome!v0.2: Added support for custom models and implemented graph view functionality.v0.1.1: Added support for the layered architecture of Torch Sequential.InstallationInstall from PyPI (Latest release)pipinstallvisualtorchInstall from sourcepipinstallgit+https://github.com/willyfh/visualtorchUsageSequentialimportvisualtorchimporttorch.nnasnn# Example of a simple CNN model using nn.Sequentialmodel=nn.Sequential(nn.Conv2d(3,16,kernel_size=3,padding=1),nn.ReLU(),nn.MaxPool2d(2,2),nn.Conv2d(16,32,kernel_size=3,padding=1),nn.ReLU(),nn.MaxPool2d(2,2),nn.Conv2d(32,64,kernel_size=3,padding=1),nn.ReLU(),nn.MaxPool2d(2,2),nn.Flatten(),nn.Linear(64*28*28,256),# Adjusted the input size for the Linear layernn.ReLU(),nn.Linear(256,10)# Assuming 10 output classes)input_shape=(1,3,224,224)visualtorch.layered_view(model,input_shape=input_shape,legend=True).show()# display using your system viewerCustom ModelIn a custom model, only the components defined within the model'sinitmethod are visualized. The operations that are defined exclusively within the forward function are not visualized.importtorch.nnasnnimporttorch.nn.functionalasFimportvisualtorch# Example of a simple CNN modelclassSimpleCNN(nn.Module):def__init__(self):super().__init__()self.conv1=nn.Conv2d(3,16,kernel_size=3,padding=1)self.conv2=nn.Conv2d(16,32,kernel_size=3,padding=1)self.conv3=nn.Conv2d(32,64,kernel_size=3,padding=1)self.fc1=nn.Linear(64*28*28,128)self.fc2=nn.Linear(128,10)defforward(self,x):x=self.conv1(x)x=F.relu(x)x=F.max_pool2d(x,2,2)x=self.conv2(x)x=F.relu(x)x=F.max_pool2d(x,2,2)x=self.conv3(x)x=F.relu(x)x=F.max_pool2d(x,2,2)print(x.shape)x=x.view(x.size(0),-1)x=self.fc1(x)x=F.relu(x)x=self.fc2(x)returnx# Create an instance of the SimpleCNNmodel=SimpleCNN()input_shape=(1,3,224,224)visualtorch.layered_view(model,input_shape=input_shape,legend=True).show()# display using your system viewerGraph Viewimporttorchimporttorch.nnasnnimportvisualtorchclassSimpleDense(nn.Module):def__init__(self):super(SimpleDense,self).__init__()self.h0=nn.Linear(4,8)self.h1=nn.Linear(8,8)self.h2=nn.Linear(8,4)self.out=nn.Linear(4,2)defforward(self,x):x=self.h0(x)x=self.h1(x)x=self.h2(x)x=self.out(x)returnxmodel=SimpleDense()input_shape=(1,4)visualtorch.graph_view(model,input_shape).show()Save the Imagevisualtorch.layered_view(model,input_shape=input_shape,legend=True,to_file='output.png')2D Viewvisualtorch.layered_view(model,input_shape=input_shape,draw_volume=False)Custom ColorUse 'fill' to change the color of the layer, and use 'outline' to change the color of the lines.fromcollectionsimportdefaultdictcolor_map=defaultdict(dict)color_map[nn.Conv2d]['fill']='LightSlateGray'# Light Slate Graycolor_map[nn.ReLU]['fill']='#87CEFA'# Light Sky Bluecolor_map[nn.MaxPool2d]['fill']='LightSeaGreen'# Light Sea Greencolor_map[nn.Flatten]['fill']='#98FB98'# Pale Greencolor_map[nn.Linear]['fill']='LightSteelBlue'# Light Steel Blueinput_shape=(1,3,224,224)visualtorch.layered_view(model,input_shape=input_shape,color_map=color_mapContributingPlease feel free to send a pull request to contribute to this project.LicenseThis poject is available as open source under the terms of theMIT License.Originally, this project was based on thevisualkeras(under the MIT license), with additional modifications inspired bypytorchviz, andpytorch-summary, both of which are also licensed under the MIT license.CitationPlease cite this project in your publications if it helps your research as follows:@misc{Hendria2024VisualTorch,author={Hendria, Willy Fitra},title={visualtorch},year={2024},publisher={GitHub},journal={GitHub repository},note={\url{https://github.com/willyfh/visualtorch}},}