pixyz.models (Model API)

Model

class pixyz.models.Model(loss, test_loss=None, distributions=[], optimizer=<class 'torch.optim.adam.Adam'>, optimizer_params={}, clip_grad_norm=None, clip_grad_value=None)[source]

Bases: object

This class is for training and testing a loss class. It requires a defined loss class, distributions to train, and optimizer for initialization.

Examples

>>> import torch
>>> from torch import optim
>>> from torch.nn import functional as F
>>> from pixyz.distributions import Bernoulli, Normal
>>> from pixyz.losses import StochasticReconstructionLoss, KullbackLeibler
...
>>> # Set distributions (Distribution API)
>>> class Inference(Normal):
...     def __init__(self):
...         super().__init__(cond_var=["x"], var=["z"], name="q")
...         self.model_loc = torch.nn.Linear(128, 64)
...         self.model_scale = torch.nn.Linear(128, 64)
...     def forward(self, x):
...         return {"loc": self.model_loc(x), "scale": F.softplus(self.model_scale(x))}
...
>>> class Generator(Bernoulli):
...     def __init__(self):
...         super().__init__(cond_var=["z"], var=["x"], name="p")
...         self.model = torch.nn.Linear(64, 128)
...     def forward(self, z):
...         return {"probs": torch.sigmoid(self.model(z))}
...
>>> p = Generator()
>>> q = Inference()
>>> prior = Normal(loc=torch.tensor(0.), scale=torch.tensor(1.),
...                var=["z"], features_shape=[64], name="p_{prior}")
...
>>> # Define a loss function (Loss API)
>>> reconst = StochasticReconstructionLoss(q, p)
>>> kl = KullbackLeibler(q, prior)
>>> loss_cls = (reconst - kl).mean()
>>> print(loss_cls)
mean \left(- D_{KL} \left[q(z|x)||p_{prior}(z) \right] - \mathbb{E}_{q(z|x)} \left[\log p(x|z) \right] \right)
>>>
>>> # Set a model (Model API)
>>> model = Model(loss=loss_cls, distributions=[p, q],
...               optimizer=optim.Adam, optimizer_params={"lr": 1e-3})
>>> # Train and test the model
>>> data = torch.randn(1, 128)  # Pseudo data
>>> train_loss = model.train({"x": data})
>>> test_loss = model.test({"x": data})
__init__(loss, test_loss=None, distributions=[], optimizer=<class 'torch.optim.adam.Adam'>, optimizer_params={}, clip_grad_norm=None, clip_grad_value=None)[source]
Parameters:
  • loss (pixyz.losses.Loss) – Loss class for training.
  • test_loss (pixyz.losses.Loss) – Loss class for testing.
  • distributions (list) – List of pixyz.distributions.Distribution.
  • optimizer (torch.optim) – Optimization algorithm.
  • optimizer_params (dict) – Parameters of optimizer
  • clip_grad_norm (float or int) – Maximum allowed norm of the gradients.
  • clip_grad_value (float or int) – Maximum allowed value of the gradients.
set_loss(loss, test_loss=None)[source]
train(train_x_dict={}, **kwargs)[source]

Train the model.

Parameters:
  • train_x_dict (dict) – Input data.
  • **kwargs
Returns:

loss – Train loss value

Return type:

torch.Tensor

test(test_x_dict={}, **kwargs)[source]

Test the model.

Parameters:
  • test_x_dict (dict) – Input data
  • **kwargs
Returns:

loss – Test loss value

Return type:

torch.Tensor

Pre-implementation models

ML

class pixyz.models.ML(p, other_distributions=[], optimizer=<class 'torch.optim.adam.Adam'>, optimizer_params={}, clip_grad_norm=False, clip_grad_value=False)[source]

Bases: pixyz.models.model.Model

Maximum Likelihood (log-likelihood)

The negative log-likelihood of a given distribution (p) is set as the loss class of this model.

__init__(p, other_distributions=[], optimizer=<class 'torch.optim.adam.Adam'>, optimizer_params={}, clip_grad_norm=False, clip_grad_value=False)[source]
Parameters:
  • p (torch.distributions.Distribution) – Classifier (distribution).
  • optimizer (torch.optim) – Optimization algorithm.
  • optimizer_params (dict) – Parameters of optimizer
  • clip_grad_norm (float or int) – Maximum allowed norm of the gradients.
  • clip_grad_value (float or int) – Maximum allowed value of the gradients.
train(train_x_dict={}, **kwargs)[source]

Train the model.

Parameters:
  • train_x_dict (dict) – Input data.
  • **kwargs
Returns:

loss – Train loss value

Return type:

torch.Tensor

test(test_x_dict={}, **kwargs)[source]

Test the model.

Parameters:
  • test_x_dict (dict) – Input data
  • **kwargs
Returns:

loss – Test loss value

Return type:

torch.Tensor

VAE

class pixyz.models.VAE(encoder, decoder, other_distributions=[], regularizer=None, optimizer=<class 'torch.optim.adam.Adam'>, optimizer_params={}, clip_grad_norm=None, clip_grad_value=None)[source]

