(mean, logstd)
| 80 | return loss_l1 |
| 81 | |
| 82 | def compute_loss_kl(mean, logstd): |
| 83 | #get input data |
| 84 | kl_loss = 0 |
| 85 | |
| 86 | #kl loss |
| 87 | |
| 88 | kl_shape = kl( mean, logstd) |
| 89 | # free bits from IAF-VAE. so that if the KL drops below a certan value, then we stop reducing the KL |
| 90 | kl_shape = torch.clamp(kl_shape, min=0.25) |
| 91 | |
| 92 | kl_loss = kl_shape.mean() |
| 93 | |
| 94 | return kl_loss |
| 95 | |
| 96 | def kl(mean, logstd): |
| 97 | kl = (-0.5 - logstd + 0.5 * mean ** 2 + 0.5 * torch.exp(2 * logstd)) |