def test_normal_errors(): model = Normal() assert model._diag_type == 'diag' model._diag_type = 'FAKE' try: model.log_prob(None) except NotImplementedError: pass try: model.sample(4) except NotImplementedError: pass try: model.entropy() except NotImplementedError: pass try: model.scale except NotImplementedError: pass model = Normal([0., 0.], [3., 1.0, 1., 3.]) assert model._diag_type == 'cholesky' try: model.cdf(5.) except NotImplementedError: pass try: model.icdf(5.) except NotImplementedError: pass
def test_normal_affine(): model = Normal(1.0, 4.0) transform = TransformDistribution(Normal(0.0, 1.0), [Affine(1.0, 2.0)]) x = model.sample(4) assert torch.all(transform.log_prob(x)- model.log_prob(x) < 1e-5) x = transform.sample(4) assert torch.all(transform.log_prob(x)- model.log_prob(x) < 1e-5) transform.get_parameters()
def sample(self, X, compute_logprob=False): dist = Normal(F.linear(X, self.W), self.noise * torch.eye(self.K), learnable=False) z = dist.sample(1).squeeze(0) if compute_logprob: return z, dist.log_prob(z) return z
class PMF(Distribution): def __init__(self, N, M, D=5, tau=None): super().__init__() self.N = N self.M = M self.D = D # latent self.U = Parameter(torch.Tensor(D, N).float()) self.V = Parameter(torch.Tensor(D, M).float()) if tau is None: self.prior = None else: self.prior = Normal(0., tau, learnable=False) self.reset_parameters() def reset_parameters(self): init.kaiming_uniform_(self.U, a=math.sqrt(5)) init.kaiming_uniform_(self.V, a=math.sqrt(5)) def prior_penalty(self): if not self.prior: return 0. return self.prior.log_prob( torch.cat([p.view(-1) for p in self.parameters()]).view(-1, 1)).sum() def reconstruct(self): return self.U.t().mm(self.V) def log_prob(self, R): if not isinstance(R, torch.Tensor): R = torch.tensor(R) R = R.view(-1, self.N * self.M).float() mean = self.reconstruct().view(-1) return Normal(mean, torch.ones_like(mean), learnable=False).log_prob(R) + self.prior_penalty() def sample(self, batch_size, noise_std=1.0): return self.reconstruct().expand( (batch_size, self.N, self.M)) + noise_std * torch.randn( (batch_size, self.N, self.M)) def fit(self, R, **kwargs): data = Data(R.view(-1, self.N * self.M)) stats = train(data, self, cross_entropy, **kwargs) return stats def mse(self, R): if not isinstance(R, torch.Tensor): R = torch.tensor(R) return (self.reconstruct() - R.float()).pow(2).mean() def mae(self, R): if not isinstance(R, torch.Tensor): R = torch.tensor(R) return (self.reconstruct() - R.float()).abs().mean()
class ProbabilisticPCA(Distribution): has_latents = True def __init__(self, D, K=2, noise=1., tau=None): super().__init__() self.D = D self.K = K self.W = Parameter(torch.Tensor(D, K).float()) self.noise = torch.tensor(noise) self.latent = Normal(torch.zeros(K), torch.ones(K), learnable=False) self.tau = tau self.prior = None if tau: self.prior = Normal(torch.zeros(K), torch.full((K, ), tau), learnable=False) self.reset_parameters() def reset_parameters(self): init.kaiming_uniform_(self.W, a=math.sqrt(5)) def prior_probability(self, z): if self.prior is None: return 0. return self.prior.log_prob(z) def log_prob(self, X, z): dist = Normal(F.linear(z, self.W), torch.full((z.size(0), self.D), self.noise), learnable=False) return dist.log_prob(X) + self.prior_probability(z) def sample(self, z=None, batch_size=1): if z is None: if self.prior is None: raise ValueError( 'PPCA has no prior distribution to sample latents from, please set tau in init' ) z = self.prior.sample(batch_size) dist = Normal(F.linear(z, self.W), torch.full((z.size(0), self.D), self.noise), learnable=False) return dist.sample(1).squeeze(0) def fit(self, X, variational_dist=None, elbo_kwargs={}, **kwargs): if variational_dist is None: variational_dist = PPCA_Variational_V2(self) data = Data(X) stats = train(data, self, ELBO(variational_dist, **elbo_kwargs), **kwargs) return stats def transform(self, X): return X.mm(self.W)
class VAE(Distribution): has_latents = True def __init__(self, encoder_args={}, decoder_args={}, prior=None, elbo_kwargs={}): super().__init__() preset_encoder_args = { 'input_dim': 1, 'hidden_sizes': [24, 24], 'activation': 'ReLU', 'output_shapes': [1, 1], 'output_activations': [None, 'Softplus'], 'distribution': partial(Normal, learnable=False) } preset_decoder_args = { 'input_dim': 1, 'hidden_sizes': [24, 24], 'activation': 'ReLU', 'output_shapes': [1], 'output_activations': [Sigmoid()], 'distribution': partial(Bernoulli, learnable=False) } preset_encoder_args.update(encoder_args) preset_decoder_args.update(decoder_args) self.encoder = ConditionalModel(**preset_encoder_args) self.decoder = ConditionalModel(**preset_decoder_args) self.criterion = ELBO(self.encoder, **elbo_kwargs) self.prior = prior if prior is None: latent_dim = preset_decoder_args['input_dim'] self.prior = Normal(torch.zeros(latent_dim), torch.ones(latent_dim), learnable=False) def log_prob(self, X, Z=None): # latent given if Z is not None: return self.decoder.log_prob(X, Z) + self.prior.log_prob(Z) Z, encoder_probs = self.encoder.sample(X, compute_logprob=True) prior_probs = self.prior.log_prob(Z) decoder_log_probs = self.decoder.log_prob(X, Z) return decoder_log_probs + prior_probs - encoder_probs def sample(self, batch_size, compute_logprob=False): Z = self.prior.sample(batch_size) return self.decoder.sample(Z, compute_logprob) def fit(self, x, use_elbo=True, **kwargs): data = Data(x) if use_elbo: return train(data, self, self.criterion, **kwargs) return train(data, self, cross_entropy, **kwargs) def parameters(self): for name, param in self.named_parameters(recurse=True): if 'encoder' in name: continue yield param
def log_prob(self, X, z): dist = Normal(F.linear(z, self.W), torch.full((z.size(0), self.D), self.noise), learnable=False) return dist.log_prob(X) + self.prior_probability(z)
def log_prob(self, z, X): dist = Normal(self.ppca.W.t().mm(X), self.ppca.noise * torch.eye(self.ppca.D), learnable=False) return dist.log_prob(z)