(self, motion, motion_mask=None, condition=None)
| 81 | num_layers=num_layers) |
| 82 | |
| 83 | def forward(self, motion, motion_mask=None, condition=None): |
| 84 | B, T = motion.shape[:2] |
| 85 | motion = motion.view(B, T, -1) |
| 86 | feature = self.skelEmbedding(motion) |
| 87 | if self.use_condition: |
| 88 | if self.output_var: |
| 89 | if self.num_class is None: |
| 90 | sigma_query = self.sigma_layer(condition) |
| 91 | else: |
| 92 | sigma_query = self.sigma_layer[condition.long()] |
| 93 | sigma_query = sigma_query.view(B, 1, -1) |
| 94 | feature = torch.cat((sigma_query, feature), dim=1) |
| 95 | if self.num_class is None: |
| 96 | mu_query = self.mu_layer(condition).view(B, 1, -1) |
| 97 | else: |
| 98 | mu_query = self.mu_layer[condition.long()].view(B, 1, -1) |
| 99 | feature = torch.cat((mu_query, feature), dim=1) |
| 100 | else: |
| 101 | query = self.query.view(1, -1, self.latent_dim).repeat(B, 1, 1) |
| 102 | feature = torch.cat((query, feature), dim=1) |
| 103 | if self.output_var: |
| 104 | motion_mask = torch.cat( |
| 105 | (torch.zeros(B, 2).to(motion.device), 1 - motion_mask), |
| 106 | dim=1).bool() |
| 107 | else: |
| 108 | motion_mask = torch.cat( |
| 109 | (torch.zeros(B, 1).to(motion.device), 1 - motion_mask), |
| 110 | dim=1).bool() |
| 111 | feature = feature.permute(1, 0, 2).contiguous() |
| 112 | feature = self.pos_encoder(feature) |
| 113 | feature = self.seqTransEncoder(feature, |
| 114 | src_key_padding_mask=motion_mask) |
| 115 | if self.use_final_proj: |
| 116 | mu = self.final_mu(feature[0]) |
| 117 | if self.output_var: |
| 118 | sigma = self.final_sigma(feature[1]) |
| 119 | return mu, sigma |
| 120 | return mu |
| 121 | else: |
| 122 | if self.output_var: |
| 123 | return feature[0], feature[1] |
| 124 | else: |
| 125 | return feature[0] |
| 126 | |
| 127 | |
| 128 | @SUBMODULES.register_module() |
nothing calls this directly
no outgoing calls
no test coverage detected