用Python编写CNN的网络结构（便于进行网络分析和论文书写）

draw_convnet

Python script for illustrating Convolutional Neural Network (ConvNet)

Example image

With flag_omit=False

With flag_omit=True

代码：

from matplotlib.patches import Circle
from matplotlib.patches import Rectangle
from matplotlib.lines import Line2D
import os
import numpy as np
import matplotlib.pyplot as plt

plt.rcdefaults()  # 初始化

NumDots = 4
NumConvMax = 8
NumFcMax = 20

White = 1.
Light = 0.7
Medium = 0.5

Dark = 0.3
Darker = 0.15
Black = 0.


def add_layer(patches, colors, size=(24, 24), num=5,
              top_left=[0, 0],
              loc_diff=[3, -3],
              ):
    # add a rectangle
    top_left = np.array(top_left)
    loc_diff = np.array(loc_diff)
    loc_start = top_left - np.array([0, size[0]])
    for ind in range(num):
        patches.append(Rectangle(loc_start + ind * loc_diff, size[1], size[0]))
        if ind % 2:
            colors.append(Medium)
        else:
            colors.append(Light)


def add_layer_with_omission(
    patches, colors, size=(
        24, 24), num=5, num_max=8, num_dots=4, top_left=[
            0, 0], loc_diff=[
                3, -3], ):
    # add a rectangle
    top_left = np.array(top_left)
    loc_diff = np.array(loc_diff)
    loc_start = top_left - np.array([0, size[0]])
    t