如何用Python绘制分层布局的神经网络拓扑图

Nice question! The default NetworkX layouts don't handle layered neural network structures well out of the box, but we can fix this either by using NetworkX's built-in multipartite layout or by manually defining node positions to get that clean input-left, hidden-middle, output-right structure.

Method 1: Using multipartite_layout (Easiest) #

This method leverages NetworkX's multipartite_layout which automatically arranges nodes into columns based on a "layer" attribute we assign to each node group.

import networkx as nx
import matplotlib.pyplot as plt

# Replace these with your actual node lists
ixx = ["Input_1", "Input_2", "Input_3"]  # Input nodes
o = ["Hidden_A", "Hidden_B"]             # Hidden nodes
WIN = ["Output"]                         # Output node

# Initialize directed graph
g = nx.DiGraph()

# Add nodes with layer attribute to define grouping
g.add_nodes_from(ixx, layer=0)  # Input layer (leftmost, column 0)
g.add_nodes_from(o, layer=1)    # Hidden layer (middle, column 1)
g.add_nodes_from(WIN, layer=2)  # Output layer (rightmost, column 2)

# Add edges: All input nodes connect to all hidden nodes
for input_node in ixx:
    for hidden_node in o:
        g.add_edge(input_node, hidden_node)

# Add edges: All hidden nodes connect to output node
for hidden_node in o:
    g.add_edge(hidden_node, WIN[0], weight=10)

# Generate layered positions
# `subset_key='layer'` tells the layout to use our layer attribute
# `align='horizontal'` ensures nodes are arranged in vertical columns
pos = nx.multipartite_layout(g, subset_key='layer', align='horizontal')

# Draw the graph with custom styling
nx.draw(
    g,
    pos,
    with_labels=True,
    node_size=1800,
    node_color="#a8d1ff",
    font_size=11,
    font_weight="bold",
    arrows=True
)

# Optional: Draw edge weights if needed
edge_labels = nx.get_edge_attributes(g, "weight")
nx.draw_networkx_edge_labels(g, pos, edge_labels=edge_labels)

plt.title("Layered Neural Network Graph")
plt.show()

Method 2: Manual Positioning (Full Control) #

If you want exact control over node positions (e.g., adjusting vertical spacing), you can manually assign (x, y) coordinates to each node.

import networkx as nx
import matplotlib.pyplot as plt

# Your node lists
ixx = ["Input_1", "Input_2", "Input_3"]
o = ["Hidden_A", "Hidden_B"]
WIN = ["Output"]

g = nx.DiGraph()
g.add_nodes_from(ixx + o + WIN)

# Add edges same as before
for input_node in ixx:
    for hidden_node in o:
        g.add_edge(input_node, hidden_node)
for hidden_node in o:
    g.add_edge(hidden_node, WIN[0], weight=10)

# Manually define positions: x=0 (input), x=1 (hidden), x=2 (output)
pos = {}

# Input layer (x=0, evenly spaced vertically)
for idx, node in enumerate(ixx):
    pos[node] = (0, len(ixx) - 1 - idx)  # Reverse index to top-align if needed

# Hidden layer (x=1)
for idx, node in enumerate(o):
    pos[node] = (1, len(o) - 1 - idx)

# Output layer (x=2, centered vertically)
pos[WIN[0]] = (2, (len(ixx)-1)/2)  # Match midpoint of input layer

# Draw the graph
nx.draw(
    g,
    pos,
    with_labels=True,
    node_size=1800,
    node_color="#a8d1ff",
    font_size=11,
    font_weight="bold",
    arrows=True
)

nx.draw_networkx_edge_labels(g, pos, edge_labels=nx.get_edge_attributes(g, "weight"))
plt.title("Custom Layered Neural Network Graph")
plt.show()

Both methods will give you the clean, structured layout you're looking for. The multipartite_layout is great for quick setup, while manual positioning lets you tweak every detail.

内容的提问来源于stack exchange，提问作者Shilpa