Skip to content

Latest commit

 

History

History
62 lines (46 loc) · 1.83 KB

README.md

File metadata and controls

62 lines (46 loc) · 1.83 KB
Raw

draw_convnet

History

Python script for illustrating Convolutional Neural Network (ConvNet). This is based on https://github.com/gwding/draw_convnet, whose statement about using / citing the code is:

It is NOT required to cite anything to use the code.

If you are not facing space limitation and it does not break the flow of the paper, you might consider adding something like "This figure is generated by adapting the code from https://github.com/gwding/draw_convnet" (maybe in the footnote).

FYI, originally I used the code to generate the convnet figure in this paper "Automatic moth detection from trap images for pest management".

Plotting keras models

This fork allows to generate diagrams for keras sequential models, as follows:

from tensorflow.keras.models import Sequential
from tensorflow.keras.layers import InputLayer, Conv2D, MaxPool2D, AvgPool2D, Flatten, Dense

model = Sequential([
    InputLayer(input_shape=(32, 32, 3)),
    Conv2D(filters=32, kernel_size=3, padding="valid"),
    MaxPool2D(pool_size=2),
    Conv2D(filters=32, kernel_size=5, padding="same"),
    AvgPool2D(pool_size=2),
    Flatten(),
    Dense(units=124),
    Dense(units=10)
])
from draw_convnet import plot_keras_convnet

plot_keras_convnet(model, font_size=6)
plt.show()  # Not required in notebooks

Alternatively, if the goal is to save the model diagram to a file, the last lines of the previous code would become:

plot_keras_convnet(model, font_size=6, to_file="convnet_keras.svg")

At the moment, the code can only deal with Sequential models with the following layers:

  • InputLayer,
  • Conv2D,
  • MaxPool2D,
  • AvgPool2D,
  • Flatten,
  • Dense

Installation

Ideally, installation should be as simple as:

python -m pip install git+https://github.com/rtavenar/draw_convnet.git