# W, H, C, N, K, S, R, Wpad, Hpad, Wstride, Hstride
layer_sizes = {
'cifar-10': 
	{
		"vgg": [
		    # Conv1
		    (28, 28, 3, 1, 64, 3, 3, 1, 1, 1, 1),
		    (28, 28, 64, 1, 64, 3, 3, 1, 1, 1, 1),
		    # # Conv2
		    (14, 14, 64, 1, 128, 3, 3, 1, 1, 1, 1),
		    (14, 14, 128, 1, 128, 3, 3, 1, 1, 1, 1),
		    # # Conv3
		    (7, 7, 128, 1, 256, 3, 3, 1, 1, 1, 1),
		    (7, 7, 256, 1, 256, 3, 3, 1, 1, 1, 1),
		    (7, 7, 256, 1, 256, 3, 3, 1, 1, 1, 1),
		    # # Conv4
		    (4, 4, 256, 1, 512, 3, 3, 1, 1, 1, 1),
		    (4, 4, 512, 1, 512, 3, 3, 1, 1, 1, 1),
		    (4, 4, 512, 1, 512, 3, 3, 1, 1, 1, 1),
		    # # Conv5
		    (2, 2, 512, 1, 512, 3, 3, 1, 1, 1, 1),
		    (2, 2, 512, 1, 512, 3, 3, 1, 1, 1, 1),
		    (2, 2, 512, 1, 512, 3, 3, 1, 1, 1, 1),
		    # # FC
		    (1, 1, 512, 1, 512, 1, 1, 0, 0, 1, 1),
		    (1, 1, 512, 1, 10, 1, 1, 0, 0, 1, 1)
		],
		'dense': [
		    # Conv1 # too dense
		    (28, 28, 3, 1, 16, 3, 3, 1, 1, 1, 1),
		    
		    # Block 1-1, 1-2, 1-3 
		    (28, 28, 16, 1, 24, 3, 3, 1, 1, 1, 1),
		    (28, 28, 40, 1, 24, 3, 3, 1, 1, 1, 1),
		    (28, 28, 64, 1, 24, 3, 3, 1, 1, 1, 1),
		    # Block 1-4, 1-5, 1-6, 1-7
		    (28, 28, 88, 1, 24, 3, 3, 1, 1, 1, 1),
		    (28, 28, 112, 1, 24, 3, 3, 1, 1, 1, 1),
		    (28, 28, 136, 1, 24, 3, 3, 1, 1, 1, 1),
		    (28, 28, 160, 1, 24, 3, 3, 1, 1, 1, 1),
		    # Block 1-8, 1-9, 1-10, 1-11
		    (28, 28, 184, 1, 24, 3, 3, 1, 1, 1, 1),
		    (28, 28, 208, 1, 24, 3, 3, 1, 1, 1, 1),
		    (28, 28, 232, 1, 24, 3, 3, 1, 1, 1, 1),
		    (28, 28, 256, 1, 256, 1, 1, 0, 0, 1, 1),

		    # Conv 2-1, 2-2, 2-3 
		    (14, 14, 256, 1, 24, 3, 3, 1, 1, 1, 1),
		    (14, 14, 280, 1, 24, 3, 3, 1, 1, 1, 1),
		    (14, 14, 304, 1, 24, 3, 3, 1, 1, 1, 1),
		    # Block 2-4, 2-5, 2-6, 2-7
		    (14, 14, 328, 1, 24, 3, 3, 1, 1, 1, 1),
		    (14, 14, 352, 1, 24, 3, 3, 1, 1, 1, 1),
		    (14, 14, 376, 1, 24, 3, 3, 1, 1, 1, 1),
		    (14, 14, 400, 1, 24, 3, 3, 1, 1, 1, 1),
		    # Block 2-8, 2-9, 2-10, 2-11
		    (14, 14, 424, 1, 24, 3, 3, 1, 1, 1, 1),
		    (14, 14, 448, 1, 24, 3, 3, 1, 1, 1, 1),
		    (14, 14, 472, 1, 24, 3, 3, 1, 1, 1, 1),
		    (14, 14, 496, 1, 496, 1, 1, 0, 0, 1, 1),

		    # Conv 3-1, 3-2, 3-3 
		    (7, 7, 496, 1, 24, 3, 3, 1, 1, 1, 1),
		    (7, 7, 520, 1, 24, 3, 3, 1, 1, 1, 1),
		    (7, 7, 544, 1, 24, 3, 3, 1, 1, 1, 1),
		    # Block 3-4, 3-5, 3-6, 3-7
		    (7, 7, 568, 1, 24, 3, 3, 1, 1, 1, 1),
		    (7, 7, 592, 1, 24, 3, 3, 1, 1, 1, 1),
		    (7, 7, 616, 1, 24, 3, 3, 1, 1, 1, 1),
		    (7, 7, 640, 1, 24, 3, 3, 1, 1, 1, 1),
		    # Block 3-8, 3-9, 3-10
		    (7, 7, 664, 1, 24, 3, 3, 1, 1, 1, 1),
		    (7, 7, 688, 1, 24, 3, 3, 1, 1, 1, 1),
		    (7, 7, 712, 1, 24, 3, 3, 1, 1, 1, 1),
		    # FC1
		    (1, 1, 736, 1, 10, 1, 1, 0, 0, 1, 1)
		],
		'wrn': [
		    # Conv1 # too dense
		    (28, 28, 3, 1, 16, 3, 3, 1, 1, 1, 1),
		    
		    # Conv 2-1, 2-2, 2-3 
		    (28, 28, 16, 1, 160, 3, 3, 1, 1, 1, 1),
		    (28, 28, 160, 1, 160, 3, 3, 1, 1, 1, 1),
		    (28, 28, 16, 1, 160, 1, 1, 0, 0, 1, 1), # bottlneck
		    # Conv 2-4, 2-5, 2-6, 2-7, 2-8, 2-9
		    (28, 28, 160, 1, 160, 3, 3, 1, 1, 1, 1),
		    (28, 28, 160, 1, 160, 3, 3, 1, 1, 1, 1),
		    (28, 28, 160, 1, 160, 3, 3, 1, 1, 1, 1),
		    (28, 28, 160, 1, 160, 3, 3, 1, 1, 1, 1),
		    (28, 28, 160, 1, 160, 3, 3, 1, 1, 1, 1),
		    (28, 28, 160, 1, 160, 3, 3, 1, 1, 1, 1),

		    # Conv 3-1, 3-2, 3-3 
		    (14, 14, 160, 1, 320, 3, 3, 1, 1, 1, 1),
		    (14, 14, 320, 1, 320, 3, 3, 1, 1, 1, 1),
		    (14, 14, 160, 1, 320, 1, 1, 0, 0, 1, 1), # bottlneck
		    # Conv 3-4, 3-5, 3-6, 3-7, 3-8, 3-9
		    (14, 14, 320, 1, 320, 3, 3, 1, 1, 1, 1),
		    (14, 14, 320, 1, 320, 3, 3, 1, 1, 1, 1),
		    (14, 14, 320, 1, 320, 3, 3, 1, 1, 1, 1),
		    (14, 14, 320, 1, 320, 3, 3, 1, 1, 1, 1),
		    (14, 14, 320, 1, 320, 3, 3, 1, 1, 1, 1),
		    (14, 14, 320, 1, 320, 3, 3, 1, 1, 1, 1),

		    # Conv 4-1, 4-2, 4-3 
		    (7, 7, 320, 1, 640, 3, 3, 1, 1, 1, 1),
		    (7, 7, 640, 1, 640, 3, 3, 1, 1, 1, 1),
		    (7, 7, 320, 1, 640, 1, 1, 0, 0, 1, 1), # bottlneck
		    # Conv 4-4, 4-5, 4-6, 4-7, 4-8, 4-9
		    (7, 7, 640, 1, 640, 3, 3, 1, 1, 1, 1),
		    (7, 7, 640, 1, 640, 3, 3, 1, 1, 1, 1),
		    (7, 7, 640, 1, 640, 3, 3, 1, 1, 1, 1),
		    (7, 7, 640, 1, 640, 3, 3, 1, 1, 1, 1),
		    (7, 7, 640, 1, 640, 3, 3, 1, 1, 1, 1),
		    (7, 7, 640, 1, 640, 3, 3, 1, 1, 1, 1),

		    # FC # too dense
		    (1, 1, 640, 1, 10, 1, 1, 0, 0, 1, 1)
	    ]
	},
'imagenet':
	{
		'resnet18': [ # stride 2 is simulated by stride 1 since unsure if timeloop can perform stride 2 simulation
		    # Conv1
		    # (224, 224, 3, 1, 64, 7, 7, 3, 3, 2, 2), # the proper dim if we handle stride
		    (112, 112, 3, 1, 64, 7, 7, 3, 3, 1, 1), # stride should be 2
		    
		    # Conv2_x
		    # (112, 112, 64, 1, 64, 3, 3, 1, 1, 2, 2), # the proper dim if we handle stride
		    (56, 56, 64, 1, 64, 3, 3, 1, 1, 1, 1), # stride should be 2
		    (56, 56, 64, 1, 64, 3, 3, 1, 1, 1, 1),
		    (56, 56, 64, 1, 64, 3, 3, 1, 1, 1, 1),
		    (56, 56, 64, 1, 64, 3, 3, 1, 1, 1, 1),
		    
		    # Conv3_x
		    # (56, 56, 64, 1, 128, 3, 3, 1, 1, 2, 2), # the proper dim if we handle stride
		    (28, 28, 64, 1, 128, 3, 3, 1, 1, 1, 1), # stride should be 2
		    (28, 28, 128, 1, 128, 3, 3, 1, 1, 1, 1),
		    (28, 28, 128, 1, 128, 3, 3, 1, 1, 1, 1),
		    (28, 28, 128, 1, 128, 3, 3, 1, 1, 1, 1),
		    
		    # Conv4_x
		    # (28, 28, 128, 1, 256, 3, 3, 1, 1, 2, 2), # the proper dim if we handle stride
		    (14, 14, 128, 1, 256, 3, 3, 1, 1, 1, 1), # stride should be 2
		    (14, 14, 256, 1, 256, 3, 3, 1, 1, 1, 1),
		    (14, 14, 256, 1, 256, 3, 3, 1, 1, 1, 1),
		    (14, 14, 256, 1, 256, 3, 3, 1, 1, 1, 1),
		    
		    # Conv5_x
		    # (14, 14, 256, 1, 512, 3, 3, 1, 1, 2, 2), # the proper dim if we handle stride
		    (7, 7, 256, 1, 512, 3, 3, 1, 1, 1, 1), # stride should be 2
		    (7, 7, 512, 1, 512, 3, 3, 1, 1, 1, 1),
		    (7, 7, 512, 1, 512, 3, 3, 1, 1, 1, 1),
		    (7, 7, 512, 1, 512, 3, 3, 1, 1, 1, 1),
		    
		    # FC
		    (1, 1, 512, 1, 1000, 1, 1, 0, 0, 1, 1)
		],
		'mobilenetv2': [ # stride 2 is simulated by stride 1 since unsure if timeloop can perform stride 2 simulation
		    # Conv1
		    # (224, 224, 3, 1, 64, 7, 7, 3, 3, 2, 2), # the proper dim if we handle stride
		    (112, 112, 3, 1, 32, 3, 3, 1, 1, 1, 1), # stride should be 2
		    
		    # Conv2_x
		    (112, 112, 32, 1, 'D', 3, 3, 1, 1, 1, 1), # Depthwise
		    (112, 112, 32, 1, 16, 1, 1, 0, 0, 1, 1), # Bottleneck

		    # Conv3_x
		    (112, 112, 16, 1, 96, 1, 1, 0, 0, 1, 1), # Bottleneck
		    # (112, 112, 96, 1, D, 3, 3, 1, 1, 2, 2), # the proper dim if we handle stride
		    (56, 56, 96, 1, 'D', 3, 3, 1, 1, 1, 1), # Depthwise
		    (56, 56, 96, 1, 24, 1, 1, 0, 0, 1, 1), # Bottleneck

		    # Conv4_x
		    (56, 56, 24, 1, 144, 1, 1, 0, 0, 1, 1), # Bottleneck
		    (56, 56, 144, 1, 'D', 3, 3, 1, 1, 1, 1), # Depthwise
		    (56, 56, 144, 1, 24, 1, 1, 0, 0, 1, 1), # Bottleneck

		    # Conv5_x
		    (56, 56, 24, 1, 144, 1, 1, 0, 0, 1, 1), # Bottleneck
		    # (112, 112, 96, 1, D, 3, 3, 1, 1, 2, 2), # the proper dim if we handle stride
		    (28, 28, 144, 1, 'D', 3, 3, 1, 1, 1, 1), # Depthwise
		    (28, 28, 144, 1, 32, 1, 1, 0, 0, 1, 1), # Bottleneck

		    # Conv6_x
		    (28, 28, 32, 1, 192, 1, 1, 0, 0, 1, 1), # Bottleneck
		    (28, 28, 192, 1, 'D', 3, 3, 1, 1, 1, 1), # Depthwise
		    (28, 28, 192, 1, 32, 1, 1, 0, 0, 1, 1), # Bottleneck

		    # Conv7_x
		    (28, 28, 32, 1, 192, 1, 1, 0, 0, 1, 1), # Bottleneck
		    (28, 28, 192, 1, 'D', 3, 3, 1, 1, 1, 1), # Depthwise
		    (28, 28, 192, 1, 32, 1, 1, 0, 0, 1, 1), # Bottleneck

		    # Conv8_x
		    (28, 28, 32, 1, 192, 1, 1, 0, 0, 1, 1), # Bottleneck
		    # (112, 112, 96, 1, D, 3, 3, 1, 1, 2, 2), # the proper dim if we handle stride
		    (14, 14, 192, 1, 'D', 3, 3, 1, 1, 1, 1), # Depthwise
		    (14, 14, 192, 1, 64, 1, 1, 0, 0, 1, 1), # Bottleneck

		    # Conv9_x
		    (14, 14, 64, 1, 384, 1, 1, 0, 0, 1, 1), # Bottleneck
		    (14, 14, 384, 1, 'D', 3, 3, 1, 1, 1, 1), # Depthwise
		    (14, 14, 384, 1, 64, 1, 1, 0, 0, 1, 1), # Bottleneck

		    # Conv10_x
		    (14, 14, 64, 1, 384, 1, 1, 0, 0, 1, 1), # Bottleneck
		    (14, 14, 384, 1, 'D', 3, 3, 1, 1, 1, 1), # Depthwise
		    (14, 14, 384, 1, 64, 1, 1, 0, 0, 1, 1), # Bottleneck

		    # Con11_x
		    (14, 14, 64, 1, 384, 1, 1, 0, 0, 1, 1), # Bottleneck
		    (14, 14, 384, 1, 'D', 3, 3, 1, 1, 1, 1), # Depthwise
		    (14, 14, 384, 1, 64, 1, 1, 0, 0, 1, 1), # Bottleneck

		    # Conv12_x
		    (14, 14, 64, 1, 384, 1, 1, 0, 0, 1, 1), # Bottleneck
		    (14, 14, 384, 1, 'D', 3, 3, 1, 1, 1, 1), # Depthwise
		    (14, 14, 384, 1, 96, 1, 1, 0, 0, 1, 1), # Bottleneck

		    # Con13_x
		    (14, 14, 96, 1, 576, 1, 1, 0, 0, 1, 1), # Bottleneck
		    (14, 14, 576, 1, 'D', 3, 3, 1, 1, 1, 1), # Depthwise
		    (14, 14, 576, 1, 96, 1, 1, 0, 0, 1, 1), # Bottleneck

		    # Conv14_x
		    (14, 14, 96, 1, 576, 1, 1, 0, 0, 1, 1), # Bottleneck
		    (14, 14, 576, 1, 'D', 3, 3, 1, 1, 1, 1), # Depthwise
		    (14, 14, 576, 1, 96, 1, 1, 0, 0, 1, 1), # Bottleneck

		    # Conv15_x
		    (14, 14, 96, 1, 576, 1, 1, 0, 0, 1, 1), # Bottleneck
		    # (112, 112, 96, 1, D, 3, 3, 1, 1, 2, 2), # the proper dim if we handle stride
		    (7, 7, 576, 1, 'D', 3, 3, 1, 1, 1, 1), # Depthwise
		    (7, 7, 576, 1, 160, 1, 1, 0, 0, 1, 1), # Bottleneck

		    # Conv16_x
		    (7, 7, 160, 1, 960, 1, 1, 0, 0, 1, 1), # Bottleneck
		    (7, 7, 960, 1, 'D', 3, 3, 1, 1, 1, 1), # Depthwise
		    (7, 7, 960, 1, 160, 1, 1, 0, 0, 1, 1), # Bottleneck

		    # Conv17_x
		    (7, 7, 160, 1, 960, 1, 1, 0, 0, 1, 1), # Bottleneck
		    (7, 7, 960, 1, 'D', 3, 3, 1, 1, 1, 1), # Depthwise
		    (7, 7, 960, 1, 160, 1, 1, 0, 0, 1, 1), # Bottleneck

		    # Conv18_x
		    (7, 7, 160, 1, 960, 1, 1, 0, 0, 1, 1), # Bottleneck
		    (7, 7, 960, 1, 'D', 3, 3, 1, 1, 1, 1), # Depthwise
		    (7, 7, 960, 1, 320, 1, 1, 0, 0, 1, 1), # Bottleneck

		    # Conv19
		    (7, 7, 320, 1, 1280, 1, 1, 0, 0, 1, 1), # Bottleneck
		    
		    # FC
		    (1, 1, 1280, 1, 1000, 1, 1, 0, 0, 1, 1)
		]
	}
}

