| 73 | |
| 74 | class Block(nn.Module): |
| 75 | def __init__(self, length, frames, dim, tokens_dim, channels_dim, adj, drop=0., |
| 76 | drop_path=0., act_layer=nn.GELU, norm_layer=nn.LayerNorm): |
| 77 | super().__init__() |
| 78 | self.norm1 = norm_layer(length) |
| 79 | |
| 80 | self.gcn_1 = Gcn(dim, dim, adj) |
| 81 | self.gcn_2 = Gcn(dim, dim, adj) |
| 82 | self.adj = adj |
| 83 | |
| 84 | if frames == 1: |
| 85 | self.mlp_1 = Mlp(in_features=length, hidden_features=tokens_dim, act_layer=act_layer, drop=drop) |
| 86 | else: |
| 87 | self.mlp_1 = Mlp_ln(in_features=length, hidden_features=tokens_dim, act_layer=act_layer, drop=drop) |
| 88 | |
| 89 | self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity() |
| 90 | self.norm2 = norm_layer(dim) |
| 91 | self.mlp_2 = Mlp(in_features=dim, hidden_features=channels_dim, act_layer=act_layer, drop=drop) |
| 92 | |
| 93 | def forward(self, x): |
| 94 | ## Spatial Graph MLP |