wassname/pytorch-ts
1
0
Fork 0
mirror of https://github.com/wassname/pytorch-ts.git synced 2026-08-11 11:24:31 +08:00
Code Issues Packages Projects Releases Wiki Activity

fix RealNVP forward and inverse

Browse Source
This commit is contained in:
Kashif Rasul
2020-01-21 22:16:52 +01:00
parent c32bc0977a
commit f09153cf5a
1 changed files with 14 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files
pts/modules/flows.py
+14 -6
View File
@@ -161,12 +161,16 @@ class LinearMaskedCoupling(nn.Module):
# run through model
s = self.s_net(mx if y is None else torch.cat([y, mx], dim=-1))
t = self.t_net(mx if y is None else torch.cat([y, mx], dim=-1))
t = self.t_net(mx if y is None else torch.cat([y, mx], dim=-1)) * (1 - self.mask)
# cf RealNVP eq 8 where u corresponds to x (here we're modeling u)
u = mx + (1 - self.mask) * (x - t) * torch.exp(-s)
log_s = torch.tanh(s) * (1 - self.mask)
u = (x - t) * torch.exp(-log_s)
# u = mx + (1 - self.mask) * (x - t) * torch.exp(-s)
# log det du/dx; cf RealNVP 8 and 6; note, sum over input_size done at model log_prob
log_abs_det_jacobian = -(1 - self.mask) * s
#log_abs_det_jacobian = -(1 - self.mask) * s
log_abs_det_jacobian = -log_s #.sum(-1, keepdim=True)
return u, log_abs_det_jacobian
@@ -176,10 +180,14 @@ class LinearMaskedCoupling(nn.Module):
# run through model
s = self.s_net(mu if y is None else torch.cat([y, mu], dim=-1))
t = self.t_net(mu if y is None else torch.cat([y, mu], dim=-1))
x = mu + (1 - self.mask) * (u * s.exp() + t) # cf RealNVP eq 7
t = self.t_net(mu if y is None else torch.cat([y, mu], dim=-1)) * (1 - self.mask)
log_abs_det_jacobian = (1 - self.mask) * s # log det dx/du
log_s = torch.tanh(s) * (1 - self.mask)
x = u * torch.exp(log_s) + t
#x = mu + (1 - self.mask) * (u * s.exp() + t) # cf RealNVP eq 7
#log_abs_det_jacobian = (1 - self.mask) * s # log det dx/du
log_abs_det_jacobian = log_s #.sum(-1, keepdim=True)
return x, log_abs_det_jacobian