layer_names = {
'cifar-10': 
	{
		'vgg': [
		    'Conv1_1', 
		    'Conv1_2',
		    'Conv2_1', 
		    'Conv2_2',
		    'Conv3_1', 'Conv3_2', 'Conv3_3',
		    'Conv4_1', 'Conv4_2', 'Conv4_3',
		    'Conv5_1', 
        'Conv5_2', 'Conv5_3',
		    'FC1', 
		    'FC2'
		],
		'dense': [
		    'Conv1', 
		    # Dense Block 1
		    'Block1_1', 'Block1_2', 
		    'Block1_3',
		    'Block1_4', 'Block1_5', 'Block1_6', 'Block1_7',
		    'Block1_8', 'Block1_9', 'Block1_10', 
		    'Block1_11',
		    # Dense Block 2
		    'Block2_1', 'Block2_2', 'Block2_3',
		    'Block2_4', 'Block2_5', 'Block2_6', 'Block2_7',
		    'Block2_8', 'Block2_9', 'Block2_10', 
		    'Block2_11',
		    # Dense Block 3 
		    'Block3_1', 'Block3_2', 'Block3_3',
		    'Block3_4', 'Block3_5', 'Block3_6', 'Block3_7',
		    'Block3_8', 'Block3_9', 'Block3_10',
		    'FC1'
		],
		'wrn': [
		    'Conv1', 
		    'Conv2_1', 'Conv2_2',
		    'Conv2_3', # bottleneck layer
		    'Conv2_4', 'Conv2_5',
		    'Conv2_6', 'Conv2_7',
		    'Conv2_8', 'Conv2_9',
		    # Conv 3
		    'Conv3_1', 'Conv3_2',
		    'Conv3_3', # bottleneck layer
		    'Conv3_4', 
		    'Conv3_5',
		    'Conv3_6', 'Conv3_7',
		    'Conv3_8', 'Conv3_9',
		    # Conv 4
		    'Conv4_1', 'Conv4_2',
		    'Conv4_3', # bottleneck layer
		    'Conv4_4',
		    'Conv4_5',
		    'Conv4_6', 'Conv4_7',
		    'Conv4_8', 'Conv4_9',
		    'FC1'
		]
	},
'imagenet':
	{
		'resnet18': [
		    'Conv1',
		    'Conv2_1',
		    'Conv2_2',
		    'Conv2_3',
		    'Conv2_4',
		    'Conv3_1', 
		    'Conv3_2',
		    'Conv3_3', 
		    'Conv3_4',
		    'Conv4_1', 
		    'Conv4_2',
		    'Conv4_3', 
		    'Conv4_4',
		    'Conv5_1', 
		    'Conv5_2',
		    'Conv5_3', 
		    'Conv5_4',
		    'FC'
		],
		'mobilenetv2': [
		    'Conv1',
		    'Conv2_1', 'Conv2_2',
		    'Conv3_1', 'Conv3_2', 'Conv3_3', 
		    'Conv4_1', 'Conv4_2', 'Conv4_3', 
		    'Conv5_1', 'Conv5_2', 'Conv5_3', 
		    'Conv6_1', 'Conv6_2', 'Conv6_3', 
		    'Conv7_1', 'Conv7_2', 'Conv7_3', 
		    'Conv8_1', 'Conv8_2', 'Conv8_3', 
		    'Conv9_1', 'Conv9_2', 'Conv9_3', 
		    'Conv10_1', 'Conv10_2', 'Conv10_3', 
		    'Conv11_1', 'Conv11_2', 'Conv11_3', 
		    'Conv12_1', 'Conv12_2', 'Conv12_3', 
		    'Conv13_1', 'Conv13_2', 'Conv13_3', 
		    'Conv14_1', 'Conv14_2', 'Conv14_3',
		    'Conv15_1', 'Conv15_2', 'Conv15_3', 
		    'Conv16_1', 'Conv16_2', 'Conv16_3', 
		    'Conv17_1', 'Conv17_2', 'Conv17_3',
		    'Conv18_1', 'Conv18_2', 'Conv18_3',
		    'Conv19',
		    'FC'
		]
	}
}

