File size: 2,657 Bytes
2cd560a |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 |
# Copyright (c) OpenMMLab. All rights reserved.
import torch
import torch.nn as nn
from torch import distributions
class RealNVP(nn.Module):
"""RealNVP: a flow-based generative model
`Density estimation using Real NVP
arXiv: <https://arxiv.org/abs/1605.08803>`_.
Code is modified from `the official implementation of RLE
<https://github.com/Jeff-sjtu/res-loglikelihood-regression>`_.
See also `real-nvp-pytorch
<https://github.com/senya-ashukha/real-nvp-pytorch>`_.
"""
@staticmethod
def get_scale_net():
"""Get the scale model in a single invertable mapping."""
return nn.Sequential(
nn.Linear(2, 64), nn.LeakyReLU(), nn.Linear(64, 64),
nn.LeakyReLU(), nn.Linear(64, 2), nn.Tanh())
@staticmethod
def get_trans_net():
"""Get the translation model in a single invertable mapping."""
return nn.Sequential(
nn.Linear(2, 64), nn.LeakyReLU(), nn.Linear(64, 64),
nn.LeakyReLU(), nn.Linear(64, 2))
@property
def prior(self):
"""The prior distribution."""
return distributions.MultivariateNormal(self.loc, self.cov)
def __init__(self):
super(RealNVP, self).__init__()
self.register_buffer('loc', torch.zeros(2))
self.register_buffer('cov', torch.eye(2))
self.register_buffer(
'mask', torch.tensor([[0, 1], [1, 0]] * 3, dtype=torch.float32))
self.s = torch.nn.ModuleList(
[self.get_scale_net() for _ in range(len(self.mask))])
self.t = torch.nn.ModuleList(
[self.get_trans_net() for _ in range(len(self.mask))])
self.init_weights()
def init_weights(self):
"""Initialization model weights."""
for m in self.modules():
if isinstance(m, nn.Linear):
nn.init.xavier_uniform_(m.weight, gain=0.01)
def backward_p(self, x):
"""Apply mapping form the data space to the latent space and calculate
the log determinant of the Jacobian matrix."""
log_det_jacob, z = x.new_zeros(x.shape[0]), x
for i in reversed(range(len(self.t))):
z_ = self.mask[i] * z
s = self.s[i](z_) * (1 - self.mask[i]) # torch.exp(s): betas
t = self.t[i](z_) * (1 - self.mask[i]) # gammas
z = (1 - self.mask[i]) * (z - t) * torch.exp(-s) + z_
log_det_jacob -= s.sum(dim=1)
return z, log_det_jacob
def log_prob(self, x):
"""Calculate the log probability of given sample in data space."""
z, log_det = self.backward_p(x)
return self.prior.log_prob(z) + log_det
|