Bases: pixyz.models.model.Model

Variational Autoencoder.

In VAE class, reconstruction loss on given distributions (encoder and decoder) is set as the default loss class. However, if you want to add additional terms, e.g., the KL divergence between encoder and prior, you need to set them to the regularizer argument, which defaults to None.

References

[Kingma+ 2013] Auto-Encoding Variational Bayes

__init__(encoder, decoder, other_distributions=[], regularizer=None, optimizer=<class 'torch.optim.adam.Adam'>, optimizer_params={}, clip_grad_norm=None, clip_grad_value=None)[source]
Parameters:
  • encoder (torch.distributions.Distribution) – Encoder distribution.
  • decoder (torch.distributions.Distribution) – Decoder distribution.
  • regularizer (torch.losses.Loss, defaults to None) – If you want to add additional terms to the loss, set them to this argument.
  • optimizer (torch.optim) – Optimization algorithm.
  • optimizer_params (dict) – Parameters of optimizer
  • clip_grad_norm (float or int) – Maximum allowed norm of the gradients.
  • clip_grad_value (float or int) – Maximum allowed value of the gradients.
train(train_x_dict={}, **kwargs)[source]

Train the model.

Parameters:
  • train_x_dict (dict) – Input data.
  • **kwargs
Returns:

loss – Train loss value

Return type:

torch.Tensor

test(test_x_dict={}, **kwargs)[source]

Test the model.

Parameters:
  • test_x_dict (dict) – Input data
  • **kwargs
Returns:

loss – Test loss value

Return type:

torch.Tensor

VI

class pixyz.models.VI(p, approximate_dist, other_distributions=[], optimizer=<class 'torch.optim.adam.Adam'>, optimizer_params={}, clip_grad_norm=None, clip_grad_value=None)[source]

Bases: pixyz.models.model.Model

Variational Inference (Amortized inference)

The ELBO for given distributions (p, approximate_dist) is set as the loss class of this model.

__init__(p, approximate_dist, other_distributions=[], optimizer=<class 'torch.optim.adam.Adam'>, optimizer_params={}, clip_grad_norm=None, clip_grad_value=None)[source]
Parameters:
  • p (torch.distributions.Distribution) – Generative model (distribution).
  • approximate_dist (torch.distributions.Distribution) – Approximate posterior distribution.
  • optimizer (torch.optim) – Optimization algorithm.
  • optimizer_params (dict) – Parameters of optimizer
  • clip_grad_norm (float or int) – Maximum allowed norm of the gradients.
  • clip_grad_value (float or int) – Maximum allowed value of the gradients.
train(train_x_dict={}, **kwargs)[source]

Train the model.

Parameters:
  • train_x_dict (dict) – Input data.
  • **kwargs
Returns:

loss – Train loss value

Return type:

torch.Tensor

test(test_x_dict={}, **kwargs)[source]

Test the model.

Parameters:
  • test_x_dict (dict) – Input data
  • **kwargs
Returns:

loss – Test loss value

Return type:

torch.Tensor

GAN

class pixyz.models.GAN(p, discriminator, optimizer=<class 'torch.optim.adam.Adam'>, optimizer_params={}, d_optimizer=<class 'torch.optim.adam.Adam'>, d_optimizer_params={}, clip_grad_norm=None, clip_grad_value=None)[source]

Bases: pixyz.models.model.Model

Generative Adversarial Network

(Adversarial) Jensen-Shannon divergence between given distributions (p_data, p) is set as the loss class of this model.

__init__(p, discriminator, optimizer=<class 'torch.optim.adam.Adam'>, optimizer_params={}, d_optimizer=<class 'torch.optim.adam.Adam'>, d_optimizer_params={}, clip_grad_norm=None, clip_grad_value=None)[source]
Parameters:
  • p (torch.distributions.Distribution) – Generative model (generator).
  • discriminator (torch.distributions.Distribution) – Critic (discriminator).
  • optimizer (torch.optim) – Optimization algorithm.
  • optimizer_params (dict) – Parameters of optimizer
  • clip_grad_norm (float or int) – Maximum allowed norm of the gradients.
  • clip_grad_value (float or int) – Maximum allowed value of the gradients.
train(train_x_dict={}, adversarial_loss=True, **kwargs)[source]

Train the model.

Parameters:
  • train_x_dict (dict, defaults to {}) – Input data.
  • adversarial_loss (bool, defaults to True) – Whether to train the discriminator.
  • **kwargs
Returns:

  • loss (torch.Tensor) – Train loss value.
  • d_loss (torch.Tensor) – Train loss value of the discriminator (if adversarial_loss is True).

test(test_x_dict={}, adversarial_loss=True, **kwargs)[source]

Train the model.

Parameters:
  • test_x_dict (dict, defaults to {}) – Input data.
  • adversarial_loss (bool, defaults to True) – Whether to return the discriminator loss.
  • **kwargs
Returns:

  • loss (torch.Tensor) – Test loss value.
  • d_loss (torch.Tensor) – Test loss value of the discriminator (if adversarial_loss is True).