Algorithm 2.
(self, x_T, cond, pre_ori='False')
| 68 | return y_noisy, noise |
| 69 | |
| 70 | def forward(self, x_T, cond, pre_ori='False'): |
| 71 | """ |
| 72 | Algorithm 2. |
| 73 | """ |
| 74 | x_t = x_T |
| 75 | cond_ = cond |
| 76 | for time_step in reversed(range(self.T)): |
| 77 | print("time_step: ", time_step) |
| 78 | t = x_t.new_ones([x_T.shape[0], ], dtype=torch.long) * time_step |
| 79 | if pre_ori == 'False': |
| 80 | mean, var = self.p_mean_variance(x_t=x_t, t=t, cond_=cond_) |
| 81 | if time_step > 0: |
| 82 | noise = torch.randn_like(x_t) |
| 83 | else: |
| 84 | noise = 0 |
| 85 | x_t = mean + torch.sqrt(var) * noise |
| 86 | assert torch.isnan(x_t).int().sum() == 0, "nan in tensor." |
| 87 | else: |
| 88 | if time_step > 0: |
| 89 | ori = self.model(torch.cat((x_t, cond_), dim=1), t) |
| 90 | eps = x_t - extract_(self.sqrt_gammas, t, ori.shape) * ori |
| 91 | eps = eps / extract_(self.sqrt_one_minus_gammas, t, eps.shape) |
| 92 | x_t = extract_(self.sqrt_gammas, t - 1, ori.shape) * ori + extract_(self.sqrt_one_minus_gammas, t - 1, eps.shape) * eps |
| 93 | else: |
| 94 | x_t = self.model(torch.cat((x_t, cond_), dim=1), t) |
| 95 | |
| 96 | x_0 = x_t |
| 97 | return x_0 |
| 98 | |
| 99 | |
| 100 | if __name__ == '__main__': |
nothing calls this directly
no test coverage detected