shapes = {
	'fw': 'cnn-layer',
	'bw': 'backward-pass',
	'wu': 'weight-update'
}

# W, H, C, N, K, S, R, Wpad, Hpad, Wstride, Hstride
layer_sizes = {
'cifar-10': 
	{
		"vgg": [
		    # Conv1
		    (28, 28, 3, 1, 64, 3, 3, 1, 1, 1, 1),
		    (28, 28, 64, 1, 64, 3, 3, 1, 1, 1, 1),
		    # # Conv2
		    (14, 14, 64, 1, 128, 3, 3, 1, 1, 1, 1),
		    (14, 14, 128, 1, 128, 3, 3, 1, 1, 1, 1),
		    # # Conv3
		    (7, 7, 128, 1, 256, 3, 3, 1, 1, 1, 1),
		    (7, 7, 256, 1, 256, 3, 3, 1, 1, 1, 1),
		    (7, 7, 256, 1, 256, 3, 3, 1, 1, 1, 1),
		    # # Conv4
		    (4, 4, 256, 1, 512, 3, 3, 1, 1, 1, 1),
		    (4, 4, 512, 1, 512, 3, 3, 1, 1, 1, 1),
		    (4, 4, 512, 1, 512, 3, 3, 1, 1, 1, 1),
		    # # Conv5
		    (2, 2, 512, 1, 512, 3, 3, 1, 1, 1, 1),
		    (2, 2, 512, 1, 512, 3, 3, 1, 1, 1, 1),
		    (2, 2, 512, 1, 512, 3, 3, 1, 1, 1, 1),
		    # # FC
		    (1, 1, 512, 1, 512, 1, 1, 0, 0, 1, 1),
		    (1, 1, 512, 1, 10, 1, 1, 0, 0, 1, 1)
		],
		'dense': [
		    # Conv1 # too dense
		    (28, 28, 3, 1, 16, 3, 3, 1, 1, 1, 1),
		    
		    # Block 1-1, 1-2, 1-3 
		    (28, 28, 16, 1, 24, 3, 3, 1, 1, 1, 1),
		    (28, 28, 40, 1, 24, 3, 3, 1, 1, 1, 1),
		    (28, 28, 64, 1, 24, 3, 3, 1, 1, 1, 1),
		    # Block 1-4, 1-5, 1-6, 1-7
		    (28, 28, 88, 1, 24, 3, 3, 1, 1, 1, 1),
		    (28, 28, 112, 1, 24, 3, 3, 1, 1, 1, 1),
		    (28, 28, 136, 1, 24, 3, 3, 1, 1, 1, 1),
		    (28, 28, 160, 1, 24, 3, 3, 1, 1, 1, 1),
		    # Block 1-8, 1-9, 1-10, 1-11
		    (28, 28, 184, 1, 24, 3, 3, 1, 1, 1, 1),
		    (28, 28, 208, 1, 24, 3, 3, 1, 1, 1, 1),
		    (28, 28, 232, 1, 24, 3, 3, 1, 1, 1, 1),
		    (28, 28, 256, 1, 256, 1, 1, 0, 0, 1, 1),

		    # Conv 2-1, 2-2, 2-3 
		    (14, 14, 256, 1, 24, 3, 3, 1, 1, 1, 1),
		    (14, 14, 280, 1, 24, 3, 3, 1, 1, 1, 1),
		    (14, 14, 304, 1, 24, 3, 3, 1, 1, 1, 1),
		    # Block 2-4, 2-5, 2-6, 2-7
		    (14, 14, 328, 1, 24, 3, 3, 1, 1, 1, 1),
		    (14, 14, 352, 1, 24, 3, 3, 1, 1, 1, 1),
		    (14, 14, 376, 1, 24, 3, 3, 1, 1, 1, 1),
		    (14, 14, 400, 1, 24, 3, 3, 1, 1, 1, 1),
		    # Block 2-8, 2-9, 2-10, 2-11
		    (14, 14, 424, 1, 24, 3, 3, 1, 1, 1, 1),
		    (14, 14, 448, 1, 24, 3, 3, 1, 1, 1, 1),
		    (14, 14, 472, 1, 24, 3, 3, 1, 1, 1, 1),
		    (14, 14, 496, 1, 496, 1, 1, 0, 0, 1, 1),

		    # Conv 3-1, 3-2, 3-3 
		    (7, 7, 496, 1, 24, 3, 3, 1, 1, 1, 1),
		    (7, 7, 520, 1, 24, 3, 3, 1, 1, 1, 1),
		    (7, 7, 544, 1, 24, 3, 3, 1, 1, 1, 1),
		    # Block 3-4, 3-5, 3-6, 3-7
		    (7, 7, 568, 1, 24, 3, 3, 1, 1, 1, 1),
		    (7, 7, 592, 1, 24, 3, 3, 1, 1, 1, 1),
		    (7, 7, 616, 1, 24, 3, 3, 1, 1, 1, 1),
		    (7, 7, 640, 1, 24, 3, 3, 1, 1, 1, 1),
		    # Block 3-8, 3-9, 3-10
		    (7, 7, 664, 1, 24, 3, 3, 1, 1, 1, 1),
		    (7, 7, 688, 1, 24, 3, 3, 1, 1, 1, 1),
		    (7, 7, 712, 1, 24, 3, 3, 1, 1, 1, 1),
		    # FC1
		    (1, 1, 736, 1, 10, 1, 1, 0, 0, 1, 1)
		],
		'wrn': [
		    # Conv1 # too dense
		    (28, 28, 3, 1, 16, 3, 3, 1, 1, 1, 1),
		    
		    # Conv 2-1, 2-2, 2-3 
		    (28, 28, 16, 1, 160, 3, 3, 1, 1, 1, 1),
		    (28, 28, 160, 1, 160, 3, 3, 1, 1, 1, 1),
		    (28, 28, 16, 1, 160, 1, 1, 0, 0, 1, 1), # bottlneck
		    # Conv 2-4, 2-5, 2-6, 2-7, 2-8, 2-9
		    (28, 28, 160, 1, 160, 3, 3, 1, 1, 1, 1),
		    (28, 28, 160, 1, 160, 3, 3, 1, 1, 1, 1),
		    (28, 28, 160, 1, 160, 3, 3, 1, 1, 1, 1),
		    (28, 28, 160, 1, 160, 3, 3, 1, 1, 1, 1),
		    (28, 28, 160, 1, 160, 3, 3, 1, 1, 1, 1),
		    (28, 28, 160, 1, 160, 3, 3, 1, 1, 1, 1),

		    # Conv 3-1, 3-2, 3-3 
		    (14, 14, 160, 1, 320, 3, 3, 1, 1, 1, 1),
		    (14, 14, 320, 1, 320, 3, 3, 1, 1, 1, 1),
		    (14, 14, 160, 1, 320, 1, 1, 0, 0, 1, 1), # bottlneck
		    # Conv 3-4, 3-5, 3-6, 3-7, 3-8, 3-9
		    (14, 14, 320, 1, 320, 3, 3, 1, 1, 1, 1),
		    (14, 14, 320, 1, 320, 3, 3, 1, 1, 1, 1),
		    (14, 14, 320, 1, 320, 3, 3, 1, 1, 1, 1),
		    (14, 14, 320, 1, 320, 3, 3, 1, 1, 1, 1),
		    (14, 14, 320, 1, 320, 3, 3, 1, 1, 1, 1),
		    (14, 14, 320, 1, 320, 3, 3, 1, 1, 1, 1),

		    # Conv 4-1, 4-2, 4-3 
		    (7, 7, 320, 1, 640, 3, 3, 1, 1, 1, 1),
		    (7, 7, 640, 1, 640, 3, 3, 1, 1, 1, 1),
		    (7, 7, 320, 1, 640, 1, 1, 0, 0, 1, 1), # bottlneck
		    # Conv 4-4, 4-5, 4-6, 4-7, 4-8, 4-9
		    (7, 7, 640, 1, 640, 3, 3, 1, 1, 1, 1),
		    (7, 7, 640, 1, 640, 3, 3, 1, 1, 1, 1),
		    (7, 7, 640, 1, 640, 3, 3, 1, 1, 1, 1),
		    (7, 7, 640, 1, 640, 3, 3, 1, 1, 1, 1),
		    (7, 7, 640, 1, 640, 3, 3, 1, 1, 1, 1),
		    (7, 7, 640, 1, 640, 3, 3, 1, 1, 1, 1),

		    # FC # too dense
		    (1, 1, 640, 1, 10, 1, 1, 0, 0, 1, 1)
	    ]
	},
'imagenet':
	{
		'resnet18': [ # stride 2 is simulated by stride 1 since unsure if timeloop can perform stride 2 simulation
		    # Conv1
		    # (224, 224, 3, 1, 64, 7, 7, 3, 3, 2, 2), # the proper dim if we handle stride
		    (112, 112, 3, 1, 64, 7, 7, 3, 3, 1, 1), # stride should be 2
		    
		    # Conv2_x
		    # (112, 112, 64, 1, 64, 3, 3, 1, 1, 2, 2), # the proper dim if we handle stride
		    (56, 56, 64, 1, 64, 3, 3, 1, 1, 1, 1), # stride should be 2
		    (56, 56, 64, 1, 64, 3, 3, 1, 1, 1, 1),
		    (56, 56, 64, 1, 64, 3, 3, 1, 1, 1, 1),
		    (56, 56, 64, 1, 64, 3, 3, 1, 1, 1, 1),
		    
		    # Conv3_x
		    # (56, 56, 64, 1, 128, 3, 3, 1, 1, 2, 2), # the proper dim if we handle stride
		    (28, 28, 64, 1, 128, 3, 3, 1, 1, 1, 1), # stride should be 2
		    (28, 28, 128, 1, 128, 3, 3, 1, 1, 1, 1),
		    (28, 28, 128, 1, 128, 3, 3, 1, 1, 1, 1),
		    (28, 28, 128, 1, 128, 3, 3, 1, 1, 1, 1),
		    
		    # Conv4_x
		    # (28, 28, 128, 1, 256, 3, 3, 1, 1, 2, 2), # the proper dim if we handle stride
		    (14, 14, 128, 1, 256, 3, 3, 1, 1, 1, 1), # stride should be 2
		    (14, 14, 256, 1, 256, 3, 3, 1, 1, 1, 1),
		    (14, 14, 256, 1, 256, 3, 3, 1, 1, 1, 1),
		    (14, 14, 256, 1, 256, 3, 3, 1, 1, 1, 1),
		    
		    # Conv5_x
		    # (14, 14, 256, 1, 512, 3, 3, 1, 1, 2, 2), # the proper dim if we handle stride
		    (7, 7, 256, 1, 512, 3, 3, 1, 1, 1, 1), # stride should be 2
		    (7, 7, 512, 1, 512, 3, 3, 1, 1, 1, 1),
		    (7, 7, 512, 1, 512, 3, 3, 1, 1, 1, 1),
		    (7, 7, 512, 1, 512, 3, 3, 1, 1, 1, 1),
		    
		    # FC
		    (1, 1, 512, 1, 1000, 1, 1, 0, 0, 1, 1)
		],
		'mobilenetv2': [ # stride 2 is simulated by stride 1 since unsure if timeloop can perform stride 2 simulation
		    # Conv1
		    # (224, 224, 3, 1, 64, 7, 7, 3, 3, 2, 2), # the proper dim if we handle stride
		    (112, 112, 3, 1, 32, 3, 3, 1, 1, 1, 1), # stride should be 2
		    
		    # Conv2_x
		    (112, 112, 32, 1, 'D', 3, 3, 1, 1, 1, 1), # Depthwise
		    (112, 112, 32, 1, 16, 1, 1, 0, 0, 1, 1), # Bottleneck

		    # Conv3_x
		    (112, 112, 16, 1, 96, 1, 1, 0, 0, 1, 1), # Bottleneck
		    # (112, 112, 96, 1, D, 3, 3, 1, 1, 2, 2), # the proper dim if we handle stride
		    (56, 56, 96, 1, 'D', 3, 3, 1, 1, 1, 1), # Depthwise
		    (56, 56, 96, 1, 24, 1, 1, 0, 0, 1, 1), # Bottleneck

		    # Conv4_x
		    (56, 56, 24, 1, 144, 1, 1, 0, 0, 1, 1), # Bottleneck
		    (56, 56, 144, 1, 'D', 3, 3, 1, 1, 1, 1), # Depthwise
		    (56, 56, 144, 1, 24, 1, 1, 0, 0, 1, 1), # Bottleneck

		    # Conv5_x
		    (56, 56, 24, 1, 144, 1, 1, 0, 0, 1, 1), # Bottleneck
		    # (112, 112, 96, 1, D, 3, 3, 1, 1, 2, 2), # the proper dim if we handle stride
		    (28, 28, 144, 1, 'D', 3, 3, 1, 1, 1, 1), # Depthwise
		    (28, 28, 144, 1, 32, 1, 1, 0, 0, 1, 1), # Bottleneck

		    # Conv6_x
		    (28, 28, 32, 1, 192, 1, 1, 0, 0, 1, 1), # Bottleneck
		    (28, 28, 192, 1, 'D', 3, 3, 1, 1, 1, 1), # Depthwise
		    (28, 28, 192, 1, 32, 1, 1, 0, 0, 1, 1), # Bottleneck

		    # Conv7_x
		    (28, 28, 32, 1, 192, 1, 1, 0, 0, 1, 1), # Bottleneck
		    (28, 28, 192, 1, 'D', 3, 3, 1, 1, 1, 1), # Depthwise
		    (28, 28, 192, 1, 32, 1, 1, 0, 0, 1, 1), # Bottleneck

		    # Conv8_x
		    (28, 28, 32, 1, 192, 1, 1, 0, 0, 1, 1), # Bottleneck
		    # (112, 112, 96, 1, D, 3, 3, 1, 1, 2, 2), # the proper dim if we handle stride
		    (14, 14, 192, 1, 'D', 3, 3, 1, 1, 1, 1), # Depthwise
		    (14, 14, 192, 1, 64, 1, 1, 0, 0, 1, 1), # Bottleneck

		    # Conv9_x
		    (14, 14, 64, 1, 384, 1, 1, 0, 0, 1, 1), # Bottleneck
		    (14, 14, 384, 1, 'D', 3, 3, 1, 1, 1, 1), # Depthwise
		    (14, 14, 384, 1, 64, 1, 1, 0, 0, 1, 1), # Bottleneck

		    # Conv10_x
		    (14, 14, 64, 1, 384, 1, 1, 0, 0, 1, 1), # Bottleneck
		    (14, 14, 384, 1, 'D', 3, 3, 1, 1, 1, 1), # Depthwise
		    (14, 14, 384, 1, 64, 1, 1, 0, 0, 1, 1), # Bottleneck

		    # Con11_x
		    (14, 14, 64, 1, 384, 1, 1, 0, 0, 1, 1), # Bottleneck
		    (14, 14, 384, 1, 'D', 3, 3, 1, 1, 1, 1), # Depthwise
		    (14, 14, 384, 1, 64, 1, 1, 0, 0, 1, 1), # Bottleneck

		    # Conv12_x
		    (14, 14, 64, 1, 384, 1, 1, 0, 0, 1, 1), # Bottleneck
		    (14, 14, 384, 1, 'D', 3, 3, 1, 1, 1, 1), # Depthwise
		    (14, 14, 384, 1, 96, 1, 1, 0, 0, 1, 1), # Bottleneck

		    # Con13_x
		    (14, 14, 96, 1, 576, 1, 1, 0, 0, 1, 1), # Bottleneck
		    (14, 14, 576, 1, 'D', 3, 3, 1, 1, 1, 1), # Depthwise
		    (14, 14, 576, 1, 96, 1, 1, 0, 0, 1, 1), # Bottleneck

		    # Conv14_x
		    (14, 14, 96, 1, 576, 1, 1, 0, 0, 1, 1), # Bottleneck
		    (14, 14, 576, 1, 'D', 3, 3, 1, 1, 1, 1), # Depthwise
		    (14, 14, 576, 1, 96, 1, 1, 0, 0, 1, 1), # Bottleneck

		    # Conv15_x
		    (14, 14, 96, 1, 576, 1, 1, 0, 0, 1, 1), # Bottleneck
		    # (112, 112, 96, 1, D, 3, 3, 1, 1, 2, 2), # the proper dim if we handle stride
		    (7, 7, 576, 1, 'D', 3, 3, 1, 1, 1, 1), # Depthwise
		    (7, 7, 576, 1, 160, 1, 1, 0, 0, 1, 1), # Bottleneck

		    # Conv16_x
		    (7, 7, 160, 1, 960, 1, 1, 0, 0, 1, 1), # Bottleneck
		    (7, 7, 960, 1, 'D', 3, 3, 1, 1, 1, 1), # Depthwise
		    (7, 7, 960, 1, 160, 1, 1, 0, 0, 1, 1), # Bottleneck

		    # Conv17_x
		    (7, 7, 160, 1, 960, 1, 1, 0, 0, 1, 1), # Bottleneck
		    (7, 7, 960, 1, 'D', 3, 3, 1, 1, 1, 1), # Depthwise
		    (7, 7, 960, 1, 160, 1, 1, 0, 0, 1, 1), # Bottleneck

		    # Conv18_x
		    (7, 7, 160, 1, 960, 1, 1, 0, 0, 1, 1), # Bottleneck
		    (7, 7, 960, 1, 'D', 3, 3, 1, 1, 1, 1), # Depthwise
		    (7, 7, 960, 1, 320, 1, 1, 0, 0, 1, 1), # Bottleneck

		    # Conv19
		    (7, 7, 320, 1, 1280, 1, 1, 0, 0, 1, 1), # Bottleneck
		    
		    # FC
		    (1, 1, 1280, 1, 1000, 1, 1, 0, 0, 1, 1)
		]
	}
}

