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 KullbackLeibler
...
>>> # Set distributions (Distribution API)
>>> class Inference(Normal):
...     def __init__(self):
...         super().__init__(var=["z"],cond_var=["x"],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__(var=["x"],cond_var=["z"],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 = -p.log_prob().expectation(q)
>>> 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

save(path)[source]¶

Save the model. The only parameters that are saved are those that are included in the distribution.: Parameters such as device, optimizer, placement of clip_grad, etc. are not saved.

Parameters:	path (str) – Target file path

load(path)[source]¶

Load the model.

Parameters:	path (str) – Target file path

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.

Examples

>>> import torch
>>> from torch import nn, optim
>>> from pixyz.distributions import Deterministic
>>> from pixyz.distributions import Normal
>>> from pixyz.models import GAN
>>> from pixyz.utils import print_latex
>>> x_dim = 128
>>> z_dim = 100
...
>>> # Set distributions (Distribution API)
...
>>> # generator model p(x|z)
>>> class Generator(Deterministic):
...     def __init__(self):
...         super(Generator, self).__init__(var=["x"], cond_var=["z"], name="p")
...         self.model = nn.Sequential(
...             nn.Linear(z_dim, x_dim),
...             nn.Sigmoid()
...         )
...     def forward(self, z):
...         x = self.model(z)
...         return {"x": x}
...
>>> # prior model p(z)
>>> prior = Normal(loc=torch.tensor(0.), scale=torch.tensor(1.),
...                var=["z"], features_shape=[z_dim], name="p_{prior}")
...
>>> # generative model
>>> p_g = Generator()
>>> p = (p_g*prior).marginalize_var("z")
...
>>> # discriminator model p(t|x)
>>> class Discriminator(Deterministic):
...     def __init__(self):
...         super(Discriminator, self).__init__(var=["t"], cond_var=["x"], name="d")
...         self.model = nn.Sequential(
...             nn.Linear(x_dim, 1),
...             nn.Sigmoid()
...         )
...     def forward(self, x):
...         t = self.model(x)
...         return {"t": t}
...
>>> d = Discriminator()
>>> # Set a model (Model API)
>>> model = GAN(p, d, optimizer_params={"lr":0.0002}, d_optimizer_params={"lr":0.0002})
>>> print(model)
Distributions (for training):
  p(x)
Loss function:
  mean(D_{JS}^{Adv} \left[p_{data}(x)||p(x) \right])
Optimizer:
  Adam (
  Parameter Group 0
      amsgrad: False
      betas: (0.9, 0.999)
      eps: 1e-08
      lr: 0.0002
      weight_decay: 0
  )
>>> # Train and test the model
>>> data = torch.randn(1, x_dim)  # Pseudo data
>>> train_loss = model.train({"x": data})
>>> test_loss = model.test({"x": data})

__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).