Skip to main content
Donate to the Python Software Foundation or Purchase a PyCharm License to Benefit the PSF! Donate Now

A machine learning figure generation library

Project description

AI Figures


AIFIG is a python library for generating figures of machine learning models.

The libary allows you to generate figures such as the following, which may be useful for use in presentations, papers etc.

AIFIG is a refactored version of some of my personal code. Functionality will naturally be limited and not suited for every use. I encourage anyone who is interested to contribute with additional features.

If you use AIFIG in a paper, you can cite the library like this (bibtex):

  author = {Sigve Rokenes},
  title = {AI-FIG},
  year = {2019},
  publisher = {GitHub},
  journal = {GitHub repository},
  howpublished = {\url{}}


AI-FIG library with svg export:

pip install aifig

If you need to export as png or pdf:

pip install svglib


Simple example

# Import library
import aifig

# Create new figure, title and author is optional
my_figure = aifig.figure("Figure 1", "Sigve Rokenes")

# Figures consist of graphs (eg. each network in a model)
my_graph = aifig.graph("gen")

# Graphs contain elements (inputs, outputs, layers)
my_graph.add(aifig.dense("input", 16))
my_graph.add(aifig.dense("hidden_1", 64))
my_graph.add(aifig.dense("hidden_2", 128))
my_graph.add(aifig.dense("hidden_3", 64))
my_graph.add(aifig.dense("output", 1))

# Add the graph to the figure at position (0,0)
my_figure.add(graph, 0, 0)

# Save the figure 
my_figure.save_png("my_figure.png", scale=1)

The above code generates this figure:

Multi-graph example (GAN model)

import aifig
figure = aifig.figure()

# Define generator network
generator_elements = [
  aifig.dense("noise_vector", 128, comment="norm_dist", simple=True),
  aifig.conv("tconv_1", 48, comment="5x5"),
  aifig.conv("tconv_2", 32, comment="5x5"),
  aifig.conv("tconv_3", 8, comment="5x5"),
  aifig.conv("tconv_4", 3, comment="5x5"),
  aifig.image("gen_result", comment="(fake image)")

# Define discriminator network
discriminator_elements = [
  aifig.image("image_input", comment="real/fake"),
  aifig.conv("conv_1", 16, comment="5x5"),
  aifig.conv("conv_2", 32, comment="5x5"),
  aifig.conv("conv_3", 48, comment="5x5"),
  aifig.dense("dense_1", 64),
  aifig.dense("output", 1),
  aifig.arrow("prediction", comment="log prob")

# Create graphs with elements
gen_graph = aifig.graph("gen", generator_elements)
dsc_graph = aifig.graph("dsc", discriminator_elements)
dat_graph = aifig.graph("dat", [aifig.image("real_image", comment="(dataset)")])

# Add graphs to figure
figure.add(gen_graph, 0, 0)
figure.add(dat_graph, 1, 0)
figure.add(dsc_graph, 0, 1)

# Connect inputs to discriminator network
figure.connect("gen", "dsc")
figure.connect("dat", "dsc")

# Save figure as png

This code generates the following figure:


A figure consists of one or more graphs. These graphs are placed in a grid using figure.add(graph, x, y). You can add elements to graphs using mygraph.add(element), and you can connect graphs with arrows using figure.connect("graph_name1", "graph_name2"). Finally, to save a figure, use my_figure.save_svg("fig.svg") or variants for different formats.

# ===================== #
#        Figure         #
# ===================== #
#  title			figure title
#  author			figure author
my_figure = aifig.figure()

# figure.add
#  graph			graph to add
#  x				x position in grid
#  y				y position in grid
my_figure.add(graph, 0, 0)

# figure.connect
#  from				name of first graph
#  to				name of second graph
#  position			grid position of arrow, use this if 
#				different arrows overlap
#  offset			arrow offset in units, useful to
#				distinguish different arrows at same position
my_figure.connect("graph1", "graph2")

# (path)
#  path				file path to save to
#  scale			upscale (png only)
#  debug			enable debug draw mode
my_figure.save_png("my_figure.png", scale=1)

# ===================== #
#         Graph         #
# ===================== #
# 	name			(required)
# 	elements		[list of elements]
# 	spacing			(between elements, default 32)

my_graph = aifig.graph("graph_name")

# ===================== #
#    Layer elements     #
# ===================== #
# 	label			text label, use None to hide
# 	size			size of layer (nodes, filters)
#	comment			additional comment text
#	size_label		set to False to hide size label
#	simple			(dense only) set True to render as simple rectangle

dense = aifig.dense()		# Dense (fully connected)
conv  = aifig.conv()		# Convolutional layer

# ===================== #
#    Simple elements    #
# ===================== #
#	label			text label, use None to hide
#	comment			additional comment text

pool  = aifig.pool()		# Pooling layer
image = aifig.image()		# Image (usually input)
arrow = aifig.arrow()		# Arrow

# ===================== #
#    Special elements   #
# ===================== #
# 	width			width of padding (use negative to reduce)

padding = aifig.padding(10) 


  • svgwrite
  • svglib (only to save as pdf/png)
  • reportlab (only to save as pdf/png)

Project details

Download files

Download the file for your platform. If you're not sure which to choose, learn more about installing packages.

Filename, size & hash SHA256 hash help File type Python version Upload date
aifig-0.1.6-py3-none-any.whl (9.8 kB) Copy SHA256 hash SHA256 Wheel py3
aifig-0.1.6.tar.gz (10.4 kB) Copy SHA256 hash SHA256 Source None

Supported by

Elastic Elastic Search Pingdom Pingdom Monitoring Google Google BigQuery Sentry Sentry Error logging AWS AWS Cloud computing DataDog DataDog Monitoring Fastly Fastly CDN SignalFx SignalFx Supporter DigiCert DigiCert EV certificate StatusPage StatusPage Status page