Neural Network Layers

import matplotlib.pyplot as plt
import numpy as np
from mpl_toolkits.mplot3d import Axes3D
from matplotlib.colors import LinearSegmentedColormap

fig = plt.figure(figsize=(12, 8), facecolor='#ffffff')
ax = fig.add_subplot(111, projection='3d')
ax.set_facecolor('#ffffff')

models = 4
layers = 6
xpos = np.arange(models)
zpos = np.arange(layers)
xpos, zpos = np.meshgrid(xpos, zpos)
xpos = xpos.flatten()
zpos = zpos.flatten()
ypos = np.zeros_like(xpos)

dx = 0.6
dz = 0.55
np.random.seed(456)
dy = np.random.exponential(500, size=24) + 100

cmap = LinearSegmentedColormap.from_list('neon', ['#5314E6', '#F527B0', '#F5D327'])
norm = plt.Normalize(min(dy), max(dy))
colors = [cmap(norm(v)) for v in dy]

ax.bar3d(xpos, ypos, zpos, dx, dy, dz, color=colors, alpha=0.9, edgecolor='#e5e7eb', linewidth=0.4)

ax.set_xlabel('Model', fontsize=11, color='#1f2937', labelpad=12)
ax.set_ylabel('Parameters (M)', fontsize=11, color='#1f2937', labelpad=12)
ax.set_zlabel('Layer', fontsize=11, color='#1f2937', labelpad=10)
ax.set_title('Neural Network Parameter Distribution', fontsize=14, color='#1f2937', fontweight='bold', pad=20)

ax.set_xticks(range(4))
ax.set_xticklabels(['ResNet', 'ViT', 'GPT', 'BERT'], fontsize=8, color='#374151')
ax.set_zticks(range(6))
ax.set_zticklabels(['Embed', 'Attn1', 'FF1', 'Attn2', 'FF2', 'Head'], fontsize=7, color='#374151')
ax.tick_params(colors='#374151', labelsize=9)

ax.xaxis.pane.fill = False
ax.yaxis.pane.fill = False
ax.zaxis.pane.fill = False
ax.xaxis.pane.set_edgecolor('#e5e7eb')
ax.yaxis.pane.set_edgecolor('#e5e7eb')
ax.zaxis.pane.set_edgecolor('#e5e7eb')
ax.grid(True, alpha=0.3, linewidth=0.5)

ax.view_init(elev=20, azim=40)
plt.tight_layout()
plt.show()