layer_names = {
'cifar-10': 
	{
		'vgg': [
		    'Conv1_1', 
		    'Conv1_2',
		    'Conv2_1', 
		    'Conv2_2',
		    'Conv3_1', 'Conv3_2', 'Conv3_3',
		    'Conv4_1', 'Conv4_2', 'Conv4_3',
		    'Conv5_1', 
        'Conv5_2', 'Conv5_3',
		    'FC1', 
		    'FC2'
		],
		'dense': [
		    'Conv1', 
		    # Dense Block 1
		    'Block1_1', 'Block1_2', 
		    'Block1_3',
		    'Block1_4', 'Block1_5', 'Block1_6', 'Block1_7',
		    'Block1_8', 'Block1_9', 'Block1_10', 
		    'Block1_11',
		    # Dense Block 2
		    'Block2_1', 'Block2_2', 'Block2_3',
		    'Block2_4', 'Block2_5', 'Block2_6', 'Block2_7',
		    'Block2_8', 'Block2_9', 'Block2_10', 
		    'Block2_11',
		    # Dense Block 3 
		    'Block3_1', 'Block3_2', 'Block3_3',
		    'Block3_4', 'Block3_5', 'Block3_6', 'Block3_7',
		    'Block3_8', 'Block3_9', 'Block3_10',
		    'FC1'
		],
		'wrn': [
		    'Conv1', 
		    'Conv2_1', 'Conv2_2',
		    'Conv2_3', # bottleneck layer
		    'Conv2_4', 'Conv2_5',
		    'Conv2_6', 'Conv2_7',
		    'Conv2_8', 'Conv2_9',
		    # Conv 3
		    'Conv3_1', 'Conv3_2',
		    'Conv3_3', # bottleneck layer
		    'Conv3_4', 
		    'Conv3_5',
		    'Conv3_6', 'Conv3_7',
		    'Conv3_8', 'Conv3_9',
		    # Conv 4
		    'Conv4_1', 'Conv4_2',
		    'Conv4_3', # bottleneck layer
		    'Conv4_4',
		    'Conv4_5',
		    'Conv4_6', 'Conv4_7',
		    'Conv4_8', 'Conv4_9',
		    'FC1'
		]
	},
'imagenet':
	{
		'resnet18': [
		    'Conv1',
		    'Conv2_1',
		    'Conv2_2',
		    'Conv2_3',
		    'Conv2_4',
		    'Conv3_1', 
		    'Conv3_2',
		    'Conv3_3', 
		    'Conv3_4',
		    'Conv4_1', 
		    'Conv4_2',
		    'Conv4_3', 
		    'Conv4_4',
		    'Conv5_1', 
		    'Conv5_2',
		    'Conv5_3', 
		    'Conv5_4',
		    'FC'
		],
		'mobilenetv2': [
		    'Conv1',
		    'Conv2_1', 'Conv2_2',
		    'Conv3_1', 'Conv3_2', 'Conv3_3', 
		    'Conv4_1', 'Conv4_2', 'Conv4_3', 
		    'Conv5_1', 'Conv5_2', 'Conv5_3', 
		    'Conv6_1', 'Conv6_2', 'Conv6_3', 
		    'Conv7_1', 'Conv7_2', 'Conv7_3', 
		    'Conv8_1', 'Conv8_2', 'Conv8_3', 
		    'Conv9_1', 'Conv9_2', 'Conv9_3', 
		    'Conv10_1', 'Conv10_2', 'Conv10_3', 
		    'Conv11_1', 'Conv11_2', 'Conv11_3', 
		    'Conv12_1', 'Conv12_2', 'Conv12_3', 
		    'Conv13_1', 'Conv13_2', 'Conv13_3', 
		    'Conv14_1', 'Conv14_2', 'Conv14_3',
		    'Conv15_1', 'Conv15_2', 'Conv15_3', 
		    'Conv16_1', 'Conv16_2', 'Conv16_3', 
		    'Conv17_1', 'Conv17_2', 'Conv17_3',
		    'Conv18_1', 'Conv18_2', 'Conv18_3',
		    'Conv19',
		    'FC'
		]
	}
}

