Pytorch2Tikz

Generate Tikz figures for neural networks implemented in pytorch. It uses LaTeX snippets from PlotNeuralNet but you can now just run your network to plot everything automatically. For examples see ./examples.

• Pytorch2Tikz
  • Example
  • Getting Started
  • Interface
    • Architecture
      • Methods
      • Methods
    • Block
    • Arguments
  • Contributions
    • Layer support
    • Custom Connection
    • Colors

Example

from pytorch2tikz import Architecture

print('Load model')
model = vgg16(True)

print('Load data')
...

print('Init architecture')
arch = Architecture(model)

print('Run model')
with torch.inference_mode():
    for image, _ in data_loader:
        image = image.to(device, non_blocking=True)
        output = model(image)

print('Write result to out.tex')
arch.save('out.tex')

Getting Started

pip install pytorch2tikz

Interface

Architecture

Architecture(module: nn.Module,
            block_offset=8,
            height_depth_factor=0.8,
            width_factor=0.8,
            linear_factor=0.8,
            image_path='./input_{i}.png',
            ignore_layers=['batchnorm', 'flatten'],
            colors=COLOR_VALUES)

Methods

Methods

def get_block(self, name: str) -> Block:
    ...

get a specific block to alter its properties

def get_tex(self) -> str:
    ...

generate the tex code

    
def save(self, file_path: str):
    ...

generate and save the tex code to the given path

Block

Block(name,
    fill: COLOR = COLOR.LINEAR,
    bandfill: COLOR = None,
    pictype = PICTYPE.BOX,
    opacity = 0.7,
    size = (10,40,40),
    default_size = DEFAULT_VALUE,
    dim = 3,
    scale_factor = np.zeros(3),
    offset: Tuple[int] = (0,0,0),
    to: Union[Tuple[int], Block] = (0,0,0),
    caption = " ",
    xlabel = True,
    ylabel = False,
    zlabel = True)

Arguments

Contributions

Thank you for share your improvements to this package!

Layer support

Please don't hesitate to add blocks for unsupported layers under pytorch2tikz/block/D<x>.py with x being the dimensionality of your layer. If your layer exists for multiple dimensions, choose Dn.py:

1. add your block definition under pytorch2tikz/block/D<x>.py
2. add mapping of type string to pytorch2tikz/mapping.py
3. add your color to pytorch2tikz/constants.py (see Colors)

Custom Connection

For custom connections that can be added in postprocessing of an architecture like residual connections, add your desired connection in pytorch2tikz/block/connections.py. See the examples there as a guidance. For existing connections there are a bunch of defined positions for each block:

Each position can be combined with (padding-)(near|far)(north|south)(east|west).

Colors

Colors are defined in pytorch2tikz/constants.py. For each color there must exist an entry in the enum COLOR and the defined value in the Dict COLOR_VALUES. Make sure your color is easily distinguishable from other layers.