110 lines
5.0 KiB
Python
110 lines
5.0 KiB
Python
from graphviz import Digraph
|
|
import os
|
|
|
|
|
|
class layer:
|
|
def __init__(self, graph, name, size, color):
|
|
self.name = name
|
|
self.size = size
|
|
self.color = color
|
|
self.graph = graph
|
|
|
|
def draw():
|
|
pass
|
|
|
|
class input_layer(layer):
|
|
def __init__(self, graph, size):
|
|
super().__init__(graph, f"Input Layer({size})", size)
|
|
self.graph.node(self.name, shape='circle', style='filled', fillcolor=self.color, label=" ")
|
|
self.graph.attr(label=f'{self.name} Layer({self.size})', fontname='Times New Roman', fontweight='bold', fontsize='36')
|
|
|
|
|
|
def draw_neural_net(input_size, hidden_sizes, num_classes, show_hidden=3):
|
|
g = Digraph('G', filename='neural_network', format='png')
|
|
g.attr(rankdir='LR', size='10,8', nodesep='1', ranksep='2', bgcolor='transparent', dpi='300')
|
|
|
|
# Input layer
|
|
with g.subgraph(name='cluster_input') as c:
|
|
c.attr(color='white')
|
|
for i in range(input_size):
|
|
c.node(f'input_{i}', shape='circle', style='filled', fillcolor='darkorange:orange', label=" ")
|
|
c.attr(label=f'Input Layer({input_size})', fontname='Times New Roman', fontweight='bold', fontsize='36')
|
|
|
|
# Hidden layers
|
|
previous_layer = 'input'
|
|
previous_layer_size = input_size
|
|
for layer_idx, hidden_size in enumerate(hidden_sizes):
|
|
with g.subgraph(name=f'cluster_hidden_{layer_idx}') as c:
|
|
c.attr(color='white')
|
|
for i in range(show_hidden):
|
|
c.node(f'hidden_{layer_idx}_{i}', shape='circle', style='filled', fillcolor='darkgreen:lightgreen', label=" ")
|
|
if hidden_size > show_hidden * 2:
|
|
c.node(f'ellipsis_{layer_idx}', shape='plaintext', label='...')
|
|
for i in range(hidden_size - show_hidden, hidden_size):
|
|
c.node(f'hidden_{layer_idx}_{i}', shape='circle', style='filled', fillcolor='darkgreen:lightgreen', label=" ")
|
|
c.attr(label=f'Hidden Layer {layer_idx + 1}({hidden_size})', fontname='Times New Roman', fontweight='bold', fontsize='36')
|
|
|
|
# Add edges from previous layer to current hidden layer
|
|
if layer_idx == 0: # Only connect input layer to first hidden layer
|
|
for i in range(previous_layer_size):
|
|
for j in range(show_hidden):
|
|
g.edge(f'{previous_layer}_{i}', f'hidden_{layer_idx}_{j}')
|
|
for j in range(hidden_size - show_hidden, hidden_size):
|
|
g.edge(f'{previous_layer}_{i}', f'hidden_{layer_idx}_{j}')
|
|
else:
|
|
for i in range(show_hidden):
|
|
for j in range(show_hidden):
|
|
g.edge(f'hidden_{layer_idx - 1}_{i}', f'hidden_{layer_idx}_{j}')
|
|
for j in range(hidden_size - show_hidden, hidden_size):
|
|
g.edge(f'hidden_{layer_idx - 1}_{i}', f'hidden_{layer_idx}_{j}')
|
|
for i in range(hidden_size - show_hidden, hidden_size):
|
|
for j in range(show_hidden):
|
|
g.edge(f'hidden_{layer_idx - 1}_{i}', f'hidden_{layer_idx}_{j}')
|
|
for j in range(hidden_size - show_hidden, hidden_size):
|
|
g.edge(f'hidden_{layer_idx - 1}_{i}', f'hidden_{layer_idx}_{j}')
|
|
|
|
previous_layer = f'hidden_{layer_idx}'
|
|
previous_layer_size = hidden_size
|
|
|
|
# Output layer
|
|
with g.subgraph(name='cluster_output') as c:
|
|
c.attr(color='white')
|
|
for i in range(num_classes):
|
|
c.node(f'output_{i}', shape='circle', style='filled', fillcolor='darkorange:orange', label=" ")
|
|
c.attr(label=f'Output Layer({num_classes})', fontname='Times New Roman', fontweight='bold', fontsize='36')
|
|
|
|
# Add edges from last hidden layer to output layer
|
|
# for i in range(previous_layer_size):
|
|
# for j in range(num_classes):
|
|
# g.edge(f'{previous_layer}_{i}', f'output_{j}')
|
|
# # Add edges
|
|
# # Add edges from input to visible hidden nodes
|
|
# for i in range(input_size):
|
|
# for j in range(show_hidden):
|
|
# g.edge(f'input_{i}', f'hidden_{j}')
|
|
# for i in range(input_size):
|
|
# for j in range(hidden_size - show_hidden, hidden_size):
|
|
# g.edge(f'input_{i}', f'hidden_{j}')
|
|
|
|
# # Add edges from visible hidden nodes to output layer
|
|
# for i in range(show_hidden):
|
|
# for j in range(num_classes):
|
|
# g.edge(f'hidden_{i}', f'output_{j}')
|
|
# for i in range(hidden_size - show_hidden, hidden_size):
|
|
# for j in range(num_classes):
|
|
# g.edge(f'hidden_{i}', f'output_{j}')
|
|
# Add edges from last hidden layer to output layer
|
|
for i in range(show_hidden):
|
|
for j in range(num_classes):
|
|
g.edge(f'{previous_layer}_{i}', f'output_{j}')
|
|
for i in range(previous_layer_size - show_hidden, previous_layer_size):
|
|
for j in range(num_classes):
|
|
g.edge(f'{previous_layer}_{i}', f'output_{j}')
|
|
|
|
return g
|
|
|
|
if __name__ == '__main__':
|
|
g = draw_neural_net(7, [60, 60], 7)
|
|
output_path = g.render(view=False)
|
|
print(output_path)
|
|
os.system(f'explorer.exe neural_network.png') |