| 72 | return x |
| 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 |
| 95 | x = rearrange(x, f'b j c -> b c j') |
| 96 | res = x |
| 97 | x = self.norm1(x) |
| 98 | |
| 99 | x_gcn_1 = rearrange(x, 'b c j-> b c 1 j') |
| 100 | x_gcn_1 = self.gcn_1(x_gcn_1) |
| 101 | x_gcn_1 = rearrange(x_gcn_1, 'b c 1 j -> b c j') |
| 102 | |
| 103 | x = res + self.drop_path(self.mlp_1(x) + x_gcn_1) |
| 104 | |
| 105 | ## Channel Graph MLP |
| 106 | x = rearrange(x, f'b c j -> b j c') |
| 107 | res = x |
| 108 | x = self.norm2(x) |
| 109 | |
| 110 | x_gcn_2 = rearrange(x, 'b j c-> b c 1 j') |
| 111 | x_gcn_2 = self.gcn_2(x_gcn_2) |
| 112 | x_gcn_2 = rearrange(x_gcn_2, 'b c 1 j -> b j c') |
| 113 | |
| 114 | x = res + self.drop_path(self.mlp_2(x) + x_gcn_2) |
| 115 | |
| 116 | return x |
| 117 | |
| 118 | |
| 119 | class Mlp_gcn(nn.Module): |