======================================================================================================================================= Layer (type:depth-idx) Input Shape Output Shape Param # ======================================================================================================================================= Model -- [1, 7440, 116] -- ├─ViViTCore: 1-1 -- [1, 112, 116, 15, 29] -- │ └─ViViT: 2-1 [1, 5, 140, 36, 64] [1, 112, 116, 15, 29] -- │ │ └─Tokenizer: 3-1 [1, 5, 140, 36, 64] [1, 116, 435, 112] 79,632 │ │ │ └─Sequential: 4-1 [1, 5, 140, 36, 64] [1, 116, 435, 112] -- │ │ │ │ └─Unfold3d: 5-1 [1, 5, 140, 36, 64] [1, 116, 435, 6125] -- │ │ │ │ └─LayerNorm: 5-2 [1, 116, 435, 6125] [1, 116, 435, 6125] 12,250 │ │ │ │ └─Linear: 5-3 [1, 116, 435, 6125] [1, 116, 435, 112] 686,112 │ │ │ │ └─LayerNorm: 5-4 [1, 116, 435, 112] [1, 116, 435, 112] 224 │ │ │ │ └─Dropout: 5-5 [1, 116, 435, 112] [1, 116, 435, 112] -- │ │ └─Transformer: 3-2 [116, 435, 112] [116, 435, 112] -- │ │ │ └─ModuleList: 4-2 -- -- -- │ │ │ │ └─ParallelTransformerBlock: 5-6 [116, 435, 112] [116, 435, 112] -- │ │ │ │ │ └─LayerNorm: 6-1 [116, 435, 112] [116, 435, 112] 224 │ │ │ │ │ └─Linear: 6-2 [116, 435, 112] [116, 435, 1720] 192,640 │ │ │ │ │ └─Linear: 6-3 [116, 435, 528] [116, 435, 112] 59,136 │ │ │ │ │ └─DropPath: 6-4 [116, 435, 112] [116, 435, 112] -- │ │ │ │ │ └─Sequential: 6-5 [116, 435, 136] [116, 435, 112] -- │ │ │ │ │ │ └─GELU: 7-1 [116, 435, 136] [116, 435, 136] -- │ │ │ │ │ │ └─Dropout: 7-2 [116, 435, 136] [116, 435, 136] -- │ │ │ │ │ │ └─Linear: 7-3 [116, 435, 136] [116, 435, 112] 15,232 │ │ │ │ │ └─DropPath: 6-6 [116, 435, 112] [116, 435, 112] -- │ │ │ │ └─ParallelTransformerBlock: 5-7 [116, 435, 112] [116, 435, 112] -- │ │ │ │ │ └─LayerNorm: 6-7 [116, 435, 112] [116, 435, 112] 224 │ │ │ │ │ └─Linear: 6-8 [116, 435, 112] [116, 435, 1720] 192,640 │ │ │ │ │ └─Linear: 6-9 [116, 435, 528] [116, 435, 112] 59,136 │ │ │ │ │ └─DropPath: 6-10 [116, 435, 112] [116, 435, 112] -- │ │ │ │ │ └─Sequential: 6-11 [116, 435, 136] [116, 435, 112] -- │ │ │ │ │ │ └─GELU: 7-4 [116, 435, 136] [116, 435, 136] -- │ │ │ │ │ │ └─Dropout: 7-5 [116, 435, 136] [116, 435, 136] -- │ │ │ │ │ │ └─Linear: 7-6 [116, 435, 136] [116, 435, 112] 15,232 │ │ │ │ │ └─DropPath: 6-12 [116, 435, 112] [116, 435, 112] -- │ │ │ │ └─ParallelTransformerBlock: 5-8 [116, 435, 112] [116, 435, 112] -- │ │ │ │ │ └─LayerNorm: 6-13 [116, 435, 112] [116, 435, 112] 224 │ │ │ │ │ └─Linear: 6-14 [116, 435, 112] [116, 435, 1720] 192,640 │ │ │ │ │ └─Linear: 6-15 [116, 435, 528] [116, 435, 112] 59,136 │ │ │ │ │ └─DropPath: 6-16 [116, 435, 112] [116, 435, 112] -- │ │ │ │ │ └─Sequential: 6-17 [116, 435, 136] [116, 435, 112] -- │ │ │ │ │ │ └─GELU: 7-7 [116, 435, 136] [116, 435, 136] -- │ │ │ │ │ │ └─Dropout: 7-8 [116, 435, 136] [116, 435, 136] -- │ │ │ │ │ │ └─Linear: 7-9 [116, 435, 136] [116, 435, 112] 15,232 │ │ │ │ │ └─DropPath: 6-18 [116, 435, 112] [116, 435, 112] -- │ │ └─Transformer: 3-3 [435, 116, 112] [435, 116, 112] -- │ │ │ └─ModuleList: 4-3 -- -- -- │ │ │ │ └─ParallelTransformerBlock: 5-9 [435, 116, 112] [435, 116, 112] -- │ │ │ │ │ └─LayerNorm: 6-19 [435, 116, 112] [435, 116, 112] 224 │ │ │ │ │ └─Linear: 6-20 [435, 116, 112] [435, 116, 1720] 192,640 │ │ │ │ │ └─Linear: 6-21 [435, 116, 528] [435, 116, 112] 59,136 │ │ │ │ │ └─DropPath: 6-22 [435, 116, 112] [435, 116, 112] -- │ │ │ │ │ └─Sequential: 6-23 [435, 116, 136] [435, 116, 112] -- │ │ │ │ │ │ └─GELU: 7-10 [435, 116, 136] [435, 116, 136] -- │ │ │ │ │ │ └─Dropout: 7-11 [435, 116, 136] [435, 116, 136] -- │ │ │ │ │ │ └─Linear: 7-12 [435, 116, 136] [435, 116, 112] 15,232 │ │ │ │ │ └─DropPath: 6-24 [435, 116, 112] [435, 116, 112] -- │ │ │ │ └─ParallelTransformerBlock: 5-10 [435, 116, 112] [435, 116, 112] -- │ │ │ │ │ └─LayerNorm: 6-25 [435, 116, 112] [435, 116, 112] 224 │ │ │ │ │ └─Linear: 6-26 [435, 116, 112] [435, 116, 1720] 192,640 │ │ │ │ │ └─Linear: 6-27 [435, 116, 528] [435, 116, 112] 59,136 │ │ │ │ │ └─DropPath: 6-28 [435, 116, 112] [435, 116, 112] -- │ │ │ │ │ └─Sequential: 6-29 [435, 116, 136] [435, 116, 112] -- │ │ │ │ │ │ └─GELU: 7-13 [435, 116, 136] [435, 116, 136] -- │ │ │ │ │ │ └─Dropout: 7-14 [435, 116, 136] [435, 116, 136] -- │ │ │ │ │ │ └─Linear: 7-15 [435, 116, 136] [435, 116, 112] 15,232 │ │ │ │ │ └─DropPath: 6-30 [435, 116, 112] [435, 116, 112] -- │ │ │ │ └─ParallelTransformerBlock: 5-11 [435, 116, 112] [435, 116, 112] -- │ │ │ │ │ └─LayerNorm: 6-31 [435, 116, 112] [435, 116, 112] 224 │ │ │ │ │ └─Linear: 6-32 [435, 116, 112] [435, 116, 1720] 192,640 │ │ │ │ │ └─Linear: 6-33 [435, 116, 528] [435, 116, 112] 59,136 │ │ │ │ │ └─DropPath: 6-34 [435, 116, 112] [435, 116, 112] -- │ │ │ │ │ └─Sequential: 6-35 [435, 116, 136] [435, 116, 112] -- │ │ │ │ │ │ └─GELU: 7-16 [435, 116, 136] [435, 116, 136] -- │ │ │ │ │ │ └─Dropout: 7-17 [435, 116, 136] [435, 116, 136] -- │ │ │ │ │ │ └─Linear: 7-18 [435, 116, 136] [435, 116, 112] 15,232 │ │ │ │ │ └─DropPath: 6-36 [435, 116, 112] [435, 116, 112] -- │ │ │ │ └─ParallelTransformerBlock: 5-12 [435, 116, 112] [435, 116, 112] -- │ │ │ │ │ └─LayerNorm: 6-37 [435, 116, 112] [435, 116, 112] 224 │ │ │ │ │ └─Linear: 6-38 [435, 116, 112] [435, 116, 1720] 192,640 │ │ │ │ │ └─Linear: 6-39 [435, 116, 528] [435, 116, 112] 59,136 │ │ │ │ │ └─DropPath: 6-40 [435, 116, 112] [435, 116, 112] -- │ │ │ │ │ └─Sequential: 6-41 [435, 116, 136] [435, 116, 112] -- │ │ │ │ │ │ └─GELU: 7-19 [435, 116, 136] [435, 116, 136] -- │ │ │ │ │ │ └─Dropout: 7-20 [435, 116, 136] [435, 116, 136] -- │ │ │ │ │ │ └─Linear: 7-21 [435, 116, 136] [435, 116, 112] 15,232 │ │ │ │ │ └─DropPath: 6-42 [435, 116, 112] [435, 116, 112] -- │ │ │ │ └─ParallelTransformerBlock: 5-13 [435, 116, 112] [435, 116, 112] -- │ │ │ │ │ └─LayerNorm: 6-43 [435, 116, 112] [435, 116, 112] 224 │ │ │ │ │ └─Linear: 6-44 [435, 116, 112] [435, 116, 1720] 192,640 │ │ │ │ │ └─Linear: 6-45 [435, 116, 528] [435, 116, 112] 59,136 │ │ │ │ │ └─DropPath: 6-46 [435, 116, 112] [435, 116, 112] -- │ │ │ │ │ └─Sequential: 6-47 [435, 116, 136] [435, 116, 112] -- │ │ │ │ │ │ └─GELU: 7-22 [435, 116, 136] [435, 116, 136] -- │ │ │ │ │ │ └─Dropout: 7-23 [435, 116, 136] [435, 116, 136] -- │ │ │ │ │ │ └─Linear: 7-24 [435, 116, 136] [435, 116, 112] 15,232 │ │ │ │ │ └─DropPath: 6-48 [435, 116, 112] [435, 116, 112] -- │ │ └─Rearrange: 3-4 [1, 116, 435, 112] [1, 112, 116, 15, 29] -- ├─MLPShifters: 1-2 -- [1, 2, 116] 513 │ └─MLPShifter: 2-2 [116, 2] [116, 2] -- │ │ └─Sequential: 3-5 [116, 2] [116, 2] -- │ │ │ └─Linear: 4-4 [116, 2] [116, 5] 15 │ │ │ └─Tanh: 4-5 [116, 5] [116, 5] -- │ │ │ └─Linear: 4-6 [116, 5] [116, 5] 30 │ │ │ └─Tanh: 4-7 [116, 5] [116, 5] -- │ │ │ └─Linear: 4-8 [116, 5] [116, 2] 12 │ │ │ └─Tanh: 4-9 [116, 2] [116, 2] -- ├─Readouts: 1-3 [1, 112, 116, 15, 29] [1, 7440, 116] 8,477,649 │ └─Gaussian2DReadout2: 2-3 [1, 112, 116, 15, 29] [1, 7440, 116] 870,480 │ │ └─Sequential: 3-6 [7440, 2] [7440, 2] -- │ │ │ └─Linear: 4-10 [7440, 2] [7440, 30] 90 │ │ │ └─ELU: 4-11 [7440, 30] [7440, 30] -- │ │ │ └─Linear: 4-12 [7440, 30] [7440, 2] 62 │ │ │ └─Tanh: 4-13 [7440, 2] [7440, 2] -- │ │ └─Sequential: 3-7 [7440, 2] [7440, 2] (recursive) │ │ │ └─Linear: 4-14 [7440, 2] [7440, 30] (recursive) │ │ │ └─ELU: 4-15 [7440, 30] [7440, 30] -- │ │ │ └─Linear: 4-16 [7440, 30] [7440, 2] (recursive) │ │ │ └─Tanh: 4-17 [7440, 2] [7440, 2] -- │ │ └─Sequential: 3-8 [7440, 2] [7440, 2] (recursive) │ │ │ └─Linear: 4-18 [7440, 2] [7440, 30] (recursive) │ │ │ └─ELU: 4-19 [7440, 30] [7440, 30] -- │ │ │ └─Linear: 4-20 [7440, 30] [7440, 2] (recursive) │ │ │ └─Tanh: 4-21 [7440, 2] [7440, 2] -- │ │ └─Sequential: 3-9 [116, 4] [116, 112] -- │ │ │ └─Linear: 4-22 [116, 4] [116, 112] 560 │ │ │ └─LayerNorm: 4-23 [116, 112] [116, 112] 224 │ │ │ └─GELU: 4-24 [116, 112] [116, 112] -- │ │ │ └─Dropout: 4-25 [116, 112] [116, 112] -- │ │ │ └─Linear: 4-26 [116, 112] [116, 112] 12,656 │ │ │ └─Tanh: 4-27 [116, 112] [116, 112] -- ├─Exponential: 1-4 [1, 7440, 116] [1, 7440, 116] -- ======================================================================================================================================= Total params: 12,278,365 Trainable params: 12,278,365 Non-trainable params: 0 Total mult-adds (Units.MEGABYTES): 679.88 ======================================================================================================================================= Input size (MB): 1.29 Forward/backward pass size (MB): 9208.73 Params size (MB): 11.40 Estimated Total Size (MB): 9221.42 =======================================================================================================================================