File size: 2,373 Bytes
b8f4b1f | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 | ============================================================================================================== Layer (type:depth-idx) Output Shape Param # ============================================================================================================== TrainingWrapper -- -- ├─FimOdeon: 1-1 -- -- │ └─TrajectoryEncoder: 2-1 -- 896 │ │ └─TransformerEncoder: 3-1 [1, 1194, 256] 1,579,520 │ └─Sequential: 2-2 -- -- │ │ └─Linear: 3-2 [1, 2400, 256] 1,024 │ │ └─ReLU: 3-3 [1, 2400, 256] -- │ │ └─Linear: 3-4 [1, 2400, 256] 65,792 │ └─AttentionOperator: 2-3 -- -- │ │ └─ModuleList: 3-5 -- 6,318,080 │ │ └─MLP: 3-6 [1, 2400, 3] 132,355 ├─UncertaintyEstimator: 1-2 -- -- │ └─AttentionOperator: 2-4 -- -- │ │ └─ModuleList: 3-7 -- 4,738,560 │ │ └─MLP: 3-8 [1, 2400, 1] 131,841 ============================================================================================================== Total params: 12,968,068 Trainable params: 12,968,068 Non-trainable params: 0 Total mult-adds (Units.MEGABYTES): 12.97 ============================================================================================================== Input size (MB): 0.09 Forward/backward pass size (MB): 771.15 Params size (MB): 51.87 Estimated Total Size (MB): 823.11 ============================================================================================================== |