| 9 | |
| 10 | class Model(nn.Module): |
| 11 | def __init__(self, args): |
| 12 | super().__init__() |
| 13 | self.graph = Graph('hm36_gt', 'spatial', pad=1) |
| 14 | self.A = nn.Parameter(torch.tensor(self.graph.A, dtype=torch.float32), requires_grad=False) |
| 15 | |
| 16 | self.embedding = nn.Linear(2*args.frames, args.channel) |
| 17 | self.mlp_gcn = Mlp_gcn(args.layers, args.channel, args.d_hid, args.token_dim, self.A, length=args.n_joints, frames=args.frames) |
| 18 | self.head = nn.Linear(args.channel, 3) |
| 19 | |
| 20 | def forward(self, x): |
| 21 | x = rearrange(x, 'b f j c -> b j (c f)').contiguous() # B 17 (2f) |