File size: 17,793 Bytes
26f9a55
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
=======================================================================================================================================
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
=======================================================================================================================================