def i_sample(self, shape=None, as_dist=False): shape = size_getter(shape) dist = TransformedDistribution( self.noise0.expand(shape), AffineTransform(self.i_mean(), self.i_scale(), event_dim=self._event_dim)) if as_dist: return dist return dist.sample()
def SampleAction(self, mean, std): # mean and ln_var are predicted by the neural network, this function mu = mean sig = std * 0.3 # constraining the standard deviation to at maximum 0.3mu u_range = self.args['U_UB'] - self.args['U_LB'] GPol = norm( mu, sig ) # defining gaussian distribution with mean and std as parameterised scale = AffineTransform(self.args['U_LB'], u_range) GPol = TransformedDistribution(GPol, scale) action = GPol.sample() # drawing randomly from normal distribution assert len(action) == 1 logGP = GPol.log_prob( action) # calculating log probability of action taken return action.cpu(), logGP
def i_sample(self, shape=None, as_dist=False): """ Samples from the initial distribution. :param shape: The number of samples :type shape: int|tuple[int] :param as_dist: Whether to return the new value as a distribution :type as_dist: bool :return: Samples from the initial distribution :rtype: torch.Tensor|float """ shape = ((shape,) if isinstance(shape, int) else shape) or torch.Size([]) loc = concater(self.f0(*parameter_caster(len(shape), *self.theta))) scale = concater(self.g0(*parameter_caster(len(shape), *self.theta))) dist = TransformedDistribution(self.noise0.expand(shape), self._transform(loc, scale)) if as_dist: return dist return dist.sample()
def test_transformed_distribution(base_batch_dim, base_event_dim, transform_dim, num_transforms, sample_shape): shape = torch.Size([2, 3, 4, 5]) base_dist = Normal(0, 1) base_dist = base_dist.expand(shape[4 - base_batch_dim - base_event_dim:]) if base_event_dim: base_dist = Independent(base_dist, base_event_dim) transforms = [ AffineTransform(torch.zeros(shape[4 - transform_dim:]), 1), ReshapeTransform((4, 5), (20, )), ReshapeTransform((3, 20), (6, 10)) ] transforms = transforms[:num_transforms] transform = ComposeTransform(transforms) # Check validation in .__init__(). if base_batch_dim + base_event_dim < transform.domain.event_dim: with pytest.raises(ValueError): TransformedDistribution(base_dist, transforms) return d = TransformedDistribution(base_dist, transforms) # Check sampling is sufficiently expanded. x = d.sample(sample_shape) assert x.shape == sample_shape + d.batch_shape + d.event_shape num_unique = len(set(x.reshape(-1).tolist())) assert num_unique >= 0.9 * x.numel() # Check log_prob shape on full samples. log_prob = d.log_prob(x) assert log_prob.shape == sample_shape + d.batch_shape # Check log_prob shape on partial samples. y = x while y.dim() > len(d.event_shape): y = y[0] log_prob = d.log_prob(y) assert log_prob.shape == d.batch_shape
def propagate(self, x, as_dist=False): """ Propagates the model forward conditional on the previous state and current parameters. :param x: The previous state :type x: torch.Tensor|float :param as_dist: Whether to return the new value as a distribution :type as_dist: bool :return: Samples from the model :rtype: torch.Tensor|float """ loc, scale = self.mean(x), self.scale(x) if isinstance(self, Observable): shape = _get_shape(loc if loc.dim() > scale.dim() else scale, self.ndim) else: shape = _get_shape(x, self.ndim) dist = TransformedDistribution(self.noise.expand(shape), self._transform(loc, scale)) if as_dist: return dist return dist.sample()
class NICE(nn.Module): def __init__(self, prior, coupling, in_out_dim, mid_dim, hidden, bottleneck, compress, device, n_layers): """Initialize a NICE. Args: coupling: number of coupling layers. in_out_dim: input/output dimensions. mid_dim: number of units in a hidden layer. hidden: number of hidden layers. device: run on cpu or gpu """ super(NICE, self).__init__() self.device = device if prior == 'gaussian': self.prior = torch.distributions.Normal( torch.tensor(0.).to(device), torch.tensor(1.).to(device)) elif prior == 'logistic': self.prior = TransformedDistribution( Uniform( torch.tensor(0.).to(device), torch.tensor(1.).to(device)), [SigmoidTransform().inv, AffineTransform(loc=0., scale=1.)]) else: raise ValueError('Prior not implemented.') self.in_out_dim = in_out_dim self.coupling = coupling self.n_layers = n_layers layer = AdditiveCoupling if coupling == 'additive' else AffineCoupling self.coupling_layers = nn.ModuleList([ layer(in_out_dim, mid_dim, hidden, i % 2) for i in range(self.n_layers) ]).to(device) self.scale = Scaling(in_out_dim).to(device) self.bottleneck_factor = compress self.bottleneck_loss = nn.MSELoss() self.bottleneck = bottleneck def f_inverse(self, z): """Transformation g: Z -> X (inverse of f). Args: z: tensor in latent space Z. Returns: transformed tensor in data space X. """ x, det = self.scale(z, reverse=True) for layer in reversed(self.coupling_layers): x, _ = layer(x, 0, reverse=True) return x def f(self, x): """Transformation f: X -> Z (inverse of g). Args: x: tensor in data space X. Returns: transformed tensor in latent space Z and log determinant Jacobian """ log_det_J = 0 for layer in self.coupling_layers: x, log_det_J = layer(x, log_det_J) z, det = self.scale(x) return z, log_det_J + det def loss(self, x): """Computes data log-likelihood. (See Section 3.3 in the NICE paper.) Args: x: input minibatch. Returns: log-likelihood of input. """ z, log_det_J = self.f(x) slices = [ z[:, i::self.bottleneck_factor] for i in range(self.bottleneck_factor) ] s = torch.stack(slices).permute(1, 0, 2) if self.bottleneck == 'redundancy': bottleneck_loss = torch.sum(torch.var(s, dim=1), dim=1) log_ll = 0.0 for slice in slices: log_ll += torch.sum(self.prior.log_prob(slice), dim=1) if self.bottleneck == 'null': winner = slices[-1] loser = torch.ones_like(winner) bottleneck_loss = 0.0 for slice in slices[:-1]: bottleneck_loss += self.bottleneck_loss(slice, loser) log_ll = torch.sum(self.prior.log_prob(winner), dim=1) log_det_J -= np.log( 256 ) * self.in_out_dim #/ self.bottleneck_factor #log det for rescaling from [0.256] (after dequantization) to [0,1] #log_ll = torch.sum(self.prior.log_prob(z), dim=1) return log_ll + log_det_J, bottleneck_loss def sample(self, size): """Generates samples. Args: size: number of samples to generate. Returns: samples from the data space X. """ z = self.prior.sample( (size, self.in_out_dim // self.bottleneck_factor)).to(self.device) z_tag = torch.zeros((size, self.in_out_dim)).to(self.device) if self.bottleneck == 'redundancy': for i in range(self.bottleneck_factor): z_tag[:, i::self.bottleneck_factor] = z if self.bottleneck == 'null': for i in range(self.bottleneck_factor - 1): z_tag[:, i::self.bottleneck_factor] = torch.ones_like(z) z_tag[:, self.bottleneck_factor - 1::self.bottleneck_factor] = z return self.f_inverse(z_tag) def forward(self, x): """Forward pass. Args: x: input minibatch. Returns: log-likelihood of input. """ x = x.to(self.device) return self.loss(x)