Expand the parameters to n_samples.
(self, logits, means, log_scales, mean_activation='tanh')
| 24 | raise ValueError(f"Unknown activation function: {self.mean_activation}") |
| 25 | |
| 26 | def expand_params(self, logits, means, log_scales, mean_activation='tanh'): |
| 27 | """ |
| 28 | Expand the parameters to n_samples. |
| 29 | """ |
| 30 | B, N, _ = means.shape # [B, N, dim*nr_mix] |
| 31 | dim = int(means.shape[-1] / self.nr_mix) |
| 32 | logits = logits.repeat(1, 1, self.n_sample).reshape(B, -1, 1, self.nr_mix) # [B, N*n_sample, 1, nr_mix] |
| 33 | |
| 34 | means = means.reshape(B, -1, dim, self.nr_mix) # [B, N, dim, nr_mix] |
| 35 | log_scales = log_scales.reshape(B, -1, dim, self.nr_mix) # [B, N, dim, nr_mix] |
| 36 | |
| 37 | means = means.repeat(1, 1, self.n_sample, 1).reshape(means.shape[0], -1, dim, self.nr_mix) # [B, N*n_sample, dim, nr_mix] |
| 38 | log_scales = log_scales.repeat(1, 1, self.n_sample, 1).reshape(means.shape[0], -1, dim, self.nr_mix) # [B, N*n_sample, dim, nr_mix] |
| 39 | |
| 40 | means = self.activate_mean(means.type(torch.float32), mean_activation) |
| 41 | log_scales = log_scales.type(torch.float32) |
| 42 | |
| 43 | logits = F.softmax(logits.type(torch.float32), dim=-1) |
| 44 | |
| 45 | return logits, means, log_scales |
| 46 | |
| 47 | def get_mix_params(self, logits, means, log_scales): |
| 48 | return means.squeeze(-1), log_scales.squeeze(-1), 1.0 |
no test coverage detected