| 64 | |
| 65 | @SUBMODULES.register_module() |
| 66 | class T2MMotionEncoder(nn.Module): |
| 67 | |
| 68 | def __init__(self, |
| 69 | input_size, |
| 70 | movement_hidden_size, |
| 71 | movement_latent_size, |
| 72 | motion_hidden_size, |
| 73 | motion_latent_size): |
| 74 | super().__init__() |
| 75 | self.movement_encoder = MovementConvEncoder( |
| 76 | input_size=input_size-4, |
| 77 | hidden_size=movement_hidden_size, |
| 78 | output_size=movement_latent_size) |
| 79 | self.motion_encoder = MotionEncoderBiGRUCo( |
| 80 | input_size=movement_latent_size, |
| 81 | hidden_size=motion_hidden_size, |
| 82 | output_size=motion_latent_size |
| 83 | ) |
| 84 | |
| 85 | def load_pretrained(self, ckpt_path): |
| 86 | checkpoint = torch.load(ckpt_path, map_location='cpu') |
| 87 | self.movement_encoder.load_state_dict(checkpoint['movement_encoder']) |
| 88 | self.motion_encoder.load_state_dict(checkpoint['motion_encoder']) |
| 89 | |
| 90 | def forward(self, motion, motion_length, motion_mask): |
| 91 | motion = motion.detach().float() |
| 92 | sort_idx = np.argsort(motion_length.data.tolist())[::-1].copy() |
| 93 | rank_idx = np.empty_like(sort_idx) |
| 94 | rank_idx[sort_idx] = np.arange(len(motion_length)) |
| 95 | motion = motion[sort_idx] |
| 96 | motion_length = motion_length[sort_idx] |
| 97 | |
| 98 | movements = self.movement_encoder(motion[..., :-4]).detach() |
| 99 | m_lens = motion_length // 4 |
| 100 | motion_embedding = self.motion_encoder(movements, m_lens) |
| 101 | motion_embedding_ordered = motion_embedding[rank_idx] |
| 102 | return motion_embedding_ordered |
| 103 | |
| 104 | |
| 105 | @SUBMODULES.register_module() |
nothing calls this directly
no outgoing calls
no test coverage detected