Gluon.probability

[xidulu]

Description

Gluon.probability is part of MXNet 2.0's new features (#16167, #17240)

It is a module under mx.gluon.probability namespace.

Module details

This module is made up of three parts:

• Distribution object

• Bijector/Transformation

• StochasticBlock

Distribution object

Distribution objects are light-weight object providing useful statistical function of a distribution.

Bijector/Transformation

Class representing an invertible or non-invertible transformation.

StochasticBlock

Blocks with loss accumlated during forward.

Demo

A few toy demos can be found in https://github.com/xidulu/mystery_demo, demos still need further refinement, but I believe they can demonstrate some basic functionality of this module.

Transformations with learnable parameters

GAN/RealNVP demonstrates how to build a deep generative model with learnable tranformations.
https://github.com/xidulu/mystery_demo/blob/master/GAN.ipynb

https://github.com/xidulu/mystery_demo/blob/master/real_nvp.ipynb

Semi-supervised VAE

https://github.com/xidulu/mystery_demo/blob/master/SS-VAE.ipynb

Black-box/Auto-differentiation VI for Probabilistic PCA

https://github.com/xidulu/mystery_demo/blob/master/PPCA.ipynb

Bayesian Neural Network with local-parameterization trick

https://github.com/xidulu/mystery_demo/blob/master/bnn_demo.ipynb

Exclusive KL v.s. Inclusive KL

https://nbviewer.jupyter.org/github/xidulu/mystery_demo/blob/master/KL_demo.ipynb

[leezu]

@mxnet-bot run ci [unix-gpu]

[mxnet-bot]

Jenkins CI successfully triggered : [unix-gpu]

[szha]

@mxnet-bot run ci [unix-gpu]

[mxnet-bot]

Jenkins CI successfully triggered : [unix-gpu]

[leezu]

@xidulu any idea why in your PR the test_np_mixedType_unary_funcs[float16-3-rint-None--5.0-5.0] test persistently (8 CI runs in a row) is found to leak ndarray references?

[xidulu]

@leezu
I am looking into this issue.

[xidulu]

@leezu
Issue solved by adding garbage_expected flag: https://github.com/apache/incubator-mxnet/pull/18403/files#diff-a6e158f7742ece308155a95c165505f6R2206

[xidulu]

@szha
Could you help merge this PR?

[feevos]

Hi, thank you VERY much for your contribution!!!!!

I have noticed an issue, and am not sure if this is a bug or not (so as to submit a report). Currently, in the return statement of forward I cannot have a model hybridized and return an object of a distribution. For example, in the VAE code from the mystery_demo, if I want to return the predictions as well as the distribution that I used to get the reconstruction, I get an error when my model is hybridized.

Example

import mxnet as mx
mx.npx.set_np()

from mxnet.gluon import HybridBlock
import mxnet.gluon.probability as mgp

class SimpleNet(HybridBlock):
    def __init__(self, **kwargs):
        super().__init__(**kwargs)        
        
    
    def forward(self,input):
        mu=0.0
        scale=1.
        return input, mgp.Normal(mu,scale)

net = SimpleNet()
net.initialize()
net.hybridize() # <=========== if hybridized code exits with error

xx = mx.np.random.rand(3,5)
out = net(xx)

error:

---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
<ipython-input-6-9e229665d8a2> in <module>
----> 1 out = net(xx)

~/Software/incubator-mxnet/python/mxnet/gluon/block.py in __call__(self, x, *args)
   1645 
   1646             with x.ctx:
-> 1647                 return self._call_cached_op(x, *args)
   1648 
   1649     def forward(self, x, *args):

~/Software/incubator-mxnet/python/mxnet/gluon/block.py in _call_cached_op(self, *args)
   1263     def _call_cached_op(self, *args):
   1264         if self._cached_op is None:
-> 1265             self._build_cache(*args)
   1266 
   1267         if self._first_forward and self._partition_if_dynamic:

~/Software/incubator-mxnet/python/mxnet/gluon/block.py in _build_cache(self, update_graph, *args)
   1136 
   1137     def _build_cache(self, *args, update_graph=True):
-> 1138         data, out = self._get_graph(*args)
   1139         data_names = {data.name: i for i, data in enumerate(data)}
   1140         params = {p.var().name: p for p in self.collect_params().values()}

~/Software/incubator-mxnet/python/mxnet/gluon/block.py in _get_graph(self, *args)
   1132                 return self._get_graph_v1(*args)
   1133             else:  # Gluon 2 based on deferred compute mode
-> 1134                 return self._get_graph_v2(*args)
   1135         return self._cached_graph
   1136 

~/Software/incubator-mxnet/python/mxnet/gluon/block.py in _get_graph_v2(self, *args)
   1122             with autograd.pause(), dc.context():
   1123                 out = super().__call__(*args)
-> 1124             flatten_out, self._out_format = _flatten(out, "output")
   1125             symbol_outputs = dc.get_symbol(flatten_out, sym_cls=type(symbol_inputs[0]))
   1126             self._cached_graph = symbol_inputs, symbol_outputs

~/Software/incubator-mxnet/python/mxnet/gluon/block.py in _flatten(args, inout_str)
    151     fmts = []
    152     for i in args:
--> 153         arg, fmt = _flatten(i, inout_str)
    154         flat.extend(arg)
    155         fmts.append(fmt)

~/Software/incubator-mxnet/python/mxnet/gluon/block.py in _flatten(args, inout_str)
    144 
    145     if not isinstance(args, (list, tuple)):
--> 146         raise ValueError("When hybridized, the input of HybridBlock {}"
    147                          " must be (nested) list of Symbol"
    148                          " or NDArray, "

ValueError: When hybridized, the input of HybridBlock output must be (nested) list of Symbol or NDArray, but got Normal(loc: size (), scale: size (), F: mxnet.ndarray, event_dim: 0) of type <class 'mxnet.gluon.probability.distributions.normal.Normal'>

Regards

[xidulu]

@feevos
Hi, in MXNet, you should avoid return anything other than NDArray from the Block, I would say returning a distribution object could be a bad idea.
Also, according to our design, a distribution object is a light-weight on-the-fly wrapper of statistics functions, that is to say, it is not designed to used as the return value or argument of a function.
My suggestion for your case would be first creating a HybridBlock called "encoder" that returns \mu and \sigma or create a function called "get_latents" which returns \mu and \sigma. After you acquire the parameters, you can create a Normal object with them and call the function you need.

z_mu, z_sigma = Encoder(x)
qz = Normal(z_mu, z_sigma)
print(qz.sample())

[feevos]

Thank you @xidulu , I fixed it I was just wondering if it was a bug or not. Again, MANY thanks for your contribution!!!!