HarisIqbal88 / PlotNeuralNet

Latex code for making neural networks diagrams
MIT License
21.74k stars 2.84k forks source link

Suggested change to_ConvConvRelu #42

Open guillermo-jimenez opened 5 years ago

guillermo-jimenez commented 5 years ago

Hi. I was working with your code (great, by the way!) and, while using "to_ConvConvRelu", I realized that it could be easily be extended to any number of Convs:

def to_ConvConvRelu( name, s_filer=256, n_filer=(64,64), offset="(0,0,0)", to="(0,0,0)", width=(2,2), height=40, depth=40, caption=" " ):
    return r"""
\pic[shift={ """+ offset +""" }] at """+ to +""" 
    {RightBandedBox={
        name="""+ name +""",
        caption="""+ caption +""",
        xlabel={{ """+ ''.join([str(n_filer[i]) + """, """ for i in range(len(n_filer)-1)])+ str(n_filer[-1]) +""" }},
        zlabel="""+ str(s_filer) +""",
        fill=\ConvColor,
        bandfill=\ConvReluColor,
        height="""+ str(height) +""",
        width={ """+ ''.join([str(width[i]) + """, """ for i in range(len(width)-1)])+ str(width[-1]) +""" },
        depth="""+ str(depth) +"""
        }
    };
"""

Sorry for the not really clean code; there's surely a way of coding it that is not as dirty but it does the trick.

Regards, Guille