200 lines
4.2 KiB
Python
200 lines
4.2 KiB
Python
# subter_lenet_arch.py
|
||
# Requires running from inside the PlotNeuralNet repo, like: python3 ../subter_lenet_arch.py
|
||
import sys, argparse
|
||
|
||
sys.path.append("../") # import pycore from repo root
|
||
|
||
from pycore.tikzeng import *
|
||
|
||
parser = argparse.ArgumentParser()
|
||
parser.add_argument("--rep_dim", type=int, default=1024, help="latent size for FC")
|
||
args = parser.parse_args()
|
||
REP = int(args.rep_dim)
|
||
|
||
# Visual scales so the huge width doesn't dominate the figure
|
||
H32, H16, H8, H1 = 26, 18, 12, 1
|
||
D2048, D1024, D512, D256, D128, D1 = 52, 36, 24, 12, 6, 1
|
||
W1, W4, W8, W16, W32 = 1, 2, 2, 4, 8
|
||
|
||
|
||
arch = [
|
||
to_head(".."),
|
||
to_cor(),
|
||
to_begin(),
|
||
to_Conv(
|
||
"latent",
|
||
n_filer="",
|
||
s_filer="latent dim",
|
||
offset="(2,0,0)",
|
||
to="(0,0,0)",
|
||
height=H8 * 1.6,
|
||
depth=D1,
|
||
width=W1,
|
||
caption=f"Latent Space",
|
||
),
|
||
# to_connection("fc1", "latent"),
|
||
# --------------------------- DECODER ---------------------------
|
||
# FC back to 16384
|
||
to_fc(
|
||
"fc3",
|
||
n_filer="{{8×128×8}}",
|
||
zlabeloffset=0.5,
|
||
offset="(2,0,0)",
|
||
to="(latent-east)",
|
||
height=H1,
|
||
depth=D512,
|
||
width=W1,
|
||
caption=f"FC",
|
||
),
|
||
to_Conv(
|
||
"unsqueeze",
|
||
s_filer="{{128×8}}",
|
||
zlabeloffset=0.4,
|
||
n_filer=32,
|
||
offset="(2,0,0)",
|
||
to="(fc3-east)",
|
||
height=H8,
|
||
depth=D128,
|
||
width=W32,
|
||
caption="Unsqueeze",
|
||
),
|
||
# to_connection("latent", "fc3"),
|
||
# Reshape to 4×8×512
|
||
to_UnPool(
|
||
"up1",
|
||
offset="(2,0,0)",
|
||
n_filer=32,
|
||
to="(unsqueeze-east)",
|
||
height=H16,
|
||
depth=D256,
|
||
width=W32,
|
||
caption="",
|
||
),
|
||
to_Conv(
|
||
"dwdeconv1",
|
||
s_filer="",
|
||
n_filer=1,
|
||
offset="(0,0,0)",
|
||
to="(up1-east)",
|
||
height=H16,
|
||
depth=D256,
|
||
width=W1,
|
||
caption="Deconv1",
|
||
),
|
||
to_Conv(
|
||
"dwdeconv2",
|
||
s_filer="{{256×16}}",
|
||
zlabeloffset=0.4,
|
||
n_filer=32,
|
||
offset="(0,0,0)",
|
||
to="(dwdeconv1-east)",
|
||
height=H16,
|
||
depth=D256,
|
||
width=W32,
|
||
caption="",
|
||
),
|
||
to_UnPool(
|
||
"up2",
|
||
offset="(2,0,0)",
|
||
to="(dwdeconv2-east)",
|
||
n_filer=32,
|
||
height=H16,
|
||
depth=D1024,
|
||
width=W32,
|
||
caption="",
|
||
),
|
||
to_Conv(
|
||
"dwdeconv3",
|
||
s_filer="",
|
||
n_filer=1,
|
||
offset="(0,0,0)",
|
||
to="(up2-east)",
|
||
height=H16,
|
||
depth=D1024,
|
||
width=W1,
|
||
caption="Deconv2",
|
||
),
|
||
to_Conv(
|
||
"dwdeconv4",
|
||
s_filer="{{1024×16}}",
|
||
zlabeloffset=0.17,
|
||
n_filer=16,
|
||
offset="(0,0,0)",
|
||
to="(dwdeconv3-east)",
|
||
height=H16,
|
||
depth=D1024,
|
||
width=W16,
|
||
caption="",
|
||
),
|
||
to_UnPool(
|
||
"up3",
|
||
offset="(2,0,0)",
|
||
n_filer=16,
|
||
to="(dwdeconv4-east)",
|
||
height=H32,
|
||
depth=D2048,
|
||
width=W16,
|
||
caption="",
|
||
),
|
||
to_Conv(
|
||
"dwdeconv5",
|
||
s_filer="",
|
||
n_filer=1,
|
||
offset="(0,0,0)",
|
||
to="(up3-east)",
|
||
height=H32,
|
||
depth=D2048,
|
||
width=W1,
|
||
caption="Deconv3",
|
||
),
|
||
to_Conv(
|
||
"dwdeconv6",
|
||
s_filer="{{2048×32}}",
|
||
zlabeloffset=0.15,
|
||
n_filer=8,
|
||
offset="(0,0,0)",
|
||
to="(dwdeconv5-east)",
|
||
height=H32,
|
||
depth=D2048,
|
||
width=W8,
|
||
caption="",
|
||
),
|
||
to_Conv(
|
||
"outconv",
|
||
s_filer="{{2048×32}}",
|
||
zlabeloffset=0.15,
|
||
n_filer=1,
|
||
offset="(2,0,0)",
|
||
to="(dwdeconv6-east)",
|
||
height=H32,
|
||
depth=D2048,
|
||
width=W1,
|
||
caption="Deconv4",
|
||
),
|
||
# to_connection("up2", "deconv2"),
|
||
# Output
|
||
to_Conv(
|
||
"out",
|
||
s_filer="{{2048×32}}",
|
||
zlabeloffset=0.15,
|
||
n_filer=1,
|
||
offset="(2,0,0)",
|
||
to="(outconv-east)",
|
||
height=H32,
|
||
depth=D2048,
|
||
width=W1,
|
||
caption="Output",
|
||
),
|
||
# to_connection("deconv2", "out"),
|
||
to_end(),
|
||
]
|
||
|
||
|
||
def main():
|
||
name = "subter_lenet_arch"
|
||
to_generate(arch, name + ".tex")
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main()
|