From 910adaf579bbf21264c3348e6ec08600164356f2 Mon Sep 17 00:00:00 2001 From: David Nabergoj Date: Sat, 19 Oct 2024 16:41:44 +0200 Subject: [PATCH] Update docs --- docs/source/api/base_distributions.rst | 14 -- docs/source/api/bijections.rst | 24 --- docs/source/api/components.rst | 7 - docs/source/api/multiscale_architectures.rst | 11 -- .../general_modeling.rst} | 54 ++++- docs/source/architectures/image_modeling.rst | 56 ++++++ docs/source/architectures/index.rst | 187 ++++++++++++++++++ .../base_distributions.rst | 61 ++++++ .../bijections/autoregressive/bijections.rst | 103 ++++++++++ .../autoregressive/conditioner_transforms.rst | 21 ++ .../autoregressive/conditioners.rst | 8 + .../autoregressive/transformers.rst | 30 +++ .../bijections/continuous/bijections.rst | 5 + .../developer_reference/bijections/index.rst | 35 ++++ .../bijections/residual/bijections.rst | 2 + .../{api => developer_reference}/flow.rst | 2 +- docs/source/developer_reference/index.rst | 25 +++ docs/source/guides/image_modeling.rst | 26 ++- .../source/guides/mathematical_background.rst | 6 +- .../source/guides/modifying_architectures.rst | 154 +++++++++++++++ docs/source/guides/tutorial.rst | 1 + docs/source/index.rst | 19 +- 22 files changed, 776 insertions(+), 75 deletions(-) delete mode 100644 docs/source/api/base_distributions.rst delete mode 100644 docs/source/api/bijections.rst delete mode 100644 docs/source/api/components.rst delete mode 100644 docs/source/api/multiscale_architectures.rst rename docs/source/{api/architectures.rst => architectures/general_modeling.rst} (63%) create mode 100644 docs/source/architectures/image_modeling.rst create mode 100644 docs/source/architectures/index.rst create mode 100644 docs/source/developer_reference/base_distributions.rst create mode 100644 docs/source/developer_reference/bijections/autoregressive/bijections.rst create mode 100644 docs/source/developer_reference/bijections/autoregressive/conditioner_transforms.rst create mode 100644 docs/source/developer_reference/bijections/autoregressive/conditioners.rst create mode 100644 docs/source/developer_reference/bijections/autoregressive/transformers.rst create mode 100644 docs/source/developer_reference/bijections/continuous/bijections.rst create mode 100644 docs/source/developer_reference/bijections/index.rst create mode 100644 docs/source/developer_reference/bijections/residual/bijections.rst rename docs/source/{api => developer_reference}/flow.rst (96%) create mode 100644 docs/source/developer_reference/index.rst create mode 100644 docs/source/guides/modifying_architectures.rst diff --git a/docs/source/api/base_distributions.rst b/docs/source/api/base_distributions.rst deleted file mode 100644 index 174ba1e..0000000 --- a/docs/source/api/base_distributions.rst +++ /dev/null @@ -1,14 +0,0 @@ -Base distribution objects -========================== - -.. autoclass:: torchflows.base_distributions.gaussian.DiagonalGaussian - :members: __init__ - -.. autoclass:: torchflows.base_distributions.gaussian.DenseGaussian - :members: __init__ - -.. autoclass:: torchflows.base_distributions.mixture.DiagonalGaussianMixture - :members: __init__ - -.. autoclass:: torchflows.base_distributions.mixture.DenseGaussianMixture - :members: __init__ diff --git a/docs/source/api/bijections.rst b/docs/source/api/bijections.rst deleted file mode 100644 index 2be7d28..0000000 --- a/docs/source/api/bijections.rst +++ /dev/null @@ -1,24 +0,0 @@ -Bijection objects -==================== - -All normalizing flow transformations are bijections. -The following classes define forward and inverse pass methods which all flow architectures inherit. - -.. autoclass:: torchflows.bijections.base.Bijection - :members: __init__, forward, inverse - -.. autoclass:: torchflows.bijections.base.BijectiveComposition - :members: __init__ - -.. autoclass:: torchflows.bijections.continuous.base.ContinuousBijection - :members: __init__, forward, inverse - -.. autoclass:: torchflows.bijections.finite.multiscale.base.MultiscaleBijection - :members: __init__ - -Inverting a bijection ---------------------- - -Each bijection can be inverted with the `invert` function. - -.. autofunction:: torchflows.bijections.base.invert \ No newline at end of file diff --git a/docs/source/api/components.rst b/docs/source/api/components.rst deleted file mode 100644 index a84b33e..0000000 --- a/docs/source/api/components.rst +++ /dev/null @@ -1,7 +0,0 @@ -Model components -=================== - -.. toctree:: - base_distributions - bijections - flow diff --git a/docs/source/api/multiscale_architectures.rst b/docs/source/api/multiscale_architectures.rst deleted file mode 100644 index b74d185..0000000 --- a/docs/source/api/multiscale_architectures.rst +++ /dev/null @@ -1,11 +0,0 @@ -Multiscale architectures -======================================================== - -Multiscale architectures are suitable for image modeling. - -.. _multiscale_architectures: - -.. autoclass:: torchflows.architectures.MultiscaleRealNVP -.. autoclass:: torchflows.architectures.MultiscaleRQNSF -.. autoclass:: torchflows.architectures.MultiscaleLRSNSF -.. autoclass:: torchflows.architectures.MultiscaleNICE diff --git a/docs/source/api/architectures.rst b/docs/source/architectures/general_modeling.rst similarity index 63% rename from docs/source/api/architectures.rst rename to docs/source/architectures/general_modeling.rst index 73b9613..b1274f3 100644 --- a/docs/source/api/architectures.rst +++ b/docs/source/architectures/general_modeling.rst @@ -1,38 +1,86 @@ -Standard architectures +API for standard architectures ============================ We lists notable implemented bijection architectures. These all inherit from the Bijection class. -.. _architectures: +.. _autoregressive_architecture_api: Autoregressive architectures -------------------------------- .. autoclass:: torchflows.architectures.RealNVP + :members: __init__ + .. autoclass:: torchflows.architectures.InverseRealNVP + :members: __init__ + .. autoclass:: torchflows.architectures.NICE + :members: __init__ + .. autoclass:: torchflows.architectures.MAF + :members: __init__ + .. autoclass:: torchflows.architectures.IAF + :members: __init__ + .. autoclass:: torchflows.architectures.CouplingRQNSF + :members: __init__ + .. autoclass:: torchflows.architectures.MaskedAutoregressiveRQNSF + :members: __init__ + .. autoclass:: torchflows.architectures.InverseAutoregressiveRQNSF + :members: __init__ + .. autoclass:: torchflows.architectures.CouplingLRS + :members: __init__ + .. autoclass:: torchflows.architectures.MaskedAutoregressiveLRS + :members: __init__ + +.. autoclass:: torchflows.architectures.InverseAutoregressiveLRS + :members: __init__ + .. autoclass:: torchflows.architectures.CouplingDSF + :members: __init__ + .. autoclass:: torchflows.architectures.UMNNMAF + :members: __init__ + +.. _continuous_architecture_api: Continuous architectures ------------------------- .. autoclass:: torchflows.architectures.DeepDiffeomorphicBijection + :members: __init__ + .. autoclass:: torchflows.architectures.RNODE + :members: __init__ + .. autoclass:: torchflows.architectures.FFJORD + :members: __init__ + .. autoclass:: torchflows.architectures.OTFlow + :members: __init__ + +.. _residual_architecture_api: Residual architectures ----------------------- .. autoclass:: torchflows.architectures.ResFlow + :members: __init__ + .. autoclass:: torchflows.architectures.ProximalResFlow + :members: __init__ + .. autoclass:: torchflows.architectures.InvertibleResNet + :members: __init__ + .. autoclass:: torchflows.architectures.PlanarFlow + :members: __init__ + .. autoclass:: torchflows.architectures.RadialFlow -.. autoclass:: torchflows.architectures.SylvesterFlow \ No newline at end of file + :members: __init__ + +.. autoclass:: torchflows.architectures.SylvesterFlow + :members: __init__ diff --git a/docs/source/architectures/image_modeling.rst b/docs/source/architectures/image_modeling.rst new file mode 100644 index 0000000..8609a27 --- /dev/null +++ b/docs/source/architectures/image_modeling.rst @@ -0,0 +1,56 @@ +API for multiscale architectures +======================================================== + +Multiscale architectures are suitable for image modeling. + +.. _multiscale_architecture_api: + + +Classic multiscale architectures +------------------------------ + +.. autoclass:: torchflows.architectures.MultiscaleNICE + :members: __init__ + +.. autoclass:: torchflows.architectures.MultiscaleRealNVP + :members: __init__ + +.. autoclass:: torchflows.architectures.MultiscaleRQNSF + :members: __init__ + +.. autoclass:: torchflows.architectures.MultiscaleLRSNSF + :members: __init__ + +.. autoclass:: torchflows.bijections.finite.multiscale.architectures.MultiscaleDeepSigmoid + :members: __init__ + +.. autoclass:: torchflows.bijections.finite.multiscale.architectures.MultiscaleDenseSigmoid + :members: __init__ + +.. autoclass:: torchflows.bijections.finite.multiscale.architectures.MultiscaleDeepDenseSigmoid + :members: __init__ + + +Glow-style multiscale architectures +------------------------------ + +.. autoclass:: torchflows.architectures.AffineGlow + :members: __init__ + +.. autoclass:: torchflows.architectures.ShiftGlow + :members: __init__ + +.. autoclass:: torchflows.bijections.finite.multiscale.architectures.RQSGlow + :members: __init__ + +.. autoclass:: torchflows.bijections.finite.multiscale.architectures.LRSGlow + :members: __init__ + +.. autoclass:: torchflows.bijections.finite.multiscale.architectures.DeepSigmoidGlow + :members: __init__ + +.. autoclass:: torchflows.bijections.finite.multiscale.architectures.DenseSigmoidGlow + :members: __init__ + +.. autoclass:: torchflows.bijections.finite.multiscale.architectures.DeepDenseSigmoidGlow + :members: __init__ diff --git a/docs/source/architectures/index.rst b/docs/source/architectures/index.rst new file mode 100644 index 0000000..61a09ee --- /dev/null +++ b/docs/source/architectures/index.rst @@ -0,0 +1,187 @@ +Full list of architectures (presets) +===================================================== + +We list all implemented NF architectures and their respective class names below. +Using these presets facilitates experimentation and modeling, however you can also modify each architecture and build new ones. + +.. _autoregressive_architecture_list: + +Autoregressive architectures +----------------------------- + +We provide the list of autoregressive architectures in the table below. +Click the architecture name to see the API and usage examples. +Check the API for all autoregressive architectures :ref:`here `. + +.. list-table:: + :header-rows: 1 + + * - Architecture + - Reference + * - :class:`NICE ` + - Dinh et al. `NICE: Non-linear Independent Components Estimation `_ (2015) + * - :class:`RealNVP ` + - Dinh et al. `Density estimation using Real NVP `_ (2017) + * - :class:`Inverse RealNVP ` + - Dinh et al. `Density estimation using Real NVP `_ (2017) + * - :class:`MAF ` + - Papamakarios et al. `Masked Autoregressive Flow for Density Estimation `_ (2018) + * - :class:`IAF ` + - Kingma et al. `Improving Variational Inference with Inverse Autoregressive Flow `_ (2017) + * - :class:`Coupling RQ-NSF ` + - Durkan et al. `Neural Spline Flows `_ (2019) + * - :class:`Masked autoregressive RQ-NSF ` + - Durkan et al. `Neural Spline Flows `_ (2019) + * - :class:`Inverse autoregressive RQ-NSF ` + - Durkan et al. `Neural Spline Flows `_ (2019) + * - :class:`Coupling LR-NSF ` + - Dolatabadi et al. `Invertible Generative Modeling using Linear Rational Splines `_ (2020) + * - :class:`Masked autoregressive LR-NSF ` + - Dolatabadi et al. `Invertible Generative Modeling using Linear Rational Splines `_ (2020) + * - :class:`Inverse autoregressive LR-NSF ` + - Dolatabadi et al. `Invertible Generative Modeling using Linear Rational Splines `_ (2020) + * - :class:`Coupling deep SF ` + - + * - :class:`Masked autoregressive deep SF ` + - + * - :class:`Inverse autoregressive deep SF ` + - + * - :class:`Coupling dense SF ` + - + * - :class:`Masked autoregressive dense SF ` + - + * - :class:`Inverse autoregressive dense SF ` + - + * - :class:`Coupling deep-dense SF ` + - + * - :class:`Masked autoregressive deep-dense SF ` + - + * - :class:`Inverse autoregressive deep-dense SF ` + - + * - :class:`Unconstrained monotonic neural network ` + - + +.. _multiscale_architecture_list: + +Multiscale architectures +----------------------------------------- +We provide the list of multiscale autoregressive architectures in the table below. +These architectures are specifically made for image modeling, but can also be used for voxels or tensors with more dimensions. +Click the architecture name to see the API and usage examples. +Check the API for all multiscale architectures :ref:`here `. + +.. list-table:: + :header-rows: 1 + + * - Architecture + - Reference + * - :class:`MultiscaleNICE ` + - Dinh et al. `NICE: Non-linear Independent Components Estimation `_ (2015) + * - :class:`Multiscale RealNVP ` + - Dinh et al. `Density estimation using Real NVP `_ (2017) + * - :class:`Multiscale RQ-NSF ` + - Durkan et al. `Neural Spline Flows `_ (2019) + * - :class:`Multiscale LR-NSF ` + - Dolatabadi et al. `Invertible Generative Modeling using Linear Rational Splines `_ (2020) + * - :class:`Multiscale deep SF ` + - + * - :class:`Multiscale dense SF ` + - + * - :class:`Multiscale deep-dense SF ` + - + * - :class:`Shift Glow ` + - + * - :class:`Affine Glow ` + - + * - :class:`RQS Glow ` + - + * - :class:`LRS Glow ` + - + * - :class:`Deep sigmoidal Glow ` + - + * - :class:`Dense sigmoidal Glow ` + - + * - :class:`Deep-dense sigmoidal Glow ` + - + +Residual architectures +---------------------------- +We provide the list of iterative residual architectures in the table below. +Click the architecture name to see the API and usage examples. +Check the API for all residual architectures :ref:`here `. + +.. list-table:: + :header-rows: 1 + + * - Architecture + - Reference + * - :class:`Invertible ResNet ` + - + * - :class:`ResFlow ` + - + * - :class:`ProximalResFlow ` + - + +We also list presets for some convolutional iterative residual architectures in the table below. +These are suitable for image modeling. + +.. list-table:: + :header-rows: 1 + + * - Architecture + - Reference + * - :class:`Convolutional invertible ResNet ` + - + * - :class:`Convolutional ResFlow ` + - + +We finally list presets for residual architectures, based on the matrix determinant lemma. +These support either forward or inverse transformation, but not both. +This means they can be used for either sampling (and variational inference) or density estimation (and maximum likelihood fits), but not both at the same time. + +.. list-table:: + :header-rows: 1 + + * - Architecture + - Reference + * - :class:`Planar flow ` + - + * - :class:`Radial flow ` + - + * - :class:`Sylvester flow ` + - + +Continuous architectures +---------------------------- +We provide the list of continuous architectures in the table below. +Click the architecture name to see the API and usage examples. +Check the API for all continuous architectures :ref:`here `. + +.. list-table:: + :header-rows: 1 + + * - Architecture + - Reference + * - :class:`DDNF ` + - + * - :class:`FFJORD ` + - + * - :class:`RNODE ` + - + * - :class:`OT-Flow ` + - + +We also list presets for convolutional continuous architectures in the table below. +These are suitable for image modeling. + +.. list-table:: + :header-rows: 1 + + * - Architecture + - Reference + * - :class:`Convolutional DDNF ` + - + * - :class:`Convolutional FFJORD ` + - + * - :class:`Convolutional RNODE ` + - \ No newline at end of file diff --git a/docs/source/developer_reference/base_distributions.rst b/docs/source/developer_reference/base_distributions.rst new file mode 100644 index 0000000..614b876 --- /dev/null +++ b/docs/source/developer_reference/base_distributions.rst @@ -0,0 +1,61 @@ +Base distributions +========================== + +Existing base distributions +----------------------------- + +.. autoclass:: torchflows.base_distributions.gaussian.DiagonalGaussian + :members: __init__ + +.. autoclass:: torchflows.base_distributions.gaussian.DenseGaussian + :members: __init__ + +.. autoclass:: torchflows.base_distributions.mixture.DiagonalGaussianMixture + :members: __init__ + +.. autoclass:: torchflows.base_distributions.mixture.DenseGaussianMixture + :members: __init__ + +Creating new base distributions +----------------------------------- + +To create a new base distribution, we must create a subclass of :class:`torch.distributions.Distribution` and :class:`torch.nn.Module`. +This class should support the methods sampling and log probability computation. +We give an example for the diagonal Gaussian base distribution: + +.. code-block:: python + + import torch + import torch.distributions + import torch.nn as nn + import math + + class DiagonalGaussian(torch.distributions.Distribution, nn.Module): + def __init__(self, loc: torch.Tensor, scale: torch.Tensor): + super().__init__(event_shape=loc.shape, validate_args=False) + self.log_2_pi = math.log(2 * math.pi) + self.register_buffer('loc', loc) + self.register_buffer('log_scale', torch.log(scale)) + + @property + def scale(self): + return torch.exp(self.log_scale) + + def sample(self, sample_shape: torch.Size = torch.Size()) -> torch.Tensor: + noise = torch.randn(size=(*sample_shape, *self.event_shape)).to(self.loc) + # Unsqueeze loc and scale to match batch shape + sample_shape_mask = [None for _ in range(len(sample_shape))] + return self.loc[sample_shape_mask] + noise * self.scale[sample_shape_mask] + + def log_prob(self, value: torch.Tensor) -> torch.Tensor: + if len(value.shape) <= len(self.event_shape): + raise ValueError("Incorrect input shape") + # Unsqueeze loc and scale to match batch shape + sample_shape_mask = [None for _ in range(len(value.shape) - len(self.event_shape))] + loc = self.loc[sample_shape_mask] + scale = self.scale[sample_shape_mask] + log_scale = self.log_scale[sample_shape_mask] + + # Compute log probability + elementwise_log_prob = -(0.5 * ((value - loc) / scale) ** 2 + 0.5 * self.log_2_pi + log_scale) + return sum_except_batch(elementwise_log_prob, self.event_shape) \ No newline at end of file diff --git a/docs/source/developer_reference/bijections/autoregressive/bijections.rst b/docs/source/developer_reference/bijections/autoregressive/bijections.rst new file mode 100644 index 0000000..d793793 --- /dev/null +++ b/docs/source/developer_reference/bijections/autoregressive/bijections.rst @@ -0,0 +1,103 @@ +Autoregressive bijections +=========================== + +Autoregressive bijections belong in one of two categories: coupling or masked autoregressive bijections. +Architectures like IAF make use of the inverse masked autoregressive bijection, which simply swaps the `forward` and `inverse` methods of its corresponding masked autoregressive counterpart. +Multiscale architectures are special cases of coupling architectures. +Each autoregressive bijection consists of a transformer (parameterized bijection that transforms a part of the input), a conditioner, and a conditioner transform (model that predicts transformer parameters). +See the :doc:`transformers` and :doc:`conditioner_transforms` sections for more details. +To improve performance, we define subclasses according to the conditioner type. +We list these subclasses in the rest of the document. + +Coupling bijections +-------------------------------------------------- + +Coupling architectures are compositions of coupling bijections, which extend the following base class: + +.. autoclass:: torchflows.bijections.finite.autoregressive.layers_base.CouplingBijection + :members: __init__ + +We give an example on how to create a custom coupling bijection using a transformer, coupling strategy, and conditioner transform: + +.. code-block:: python + + from torchflows.bijections.finite.autoregressive.layers_base import CouplingBijection + from torchflows.bijections.finite.autoregressive.transformers.linear.affine import Affine + from torchflows.bijections.finite.autoregressive.conditioning.coupling_masks import HalfSplit + from torchflows.bijections.finite.autoregressive.conditioning.transforms import ResidualFeedForward + + class AffineCoupling(CouplingBijection): + def __init__(self, event_shape, **kwargs): + coupling = HalfSplit(event_shape) + super().__init__( + event_shape, + transformer_class=Affine, + coupling=HalfSplit(event_shape), + conditioner_transform_class=ResidualFeedForward + ) + + event_shape = (10,) # say we have vectors of size 10 + bijection = AffineCoupling(event_shape) # create the bijection + +Masked autoregressive bijections +---------------------------------------- + +Masked autoregressive and inverse autoregressive architectures are compositions of their respective bijections, extending one of the following classes: + +.. autoclass:: torchflows.bijections.finite.autoregressive.layers_base.MaskedAutoregressiveBijection + :members: __init__ + +.. autoclass:: torchflows.bijections.finite.autoregressive.layers_base.InverseMaskedAutoregressiveBijection + :members: __init__ + +We give an example on how to create a custom coupling bijection using a transformer, coupling strategy, and conditioner transform: + +.. code-block:: python + + from torchflows.bijections.finite.autoregressive.layers_base import CouplingBijection + from torchflows.bijections.finite.autoregressive.transformers.linear.affine import Affine + from torchflows.bijections.finite.autoregressive.conditioning.transforms import ResidualFeedForward + + class AffineForwardMaskedAutoregressive(MaskedAutoregressiveBijection): + def __init__(self, event_shape, **kwargs): + super().__init__( + event_shape, + transformer_class=Affine, + conditioner_transform_class=ResidualFeedForward + ) + + class AffineInverseMaskedAutoregressive(InverseMaskedAutoregressiveBijection): + def __init__(self, event_shape, **kwargs): + super().__init__( + event_shape, + transformer_class=Affine, + conditioner_transform_class=ResidualFeedForward + ) + + + # say we have 100 vectors of size 10 + event_shape = (10,) + x = torch.randn(size=(100, *event_shape)) + + bijection = AffineCoupling(event_shape) # create the bijection + z, log_det_forward = bijection.forward(x) + y, log_det_inverse = bijection.inverse(z) + + +Multiscale autoregressive bijections +-------------------------------------------------- + +Multiscale architectures are coupling architectures which are specialized for image modeling, extending the class below: + +.. autoclass:: torchflows.bijections.finite.multiscale.base.MultiscaleBijection + :members: __init__ + +See also +------------ + +.. toctree:: + :maxdepth: 1 + + transformers + conditioners + conditioner_transforms \ No newline at end of file diff --git a/docs/source/developer_reference/bijections/autoregressive/conditioner_transforms.rst b/docs/source/developer_reference/bijections/autoregressive/conditioner_transforms.rst new file mode 100644 index 0000000..98873fc --- /dev/null +++ b/docs/source/developer_reference/bijections/autoregressive/conditioner_transforms.rst @@ -0,0 +1,21 @@ +List of conditioner transforms +============================== + +.. autoclass:: torchflows.bijections.finite.autoregressive.conditioning.transforms.FeedForward +.. autoclass:: torchflows.bijections.finite.autoregressive.conditioning.transforms.Linear +.. autoclass:: torchflows.bijections.finite.autoregressive.conditioning.transforms.ResidualFeedForward +.. autoclass:: torchflows.bijections.finite.autoregressive.conditioning.transforms.Constant +.. autoclass:: torchflows.bijections.finite.autoregressive.conditioning.transforms.MADE +.. autoclass:: torchflows.bijections.finite.autoregressive.conditioning.transforms.LinearMADE + +Conditioner combinations +-------------------------- +.. autoclass:: torchflows.bijections.finite.autoregressive.conditioning.transforms.CombinedConditioner +.. autoclass:: torchflows.bijections.finite.autoregressive.conditioning.transforms.RegularizedCombinedConditioner +.. autoclass:: torchflows.bijections.finite.autoregressive.conditioning.transforms.RegularizedGraphicalConditioner + +Base classes +--------------- +.. autoclass:: torchflows.bijections.finite.autoregressive.conditioning.transforms.ConditionerTransform +.. autoclass:: torchflows.bijections.finite.autoregressive.conditioning.transforms.ElementwiseConditionerTransform +.. autoclass:: torchflows.bijections.finite.autoregressive.conditioning.transforms.TensorConditionerTransform diff --git a/docs/source/developer_reference/bijections/autoregressive/conditioners.rst b/docs/source/developer_reference/bijections/autoregressive/conditioners.rst new file mode 100644 index 0000000..24f5e73 --- /dev/null +++ b/docs/source/developer_reference/bijections/autoregressive/conditioners.rst @@ -0,0 +1,8 @@ +List of conditioners +======================================== + +.. autoclass:: torchflows.bijections.finite.autoregressive.conditioning.coupling_masks.PartialCoupling +.. autoclass:: torchflows.bijections.finite.autoregressive.conditioning.coupling_masks.Coupling +.. autoclass:: torchflows.bijections.finite.autoregressive.conditioning.coupling_masks.GraphicalCoupling +.. autoclass:: torchflows.bijections.finite.autoregressive.conditioning.coupling_masks.HalfSplit +.. autofunction:: torchflows.bijections.finite.autoregressive.conditioning.coupling_masks.make_coupling diff --git a/docs/source/developer_reference/bijections/autoregressive/transformers.rst b/docs/source/developer_reference/bijections/autoregressive/transformers.rst new file mode 100644 index 0000000..a95a5f7 --- /dev/null +++ b/docs/source/developer_reference/bijections/autoregressive/transformers.rst @@ -0,0 +1,30 @@ +List of transformers +================================ + +Torchflows supports several transformers to be used in autoregressive and multiscale normalizing flows. + +Linear transformers +-------------------- +.. autoclass:: torchflows.bijections.finite.autoregressive.transformers.linear.affine.Affine +.. autoclass:: torchflows.bijections.finite.autoregressive.transformers.linear.affine.InverseAffine +.. autoclass:: torchflows.bijections.finite.autoregressive.transformers.linear.affine.Shift +.. autoclass:: torchflows.bijections.finite.autoregressive.transformers.linear.convolution.Invertible1x1ConvolutionTransformer +.. autoclass:: torchflows.bijections.finite.autoregressive.transformers.linear.matrix.LUTransformer + +Spline transformers +-------------------------------- +.. autoclass:: torchflows.bijections.finite.autoregressive.transformers.spline.linear.Linear +.. autoclass:: torchflows.bijections.finite.autoregressive.transformers.spline.linear_rational.LinearRational +.. autoclass:: torchflows.bijections.finite.autoregressive.transformers.spline.rational_quadratic.RationalQuadratic + +Combination transformers +--------------------------------------- + +.. autoclass:: torchflows.bijections.finite.autoregressive.transformers.combination.sigmoid.Sigmoid +.. autoclass:: torchflows.bijections.finite.autoregressive.transformers.combination.sigmoid.DenseSigmoid +.. autoclass:: torchflows.bijections.finite.autoregressive.transformers.combination.sigmoid.DeepSigmoid +.. autoclass:: torchflows.bijections.finite.autoregressive.transformers.combination.sigmoid.DeepDenseSigmoid + +Integration transformers +--------------------------------- +.. autoclass:: torchflows.bijections.finite.autoregressive.transformers.integration.unconstrained_monotonic_neural_network.UnconstrainedMonotonicNeuralNetwork diff --git a/docs/source/developer_reference/bijections/continuous/bijections.rst b/docs/source/developer_reference/bijections/continuous/bijections.rst new file mode 100644 index 0000000..3323007 --- /dev/null +++ b/docs/source/developer_reference/bijections/continuous/bijections.rst @@ -0,0 +1,5 @@ +Continuous bijections +=========================== + +.. autoclass:: torchflows.bijections.continuous.base.ContinuousBijection + :members: __init__, forward, inverse \ No newline at end of file diff --git a/docs/source/developer_reference/bijections/index.rst b/docs/source/developer_reference/bijections/index.rst new file mode 100644 index 0000000..f5cffd5 --- /dev/null +++ b/docs/source/developer_reference/bijections/index.rst @@ -0,0 +1,35 @@ +Bijections +==================== + +All normalizing flow transformations are bijections and compositions thereof. + +Base bijections +------------------ + +The following classes define forward and inverse pass methods which all bijections inherit. + +.. autoclass:: torchflows.bijections.base.Bijection + :members: __init__, forward, inverse + +.. autoclass:: torchflows.bijections.base.BijectiveComposition + :members: __init__ + + +Bijection subclasses for different NF families +------------------------------------------------------------------ +To improve efficiency of forward and inverse passes in NF layers, we subclass the base bijections with respect to different families of NF architectures. +On the pages below, we list base classes for each family, and provide a list of already implemented classes. + +.. toctree:: + :maxdepth: 1 + + autoregressive/bijections + residual/bijections + continuous/bijections + +Inverting a bijection +------------------------------ + +Each bijection can be inverted with the `invert` function. + +.. autofunction:: torchflows.bijections.base.invert \ No newline at end of file diff --git a/docs/source/developer_reference/bijections/residual/bijections.rst b/docs/source/developer_reference/bijections/residual/bijections.rst new file mode 100644 index 0000000..22e7aa1 --- /dev/null +++ b/docs/source/developer_reference/bijections/residual/bijections.rst @@ -0,0 +1,2 @@ +Residual bijections +=========================== diff --git a/docs/source/api/flow.rst b/docs/source/developer_reference/flow.rst similarity index 96% rename from docs/source/api/flow.rst rename to docs/source/developer_reference/flow.rst index 0fe1d1d..f4cdfff 100644 --- a/docs/source/api/flow.rst +++ b/docs/source/developer_reference/flow.rst @@ -1,4 +1,4 @@ -Flow objects +Flow wrappers =============================== The `Flow` object contains a base distribution and a bijection. diff --git a/docs/source/developer_reference/index.rst b/docs/source/developer_reference/index.rst new file mode 100644 index 0000000..729d00d --- /dev/null +++ b/docs/source/developer_reference/index.rst @@ -0,0 +1,25 @@ +Developer reference +=========================== + +This section describes how to create NF architectures and NF components in Torchflows. +NFs consist of two main components: + +* a base distribution, +* a bijection. + +In Torchflows, we further wrap these two with the :class:`torchflows.flows.Flow` object or one of its subclasses to enable e.g., fitting NFs, computing the log probability density, and sampling. + +At its core, each of these components is a PyTorch module which extends existing base classes: + +* :class:`torch.distributions.Distribution` and :class:`torch.nn.Module` for base distributions, +* :class:`torchflows.bijections.base.Bijection` for bijections, +* :class:`torchflows.flows.BaseFlow` for flow wrappers. + +Check the following pages for existing subclasses and to learn to create new subclasses for your modeling and research needs: + +.. toctree:: + :maxdepth: 1 + + base_distributions + bijections/index + flow diff --git a/docs/source/guides/image_modeling.rst b/docs/source/guides/image_modeling.rst index 757365e..134d8d5 100644 --- a/docs/source/guides/image_modeling.rst +++ b/docs/source/guides/image_modeling.rst @@ -3,9 +3,11 @@ Image modeling When modeling images, we can use specialized multiscale architectures which use convolutional neural network conditioners and specialized coupling schemes. These architectures expect event shapes to be *(channels, height, width)*. +See the :ref:`list of multiscale architecture presets here `. -.. note:: - Multiscale architectures are currently undergoing improvements. +Basic multiscale architectures +--------------------------------------- +We provide some basic multiscale presets and give an example for the RealNVP variant below: .. code-block:: python @@ -19,4 +21,24 @@ These architectures expect event shapes to be *(channels, height, width)*. torch.manual_seed(0) training_images = torch.randn(size=(n_images, *image_shape)) # synthetic data flow = Flow(MultiscaleRealNVP(image_shape)) + flow.fit(training_images, show_progress=True) + +Glow-style multiscale architectures +------------------------------------------- + +Glow-style architectures are extensions of basic multiscale architectures which use an additional invertible 1x1 convolution in each layer. +We give an example for Glow with affine transformers below: + +.. code-block:: python + + import torch + from torchflows.flows import Flow + from torchflows.architectures import AffineGlow + + image_shape = (3, 28, 28) + n_images = 100 + + torch.manual_seed(0) + training_images = torch.randn(size=(n_images, *image_shape)) # synthetic data + flow = Flow(AffineGlow(image_shape)) flow.fit(training_images, show_progress=True) \ No newline at end of file diff --git a/docs/source/guides/mathematical_background.rst b/docs/source/guides/mathematical_background.rst index cb4dd94..d1807c7 100644 --- a/docs/source/guides/mathematical_background.rst +++ b/docs/source/guides/mathematical_background.rst @@ -1,5 +1,9 @@ +Mathematical background +============================= + + What is a normalizing flow -========================== +-------------------------------------------- A normalizing flow (NF) is a flexible trainable distribution. It is defined as a bijective transformation of a simple distribution, such as a standard Gaussian. diff --git a/docs/source/guides/modifying_architectures.rst b/docs/source/guides/modifying_architectures.rst new file mode 100644 index 0000000..f1f2249 --- /dev/null +++ b/docs/source/guides/modifying_architectures.rst @@ -0,0 +1,154 @@ +Modifying normalizing flow architectures +============================================ + +We sometimes wish to experiment with bijection parameters to improve NF performance on a given dataset. +We give a few examples on how to achieve this with Torchflows. + +Passing hyperparameters to existing architecture constructors +------------------------------------------------------------------- +We can make basic modifications to an existing NF architecture by passing it certain keyword arguments. +The permitted keyword arguments depend on the architecture. +Suppose we are working with RealNVP, which is a composition of several affine coupling layers. +We wish our RealNVP instance to have 5 affine coupling layers. +Each affine coupling layer should use a feed-forward neural network conditioner with 5 layers, as well as 10 hidden neurons and the ReLU activation in each layer. + +.. code-block:: python + + import torch.nn as nn + from torchflows.flows import Flow + from torchflows.architectures import RealNVP + from torchflows.bijections.finite.autoregressive.conditioning.transforms import FeedForward + + event_shape = (10,) + custom_hyperparameters = { + 'n_layers': 5, + 'conditioner_transform_class': FeedForward, + 'conditioner_kwargs': { + 'n_layers': 5, + 'n_hidden': 10, + 'nonlinearity': nn.ReLU + } + } + bijection = RealNVP(event_shape, **custom_hyperparameters) + flow = Flow(bijection) + +`Autoregressive architectures `_ can receive hyperparameters through the following keyword arguments: + +* ``n_layers``: the number of affine coupling layer; +* ``conditioner_transform_class``: the conditioner type to use in each layer; +* ``conditioner_kwargs``: conditioner keyword arguments for each layer; +* ``transformer_kwargs``: transformer keyword arguments for each layer. + +The specific keyword arguments depend on which conditioner and transformer we are using. +Check the list of implemented conditioner transforms and their constructors :doc:`here <../developer_reference/bijections/autoregressive/conditioner_transforms>`. +See which transformers are used in each architecture :ref:`here `. + +Coupling architectures can also receive: + +* ``edge_list``: an edge list of conditional dimension interactions; +* ``coupling_kwargs``: keyword arguments for :func:`make_coupling` in each layer. + +To see how other architectures use keyword arguments, consider checking the :doc:`list of architectures <../architectures/index>` + +Composing existing bijections with custom hyperparameters +------------------------------------------------------------- +In the previous section, we learned how to modify a preset architecture by passing some hyperparameters. +In residual and autoregressive NFs, this approach will use the same hyperparameters for each layer of the NF. +For more customization, we can create individual layers and compose them into a custom architecture. +Suppose we wish to create a NF with five layers: + +* two affine coupling layers, +* a rational quadratic spline coupling layer, +* an invertible residual network layer, +* an elementwise shift layer. + +The above model can be coded as follows: + +.. code-block:: python + + import torch + from torchflows.flows import Flow + from torchflows.bijections.base import BijectiveComposition + from torchflows.bijections.finite.autoregressive.layers import AffineCoupling, RQSCoupling, ElementwiseShift + from torchflows.bijections.finite.residual.iterative import InvertibleResNetBlock + + torch.manual_seed(0) + event_shape = (10,) + bijection = BijectiveComposition( + event_shape, + [ + AffineCoupling(event_shape), + AffineCoupling(event_shape), + RQSCoupling(event_shape), + InvertibleResNetBlock(event_shape), + ElementwiseShift(event_shape), + ] + ) + flow = Flow(bijection) + + x_new = flow.sample((10,)) + log_prob = flow.log_prob(x_new) + +We can also customize each layer with custom hyperparameters, for example: + +.. code-block:: python + + import torch + from torchflows.flows import Flow + from torchflows.bijections.base import BijectiveComposition + from torchflows.bijections.finite.autoregressive.layers import AffineCoupling, RQSCoupling, ElementwiseShift + from torchflows.bijections.finite.residual.iterative import InvertibleResNetBlock + + torch.manual_seed(0) + event_shape = (10,) + bijection = BijectiveComposition( + event_shape, + [ + AffineCoupling(event_shape, conditioner_kwargs={'n_hidden': 5, 'n_layers': 10}), + AffineCoupling(event_shape), + RQSCoupling(event_shape, conditioner_kwargs={'n_layers': 1}), + InvertibleResNetBlock(event_shape, hidden_size=4, n_hidden_layers=3), + ElementwiseShift(event_shape), + ] + ) + flow = Flow(bijection) + + x_new = flow.sample((10,)) + log_prob = flow.log_prob(x_new) + +.. note:: + + Due to the large number of bijections in the library, argument names are not always consistent across bijections. + Check bijection constructors to make sure you are using correct argument names. + We are working to improve this in a future release. + +Composing NF architectures +---------------------------------------- + +Since each NF transformation is a bijection, we can compose them as any other. +We give an example below, where we compose RealNVP, coupling RQ-NSF, FFJORD, and ResFlow: + +.. code-block:: python + + import torch + from torchflows.flows import Flow + from torchflows.bijections.base import BijectiveComposition + from torchflows.bijections.finite.autoregressive.architectures import RealNVP, CouplingRQNSF + from torchflows.bijections.finite.residual.architectures import ResFlow + from torchflows.bijections.continuous.ffjord import FFJORD + + torch.manual_seed(0) + event_shape = (10,) + bijection = BijectiveComposition( + event_shape, + [ + RealNVP(event_shape), + CouplingRQNSF(event_shape), + FFJORD(event_shape), + ResFlow(event_shape) + ] + ) + flow = Flow(bijection) + + x_new = flow.sample((10,)) + log_prob = flow.log_prob(x_new) diff --git a/docs/source/guides/tutorial.rst b/docs/source/guides/tutorial.rst index f9f63b1..a3fd141 100644 --- a/docs/source/guides/tutorial.rst +++ b/docs/source/guides/tutorial.rst @@ -11,3 +11,4 @@ We provide tutorials and notebooks for typical Torchflows use cases. image_modeling choosing_base_distributions cuda + modifying_architectures diff --git a/docs/source/index.rst b/docs/source/index.rst index d2f9b46..eb2210f 100644 --- a/docs/source/index.rst +++ b/docs/source/index.rst @@ -25,21 +25,16 @@ Torchflows can be installed easily using pip: For other install options, see the :ref:`install ` section. -Guides -========= +Table of contents +---------------------------- .. toctree:: + :maxdepth: 2 guides/installing guides/tutorial - -API -==== - -.. toctree:: - :maxdepth: 3 - - api/components - api/architectures - api/multiscale_architectures + architectures/index + architectures/general_modeling + architectures/image_modeling + developer_reference/index