Package: pytorch2tikz
Pytorch2TikzGenerate Tikz figures for neural networks implemented in pytorch. It uses LaTeX snippets fromPlotNeuralNetbut you can now just run your network to plot everything automatically. For examples see./examples.Pytorch2TikzExampleGetting StartedInterfaceArchitectureMethodsMethodsBlockArgumentsContributionsLayer supportCustom ConnectionColorsExamplefrompytorch2tikzimportArchitectureprint('Load model')model=vgg16(True)print('Load data')...print('Init architecture')arch=Architecture(model)print('Run model')withtorch.inference_mode():forimage,_indata_loader:image=image.to(device,non_blocking=True)output=model(image)print('Write result to out.tex')arch.save('out.tex')Getting Startedpip install pytorch2tikzInterfaceArchitectureArchitecture(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)MethodsArgumentdescriptionmoduleis the Model to plotblock_offsetoffset to the next block; A block is created when the input dimensions changeheight_depth_factorscale the change of the next layer (last 2 dimensions); typically used to make the network a bit more compactwidth_factorscale the change of the next layer (first dimension); typically used to make the network a bit more compactlinear_factorused when there is a drastic change in the last dimension (e.g. moving from conv to linear layers)image_pathoutput path for recognized input images.{i}gets replaced by the current layer indexignore_layersdefine layers that should not be plotted. This can be a list of any substring of thetype(class)(e.g. torch.nn.modules.batchnorm.BatchNorm)colorsenum of colors. For an example check out./pytorch2tikz/constantsMethodsdefget_block(self,name:str)->Block:...get a specific block to alter its propertiesdefget_tex(self)->str:...generate the tex codedefsave(self,file_path:str):...generate and save the tex code to the given pathBlockBlock(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)ArgumentsArgumentDescriptionnamearbitrary name of the block. Should be unique, and typically the layers id is usedfillfilling color as hex string, e.g.#000000bandfillfilling of subcolor at the right end of a box.pictypeshould bePICTYPE.RIGHTBANDEDBOXot be displayedpictypeone of [PICTYPE.BOX,PICTYPE.RIGHTBANDEDBOX]opacityopacity of the fillingsizesize of the boxdefault_sizeSize used for dimensions which are "flat": e.g. for 1D inputs the size (default, default, size) is used.dimdimensionality of the block, e.g. 1 for Linear layers, 3 for conv2d layers (channels x dim1 x dim2)scale_factorscale factors to alter the size when outputting tex to make the figure more compactoffsetoffset to the references position/block intotoposition tuple or block used for relative positioningcapitioncaption of the block. Use an empty string if no caption is wantedxlabeldisplay label for 1st dimensionylabeldisplay label for 2nd dimensionzlabeldisplay label for 3nd dimensionContributionsThank you for share your improvements to this package!Layer supportPlease don't hesitate to add blocks for unsupported layers underpytorch2tikz/block/D<x>.pywithxbeing the dimensionality of your layer. If your layer exists for multiple dimensions, chooseDn.py:add your block definition underpytorch2tikz/block/D<x>.pyadd mapping of type string topytorch2tikz/mapping.pyadd your color topytorch2tikz/constants.py(seeColors)Custom ConnectionFor custom connections that can be added in postprocessing of an architecture like residual connections, add your desired connection inpytorch2tikz/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).ColorsColors are defined inpytorch2tikz/constants.py. For each color there must exist an entry in the enumCOLORand the defined value in the DictCOLOR_VALUES. Make sure your color is easily distinguishable from other layers.