def __init__(self, df, scale, validate_args=None):
self._dim = scale.shape[-1]
assert df > self._dim - 1
self.df = df
self.cholesky_factor = transforms.LowerCholeskyTransform()(scale)
self.chi_sqd_dists = [dist.Chi2(self.df - i) for i in range(self._dim)]
batch_shape, event_shape = scale.shape[:-2], scale.shape[-2:]
super().__init__(batch_shape, event_shape, validate_args)
def __init__(self, z_dim=10, data_dim=64 * 64):
super().__init__()
self._z_dim = z_dim
self._data_dim = data_dim
self.encoder_convs = nn.Sequential(
nn.Conv2d(1, 32, 4, 2, 1),
nn.ReLU(),
nn.Conv2d(32, 32, 4, 2, 1),
nn.ReLU(),
nn.Conv2d(32, 64, 4, 2, 1),
nn.ReLU(),
nn.Conv2d(64, 64, 4, 2, 1),
nn.ReLU(),
nn.Conv2d(64, 128, 1),
nn.ReLU(),
)
self.encoder_linears = nn.Sequential(
nn.Linear(2048, 128),
nn.ReLU(),
)
self.mu_encoder = nn.Linear(128, self._z_dim)
self.scale_encoder = nn.Linear(128, self._z_dim**2)
self.decoder_linears = nn.Sequential(
nn.Linear(self._z_dim, 128),
nn.ReLU(),
nn.Linear(128, 4 * 4 * 64),
nn.ReLU(),
)
self.decoder_convs = nn.Sequential(
nn.ConvTranspose2d(64, 64, 4, 2, 1),
nn.ReLU(),
nn.ConvTranspose2d(64, 32, 4, 2, 1),
nn.ReLU(),
nn.ConvTranspose2d(32, 32, 4, 2, 1),
nn.ReLU(),
nn.ConvTranspose2d(32, 1, 4, 2, 1),
)
self.lower_choleskyize = transforms.LowerCholeskyTransform()
self.register_buffer('cov_loc', torch.eye(self.z_dim))
self.register_buffer('cov_df', torch.tensor([self.z_dim + 1]))
def _transform_to_lower_cholesky(constraint):
return transforms.LowerCholeskyTransform()
@transform_to.register(constraints.less_than)
def _transform_to_less_than(constraint):
loc = constraint.upper_bound
scale = loc.new([-1]).expand_as(loc)
return transforms.ComposeTransform(
[transforms.ExpTransform(),
transforms.AffineTransform(loc, scale)])
biject_to.register(constraints.unit_interval, transforms.SigmoidTransform())
transform_to.register(constraints.unit_interval, transforms.SigmoidTransform())
@biject_to.register(constraints.interval)
@transform_to.register(constraints.interval)
def _transform_to_interval(constraint):
loc = constraint.lower_bound
scale = constraint.upper_bound - constraint.lower_bound
return transforms.ComposeTransform([
transforms.SigmoidTransform(),
transforms.AffineTransform(loc, scale)
])
biject_to.register(constraints.simplex, transforms.StickBreakingTransform())
transform_to.register(constraints.simplex, transforms.BoltzmannTransform())
# TODO define a bijection for LowerCholeskyTransform
transform_to.register(constraints.lower_cholesky,
transforms.LowerCholeskyTransform())
def test_init(self):
x = torch.eye(3) + torch.ones(3, 3)
Param(x)
Param(x, transform=transforms.ExpTransform())
Param(x, transform=transforms.LowerCholeskyTransform())