shapes = {
	'fw': 'cnn-layer',
	'bw': 'backward-pass',
	'wu': 'weight-update'
}

# W, H, C, N, K, S, R, Wpad, Hpad, Wstride, Hstride # Conv1 # # Conv2 # # Conv3 # # Conv4 # # Conv5 # # FC # Conv1 # too dense # Block 1-1, 1-2, 1-3 # Block 1-4, 1-5, 1-6, 1-7 # Block 1-8, 1-9, 1-10, 1-11 # Conv 2-1, 2-2, 2-3 # Block 2-4, 2-5, 2-6, 2-7 # Block 2-8, 2-9, 2-10, 2-11 # Conv 3-1, 3-2, 3-3 # Block 3-4, 3-5, 3-6, 3-7 # Block 3-8, 3-9, 3-10 # FC1 # Conv1 # too dense # Conv 2-1, 2-2, 2-3 # bottlneck # Conv 2-4, 2-5, 2-6, 2-7, 2-8, 2-9 # Conv 3-1, 3-2, 3-3 # bottlneck # Conv 3-4, 3-5, 3-6, 3-7, 3-8, 3-9 # Conv 4-1, 4-2, 4-3 # bottlneck # Conv 4-4, 4-5, 4-6, 4-7, 4-8, 4-9 # FC # too dense # stride 2 is simulated by stride 1 since unsure if timeloop can perform stride 2 simulation # Conv1 # (224, 224, 3, 1, 64, 7, 7, 3, 3, 2, 2), # the proper dim if we handle stride # stride should be 2 # Conv2_x # (112, 112, 64, 1, 64, 3, 3, 1, 1, 2, 2), # the proper dim if we handle stride # stride should be 2 # Conv3_x # (56, 56, 64, 1, 128, 3, 3, 1, 1, 2, 2), # the proper dim if we handle stride # stride should be 2 # Conv4_x # (28, 28, 128, 1, 256, 3, 3, 1, 1, 2, 2), # the proper dim if we handle stride # stride should be 2 # Conv5_x # (14, 14, 256, 1, 512, 3, 3, 1, 1, 2, 2), # the proper dim if we handle stride # stride should be 2 # FC # stride 2 is simulated by stride 1 since unsure if timeloop can perform stride 2 simulation # Conv1 # (224, 224, 3, 1, 64, 7, 7, 3, 3, 2, 2), # the proper dim if we handle stride # stride should be 2 # Conv2_x # Depthwise # Bottleneck # Conv3_x # Bottleneck # (112, 112, 96, 1, D, 3, 3, 1, 1, 2, 2), # the proper dim if we handle stride # Depthwise # Bottleneck # Conv4_x # Bottleneck # Depthwise # Bottleneck # Conv5_x # Bottleneck # (112, 112, 96, 1, D, 3, 3, 1, 1, 2, 2), # the proper dim if we handle stride # Depthwise # Bottleneck # Conv6_x # Bottleneck # Depthwise # Bottleneck # Conv7_x # Bottleneck # Depthwise # Bottleneck # Conv8_x # Bottleneck # (112, 112, 96, 1, D, 3, 3, 1, 1, 2, 2), # the proper dim if we handle stride # Depthwise # Bottleneck # Conv9_x # Bottleneck # Depthwise # Bottleneck # Conv10_x # Bottleneck # Depthwise # Bottleneck # Con11_x # Bottleneck # Depthwise # Bottleneck # Conv12_x # Bottleneck # Depthwise # Bottleneck # Con13_x # Bottleneck # Depthwise # Bottleneck # Conv14_x # Bottleneck # Depthwise # Bottleneck # Conv15_x # Bottleneck # (112, 112, 96, 1, D, 3, 3, 1, 1, 2, 2), # the proper dim if we handle stride # Depthwise # Bottleneck # Conv16_x # Bottleneck # Depthwise # Bottleneck # Conv17_x # Bottleneck # Depthwise # Bottleneck # Conv18_x # Bottleneck # Depthwise # Bottleneck # Conv19 # Bottleneck # FC # Dense Block 1 # Dense Block 2 # Dense Block 3 # bottleneck layer # Conv 3 # bottleneck layer # Conv 4 # bottleneck layer