def logistic_distribution(loc: Tensor, scale: Tensor): base_distribution = td.Uniform(loc.new_zeros(1), scale.new_zeros(1)) transforms = [ td.SigmoidTransform().inv, td.AffineTransform(loc=loc, scale=scale) ] return td.TransformedDistribution(base_distribution, transforms)
def __init__(self, loc, scale, **kwargs): loc, scale = map(torch.as_tensor, (loc, scale)) base_distribution = ptd.Uniform(torch.zeros_like(loc), torch.ones_like(loc), **kwargs) transforms = [ ptd.SigmoidTransform().inv, ptd.AffineTransform(loc=loc, scale=scale), ] super().__init__(base_distribution, transforms)
def __init__(self, loc: torch.Tensor, scale: torch.Tensor): self.loc, self.scale = dist.utils.broadcast_all(loc, scale) zero, one = torch.Tensor([0, 1]).type_as(loc) base_distribution = dist.Uniform(zero, one).expand(self.loc.shape) transforms = [ dist.SigmoidTransform().inv, dist.AffineTransform(loc=self.loc, scale=self.scale) ] super(Logistic, self).__init__(base_distribution, transforms)
samples = self.prior().rsample((batch_size, )) samples = samples.view((batch_size, ) + self.final_shape) return self.inverse(samples, max_iter=max_iter) def set_num_terms(self, n_terms): for block in self.stack: for layer in block.stack: layer.numSeriesTerms = n_terms if __name__ == "__main__": scale = 1. loc = 0. base_distribution = distributions.Uniform(0., 1.) transforms_1 = [ distributions.SigmoidTransform().inv, distributions.AffineTransform(loc=loc, scale=scale) ] logistic_1 = distributions.TransformedDistribution(base_distribution, transforms_1) transforms_2 = [ LogisticTransform(), distributions.AffineTransform(loc=loc, scale=scale) ] logistic_2 = distributions.TransformedDistribution(base_distribution, transforms_2) x = torch.zeros(2) print(logistic_1.log_prob(x), logistic_2.log_prob(x)) 1 / 0
# only send batch_size to prior, prior has final_shape as attribute samples = self.prior().rsample((batch_size,)) samples = samples.view((batch_size,) + self.final_shape) return self.inverse(samples, max_iter=max_iter) def set_num_terms(self, n_terms): for block in self.stack: for layer in block.stack: layer.numSeriesTerms = n_terms if __name__ == "__main__": scale = 1. loc = 0. base_distribution = distributions.Uniform(0., 1.) transforms_1 = [distributions.SigmoidTransform().inv, distributions.AffineTransform(loc=loc, scale=scale)] logistic_1 = distributions.TransformedDistribution(base_distribution, transforms_1) transforms_2 = [LogisticTransform(), distributions.AffineTransform(loc=loc, scale=scale)] logistic_2 = distributions.TransformedDistribution(base_distribution, transforms_2) x = torch.zeros(2) print(logistic_1.log_prob(x), logistic_2.log_prob(x)) 1/0 diff = lambda x, y: (x - y).abs().sum() batch_size = 13 channels = 3 h, w = 32, 32 in_shape = (batch_size, channels, h, w) x = torch.randn((batch_size, channels, h, w), requires_grad=True)
def sample(args): """ Performs the following: 1. construct model object & load state dict from saved model; 2. make H x W samples from a set of gaussian or logistic prior on the latent space; 3. save to disk as a grid of images. """ # parse settings: if args.dataset == 'mnist': input_dim = 28 * 28 img_height = 28 img_width = 28 img_depth = 1 if args.dataset == 'svhn': input_dim = 32 * 32 * 3 img_height = 32 img_width = 32 img_depth = 3 if args.dataset == 'cifar10': input_dim = 32 * 32 * 3 img_height = 32 img_width = 32 img_depth = 3 if args.dataset == 'tfd': raise NotImplementedError( "[sample] Toronto Faces Dataset unsupported right now. Sorry!") input_dim = None img_height = None img_width = None img_depth = None # shut off gradients for sampling: torch.set_grad_enabled(False) # build model & load state dict: nice = NICEModel(input_dim, args.nhidden, args.nlayers) if args.model_path is not None: nice.load_state_dict(torch.load(args.model_path, map_location='cpu')) print("[sample] Loaded model from file.") nice.eval() # sample a batch: if args.prior == 'logistic': LOGISTIC_LOC = 0.0 LOGISTIC_SCALE = (3. / (np.pi**2)) # (sets variance to 1) logistic = dist.TransformedDistribution(dist.Uniform(0.0, 1.0), [ dist.SigmoidTransform().inv, dist.AffineTransform(loc=LOGISTIC_LOC, scale=LOGISTIC_SCALE) ]) print( "[sample] sampling from logistic prior with loc={0:.4f}, scale={1:.4f}." .format(LOGISTIC_LOC, LOGISTIC_SCALE)) ys = logistic.sample(torch.Size([args.nrows * args.ncols, input_dim])) xs = nice.inverse(ys) if args.prior == 'gaussian': print("[sample] sampling from gaussian prior with loc=0.0, scale=1.0.") ys = torch.randn(args.nrows * args.ncols, input_dim) xs = nice.inverse(ys) # format sample into images of correct shape: image_batch = unflatten_images(xs, img_depth, img_height, img_width) # arrange into a grid and save to file: torchvision.utils.save_image(image_batch, args.save_image_path, nrow=args.nrows) print("[sample] Saved {0}-by-{1} sampled images to {2}.".format( args.nrows, args.ncols, args.save_image_path))