mxnet.gluon.probability.block.stochastic_block

Stochastic block class.

Classes

StochasticBlock(**kwargs)

StochasticBlock extends HybridBlock to support accumulating loss in the forward phase, which is extremely useful in building Bayesian Neural Network, where the loss function is composed of a classification loss and a KL loss.

StochasticSequential(**kwargs)

Stack StochasticBlock sequentially.

class mxnet.gluon.probability.block.stochastic_block.StochasticBlock(**kwargs)[source]

Bases: HybridBlock

StochasticBlock extends HybridBlock to support accumulating loss in the forward phase, which is extremely useful in building Bayesian Neural Network, where the loss function is composed of a classification loss and a KL loss.

static collectLoss(func)[source]

To accumulate loss during the forward phase, one could first decorate forward with StochasticBlock.collectLoss, and then collect the loss tensor `x by calling self.add_loss(x). For example, in the following forward function, we generate samples from a Gaussian parameterized by loc and scale and accumulate the KL-divergence between it and its prior into the block’s loss storage.: @StochasticBlock.collectLoss def forward(self, loc, scale):

qz = mgp.Normal(loc, scale) # prior pz = mgp.Normal(np.zeros_like(loc), np.ones_like(scale)) self.add_loss(mgp.kl_divergence(qz, pz)) return qz.sample()

class mxnet.gluon.probability.block.stochastic_block.StochasticSequential(**kwargs)[source]

Bases: StochasticBlock

Stack StochasticBlock sequentially.

add(*blocks)[source]

Adds block on top of the stack.

forward(x, *args)[source]

Overrides the forward computation. Arguments must be mxnet.numpy.ndarray.