from abc import ABC, abstractmethod

import re

import numpy as np

from ..wrappers import Dropout

from ..utils import calc_pad_dims_2D

from ..activations import Tanh, Sigmoid, ReLU, LeakyReLU, Affine

from ..layers import (

DotProductAttention,

FullyConnected,

Multiply,

LSTMCell,

Add,

class ModuleBase(ABC):

def __init__(self):

self.X = None

self.trainable = True

@abstractmethod

def _init_params(self, **kwargs):

raise NotImplementedError

@abstractmethod

def forward(self, z, **kwargs):

raise NotImplementedError

@abstractmethod

def backward(self, out, **kwargs):

raise NotImplementedError

@property

def components(self):

comps = []

for c in self.hyperparameters["component_ids"]:

if hasattr(self, c):

comps.append(getattr(self, c))

return comps

def freeze(self):

self.trainable = False

for c in self.components:

c.freeze()

def unfreeze(self):

self.trainable = True

for c in self.components:

c.unfreeze()

def update(self, cur_loss=None):

assert self.trainable, "Layer is frozen"

for c in self.components:

def flush_gradients(self):

self.X = []

self._dv = {}

for c in self.components:

for k, v in c.derived_variables.items():

c.derived_variables[k] = None

for k, v in c.gradients.items():

c.gradients[k] = np.zeros_like(v)

def set_params(self, summary_dict):

cids = self.hyperparameters["component_ids"]

for k, v in summary_dict["parameters"].items():

if k == "components":

for c, cd in summary_dict["parameters"][k].items():

if c in cids:

getattr(self, c).set_params(cd)

elif k in self.parameters:

self.parameters[k] = v

for k, v in summary_dict["hyperparameters"].items():

if k == "components":

for c, cd in summary_dict["hyperparameters"][k].items():

if c in cids:

getattr(self, c).set_params(cd)

if k in self.hyperparameters:

if k == "act_fn" and v == "ReLU":

self.hyperparameters[k] = ReLU()

elif v == "act_fn" and v == "Sigmoid":

self.hyperparameters[k] = Sigmoid()

elif v == "act_fn" and v == "Tanh":

self.hyperparameters[k] = Tanh()

elif v == "act_fn" and "Affine" in v:

r = r"Affine\(slope=(.*), intercept=(.*)\)"

slope, intercept = re.match(r, v).groups()

self.hyperparameters[k] = Affine(float(slope), float(intercept))

elif v == "act_fn" and "Leaky ReLU" in v:

r = r"Leaky ReLU\(alpha=(.*)\)"

alpha = re.match(r, v).groups()[0]

self.hyperparameters[k] = LeakyReLU(float(alpha))

else:

self.hyperparameters[k] = v

def summary(self):

return {

"parameters": self.parameters,

"layer": self.hyperparameters["layer"],

"hyperparameters": self.hyperparameters,

class WavenetResidualModule(ModuleBase):

def __init__(

self,

dilation,

optimizer=None,

A WaveNet-like residual block with causal dilated convolutions.

.. code-block:: text

*Skip path in* >-------------------------------------------> + ---> *Skip path out*

Causal      |--> Tanh --|                  |

*Main    |--> Dilated Conv1D -|           * --> 1x1 Conv1D --|

path >--|                    |--> Sigm --|                  |

in*     |-------------------------------------------------> + ---> *Main path out*

*Residual path*

On the final block, the output of the skip path is further processed to

produce the network predictions.

References

.. [1] van den Oord et al. (2016). "Wavenet: a generative model for raw

audio". https://arxiv.org/pdf/1609.03499.pdf

Parameters

ch_residual : int

The number of output channels for the 1x1

:class:`~numpy_ml.neural_nets.layers.Conv1D` layer in the main path.

ch_dilation : int

The number of output channels for the causal dilated

dilation : int

The dilation rate for the causal dilated

kernel_width : int

The width of the causal dilated

:class:`~numpy_ml.neural_nets.layers.Conv1D` kernel in the main

path.

init : {'glorot_normal', 'glorot_uniform', 'he_normal', 'he_uniform'}

The weight initialization strategy. Default is 'glorot_uniform'.

optimizer : str or :doc:`Optimizer <numpy_ml.neural_nets.optimizers>` object or None

The optimization strategy to use when performing gradient updates

within the :meth:`update` method.  If None, use the

:class:`~numpy_ml.neural_nets.optimizers.SGD` optimizer with default

parameters. Default is None.

self.init = init

self.dilation = dilation

self.optimizer = optimizer

self.ch_residual = ch_residual

self.ch_dilation = ch_dilation

self.kernel_width = kernel_width

def _init_params(self):

self._dv = {}

self.conv_dilation = Conv1D(

pad="causal",

init=self.init,

dilation=self.dilation,

optimizer=self.optimizer,

act_fn=Affine(slope=1, intercept=0),

self.tanh = Tanh()

self.sigm = Sigmoid()

self.multiply_gate = Multiply(act_fn=Affine(slope=1, intercept=0))

self.conv_1x1 = Conv1D(

pad="same",

init=self.init,

optimizer=self.optimizer,

act_fn=Affine(slope=1, intercept=0),

self.add_residual = Add(act_fn=Affine(slope=1, intercept=0))

self.add_skip = Add(act_fn=Affine(slope=1, intercept=0))

@property

def parameters(self):

"""A dictionary of the module parameters."""

return {

"components": {

"conv_1x1": self.conv_1x1.parameters,

"add_skip": self.add_skip.parameters,

"add_residual": self.add_residual.parameters,

"conv_dilation": self.conv_dilation.parameters,

"multiply_gate": self.multiply_gate.parameters,

@property

def hyperparameters(self):

"""A dictionary of the module hyperparameters"""

return {

"layer": "WavenetResidualModule",

"init": self.init,

"dilation": self.dilation,

"optimizer": self.optimizer,

"ch_residual": self.ch_residual,

"ch_dilation": self.ch_dilation,

"kernel_width": self.kernel_width,

"component_ids": [

"components": {

"conv_1x1": self.conv_1x1.hyperparameters,

"add_skip": self.add_skip.hyperparameters,

"add_residual": self.add_residual.hyperparameters,

"conv_dilation": self.conv_dilation.hyperparameters,

"multiply_gate": self.multiply_gate.hyperparameters,

@property

def derived_variables(self):

"""A dictionary of intermediate values computed during the

forward/backward passes."""

dv = {

"conv_1x1_out": None,

"conv_dilation_out": None,

"multiply_gate_out": None,

"components": {

"conv_1x1": self.conv_1x1.derived_variables,

"add_skip": self.add_skip.derived_variables,

"add_residual": self.add_residual.derived_variables,

"conv_dilation": self.conv_dilation.derived_variables,

"multiply_gate": self.multiply_gate.derived_variables,

return dv

@property

def gradients(self):

"""A dictionary of the module parameter gradients."""

return {

"components": {

"conv_1x1": self.conv_1x1.gradients,

"add_skip": self.add_skip.gradients,

"add_residual": self.add_residual.gradients,

"conv_dilation": self.conv_dilation.gradients,

"multiply_gate": self.multiply_gate.gradients,

def forward(self, X_main, X_skip=None):

Compute the module output on a single minibatch.

Parameters

X_main : :py:class:`ndarray <numpy.ndarray>` of shape `(n_ex, in_rows, in_cols, in_ch)`

The input volume consisting of `n_ex` examples, each with dimension

(`in_rows`, `in_cols`, `in_ch`).

X_skip : :py:class:`ndarray <numpy.ndarray>` of shape `(n_ex, in_rows, in_cols, in_ch)`, or None

The output of the preceding skip-connection if this is not the

first module in the network.

Returns

Y_main : :py:class:`ndarray <numpy.ndarray>` of shape `(n_ex, out_rows, out_cols, out_ch)`

The output of the main pathway.

Y_skip : :py:class:`ndarray <numpy.ndarray>` of shape `(n_ex, out_rows, out_cols, out_ch)`

The output of the skip-connection pathway.

self.X_main, self.X_skip = X_main, X_skip

conv_dilation_out = self.conv_dilation.forward(X_main)

tanh_gate = self.tanh.fn(conv_dilation_out)

sigm_gate = self.sigm.fn(conv_dilation_out)

multiply_gate_out = self.multiply_gate.forward([tanh_gate, sigm_gate])

conv_1x1_out = self.conv_1x1.forward(multiply_gate_out)

# if this is the first wavenet block, initialize the "previous" skip

# connection sum to 0

self.X_skip = np.zeros_like(conv_1x1_out) if X_skip is None else X_skip

Y_skip = self.add_skip.forward([X_skip, conv_1x1_out])

Y_main = self.add_residual.forward([X_main, conv_1x1_out])

self._dv["tanh_out"] = tanh_gate

self._dv["sigm_out"] = sigm_gate

self._dv["conv_dilation_out"] = conv_dilation_out

self._dv["multiply_gate_out"] = multiply_gate_out

self._dv["conv_1x1_out"] = conv_1x1_out

return Y_main, Y_skip

def backward(self, dY_skip, dY_main=None):

dX_skip, dConv_1x1_out = self.add_skip.backward(dY_skip)

# if this is the last wavenet block, dY_main will be None. if not,

# calculate the error contribution from dY_main and add it to the

# contribution from the skip path

dX_main = np.zeros_like(self.X_main)

if dY_main is not None:

dX_main, dConv_1x1_main = self.add_residual.backward(dY_main)

dConv_1x1_out += dConv_1x1_main

dMultiply_out = self.conv_1x1.backward(dConv_1x1_out)

dTanh_out, dSigm_out = self.multiply_gate.backward(dMultiply_out)

conv_dilation_out = self.derived_variables["conv_dilation_out"]

dTanh_in = dTanh_out * self.tanh.grad(conv_dilation_out)

dSigm_in = dSigm_out * self.sigm.grad(conv_dilation_out)

dDilation_out = dTanh_in + dSigm_in

conv_back = self.conv_dilation.backward(dDilation_out)

dX_main += conv_back

self._dv["dLdTanh"] = dTanh_out

self._dv["dLdSigmoid"] = dSigm_out

self._dv["dLdConv_1x1"] = dConv_1x1_out

self._dv["dLdMultiply"] = dMultiply_out

self._dv["dLdConv_dilation"] = dDilation_out

return dX_main, dX_skip

class SkipConnectionIdentityModule(ModuleBase):

def __init__(

self,

optimizer=None,

A ResNet-like "identity" shortcut module.

Notes

The identity module enforces `same` padding during each convolution to

ensure module output has same dims as its input.

.. code-block:: text

X -> Conv2D -> Act_fn -> BatchNorm2D -> Conv2D -> BatchNorm2D -> + -> Act_fn

\______________________________________________________________/

References

.. [1] He et al. (2015). "Deep residual learning for image

recognition." https://arxiv.org/pdf/1512.03385.pdf

Parameters

out_ch : int

The number of filters/kernels to compute in the first convolutional

layer.

kernel_shape1 : 2-tuple

The dimension of a single 2D filter/kernel in the first

convolutional layer.

kernel_shape2 : 2-tuple

The dimension of a single 2D filter/kernel in the second

convolutional layer.

stride1 : int

The stride/hop of the convolution kernels in the first

convolutional layer. Default is 1.

stride2 : int

The stride/hop of the convolution kernels in the second

convolutional layer. Default is 1.

act_fn : :doc:`Activation <numpy_ml.neural_nets.activations>` object or None

The activation function for computing Y[t]. If None, use the

identity :math:`f(x) = x` by default. Default is None.

epsilon : float

A small smoothing constant to use during

:class:`~numpy_ml.neural_nets.layers.BatchNorm2D` computation to

avoid divide-by-zero errors. Default is 1e-5.

momentum : float

The momentum term for the running mean/running std calculations in

the :class:`~numpy_ml.neural_nets.layers.BatchNorm2D` layers.  The

closer this is to 1, the less weight will be given to the mean/std

of the current batch (i.e., higher smoothing). Default is 0.9.

The optimization strategy to use when performing gradient updates

within the :meth:`update` method.  If None, use the

:class:`~numpy_ml.neural_nets.optimizers.SGD` optimizer with

default parameters. Default is None.

The weight initialization strategy. Default is 'glorot_uniform'.

self.init = init

self.in_ch = None

self.out_ch = out_ch

self.epsilon = epsilon

self.stride1 = stride1

self.stride2 = stride2

self.optimizer = optimizer

self.momentum = momentum

self.kernel_shape1 = kernel_shape1

self.kernel_shape2 = kernel_shape2

self.act_fn = Affine(slope=1, intercept=0) if act_fn is None else act_fn

def _init_params(self):

self._dv = {}

self.conv1 = Conv2D(

pad="same",

init=self.init,

optimizer=self.optimizer,

# we can't initialize `conv2` without X's dimensions; see `forward`

# for further details

self.batchnorm1 = BatchNorm2D(epsilon=self.epsilon, momentum=self.momentum)

self.batchnorm2 = BatchNorm2D(epsilon=self.epsilon, momentum=self.momentum)

self.add3 = Add(self.act_fn)

def _init_conv2(self):

self.conv2 = Conv2D(

pad="same",

init=self.init,

optimizer=self.optimizer,

act_fn=Affine(slope=1, intercept=0),

@property

def parameters(self):

"""A dictionary of the module parameters."""

return {

"components": {

"add3": self.add3.parameters,

"conv1": self.conv1.parameters,

"conv2": self.conv2.parameters,

"batchnorm1": self.batchnorm1.parameters,

"batchnorm2": self.batchnorm2.parameters,

@property

def hyperparameters(self):

"""A dictionary of the module hyperparameters."""

return {

"layer": "SkipConnectionIdentityModule",

"init": self.init,

"in_ch": self.in_ch,

"out_ch": self.out_ch,

"epsilon": self.epsilon,

"stride1": self.stride1,

"stride2": self.stride2,

"momentum": self.momentum,

"optimizer": self.optimizer,

"act_fn": str(self.act_fn),

"kernel_shape1": self.kernel_shape1,

"kernel_shape2": self.kernel_shape2,

"component_ids": ["conv1", "batchnorm1", "conv2", "batchnorm2", "add3"],

"components": {

"add3": self.add3.hyperparameters,

"conv1": self.conv1.hyperparameters,

"conv2": self.conv2.hyperparameters,

"batchnorm1": self.batchnorm1.hyperparameters,

"batchnorm2": self.batchnorm2.hyperparameters,

@property

def derived_variables(self):

"""A dictionary of intermediate values computed during the

forward/backward passes."""

dv = {

"conv1_out": None,

"conv2_out": None,

"batchnorm1_out": None,

"batchnorm2_out": None,

"components": {

"add3": self.add3.derived_variables,

"conv1": self.conv1.derived_variables,

"conv2": self.conv2.derived_variables,

"batchnorm1": self.batchnorm1.derived_variables,

"batchnorm2": self.batchnorm2.derived_variables,

return dv

@property

def gradients(self):

"""A dictionary of the accumulated module parameter gradients."""

return {

"components": {

"add3": self.add3.gradients,

"conv1": self.conv1.gradients,

"conv2": self.conv2.gradients,

"batchnorm1": self.batchnorm1.gradients,

"batchnorm2": self.batchnorm2.gradients,

def forward(self, X, retain_derived=True):

Compute the module output given input volume `X`.

Parameters

The input volume consisting of `n_ex` examples, each with dimension

(`in_rows`, `in_cols`, `in_ch`).

retain_derived : bool

Whether to retain the variables calculated during the forward pass

for use later during backprop. If False, this suggests the layer

will not be expected to backprop through wrt. this input. Default

is True.

Returns

Y : :py:class:`ndarray <numpy.ndarray>` of shape (n_ex, out_rows, out_cols, out_ch)

The module output volume.

if not hasattr(self, "conv2"):

self.in_ch = X.shape[3]

conv1_out = self.conv1.forward(X, retain_derived)

bn1_out = self.batchnorm1.forward(conv1_out, retain_derived)

conv2_out = self.conv2.forward(bn1_out, retain_derived)

bn2_out = self.batchnorm2.forward(conv2_out, retain_derived)

Y = self.add3.forward([X, bn2_out], retain_derived)

if retain_derived:

self._dv["conv1_out"] = conv1_out

self._dv["conv2_out"] = conv2_out

self._dv["batchnorm1_out"] = bn1_out

self._dv["batchnorm2_out"] = bn2_out

return Y

def backward(self, dLdY, retain_grads=True):

Compute the gradient of the loss with respect to the layer parameters.

Parameters

dLdy : :py:class:`ndarray <numpy.ndarray>` of shape (`n_ex, out_rows, out_cols, out_ch`) or list of arrays

The gradient(s) of the loss with respect to the module output(s).

retain_grads : bool

Whether to include the intermediate parameter gradients computed

during the backward pass in the final parameter update. Default is

True.

Returns

dX : :py:class:`ndarray <numpy.ndarray>` of shape (n_ex, in_rows, in_cols, in_ch)

The gradient of the loss with respect to the module input volume.

dX, dBn2_out = self.add3.backward(dLdY, retain_grads)

dConv2_out = self.batchnorm2.backward(dBn2_out, retain_grads)

dBn1_out = self.conv2.backward(dConv2_out, retain_grads)

dConv1_out = self.batchnorm1.backward(dBn1_out, retain_grads)

dX += self.conv1.backward(dConv1_out, retain_grads)

self._dv["dLdAdd3_X"] = dX

self._dv["dLdBn2"] = dBn2_out

self._dv["dLdBn1"] = dBn1_out

self._dv["dLdConv2"] = dConv2_out

self._dv["dLdConv1"] = dConv1_out

return dX

class SkipConnectionConvModule(ModuleBase):

def __init__(

self,

optimizer=None,

A ResNet-like "convolution" shortcut module.

Notes

In contrast to :class:`SkipConnectionIdentityModule`, the additional

`conv2d_skip` and `batchnorm_skip` layers in the shortcut path allow

adjusting the dimensions of `X` to match the output of the main set of

convolutions.

.. code-block:: text

\_____________________ Conv2D -> Batchnorm2D __________________/

References

.. [1] He et al. (2015). "Deep residual learning for image

recognition." https://arxiv.org/pdf/1512.03385.pdf

Parameters

out_ch1 : int

The number of filters/kernels to compute in the first convolutional

layer.

out_ch2 : int

The number of filters/kernels to compute in the second

convolutional layer.

kernel_shape1 : 2-tuple

The dimension of a single 2D filter/kernel in the first

convolutional layer.

kernel_shape2 : 2-tuple

The dimension of a single 2D filter/kernel in the second

convolutional layer.

kernel_shape_skip : 2-tuple

The dimension of a single 2D filter/kernel in the "skip"

convolutional layer.

stride1 : int

The stride/hop of the convolution kernels in the first

convolutional layer. Default is 1.

stride2 : int

The stride/hop of the convolution kernels in the second

convolutional layer. Default is 1.

stride_skip : int

The stride/hop of the convolution kernels in the "skip"

convolutional layer. Default is 1.

pad1 : int, tuple, or 'same'

The number of rows/columns of 0's to pad the input to the first

convolutional layer with. Default is 0.

pad2 : int, tuple, or 'same'

The number of rows/columns of 0's to pad the input to the second

convolutional layer with. Default is 0.

The activation function for computing ``Y[t]``. If None, use the

identity :math:`f(x) = x` by default. Default is None.

epsilon : float

A small smoothing constant to use during

:class:`~numpy_ml.neural_nets.layers.BatchNorm2D` computation to

avoid divide-by-zero errors. Default is 1e-5.

momentum : float

The momentum term for the running mean/running std calculations in

the :class:`~numpy_ml.neural_nets.layers.BatchNorm2D` layers.  The

closer this is to 1, the less weight will be given to the mean/std

of the current batch (i.e., higher smoothing). Default is 0.9.

init : str

The weight initialization strategy. Valid entries are

{'glorot_normal', 'glorot_uniform', 'he_normal', 'he_uniform'}.

The optimization strategy to use when performing gradient updates

within the :class:`update` method.  If None, use the

:class:`~numpy_ml.neural_nets.optimizers.SGD` optimizer with

default parameters. Default is None.

self.init = init

self.pad1 = pad1

self.pad2 = pad2

self.in_ch = None

self.out_ch1 = out_ch1

self.out_ch2 = out_ch2

self.epsilon = epsilon

self.stride1 = stride1

self.stride2 = stride2

self.momentum = momentum

self.optimizer = optimizer

self.stride_skip = stride_skip

self.kernel_shape1 = kernel_shape1

self.kernel_shape2 = kernel_shape2

self.kernel_shape_skip = kernel_shape_skip

def _init_params(self, X=None):

self._dv = {}

self.conv1 = Conv2D(

init=self.init,

optimizer=self.optimizer,

self.conv2 = Conv2D(

init=self.init,

optimizer=self.optimizer,

act_fn=Affine(slope=1, intercept=0),

# we can't initialize `conv_skip` without X's dimensions; see `forward`

# for further details

self.batchnorm_skip = BatchNorm2D(epsilon=self.epsilon, momentum=self.momentum)

self.add3 = Add(self.act_fn)

def _calc_skip_padding(self, X):

pads = []

for p in [self.pad1, self.pad2]:

if isinstance(p, int):

pads.append((p, p, p, p))

elif isinstance(p, tuple) and len(p) == 2:

pads.append((p[0], p[0], p[1], p[1]))

self.pad1, self.pad2 = pads

# compute the dimensions of the convolution1 output

s1 = self.stride1

fr1, fc1 = self.kernel_shape1

_, in_rows, in_cols, _ = X.shape

pr11, pr12, pc11, pc12 = self.pad1

out_rows1 = np.floor(1 + (in_rows + pr11 + pr12 - fr1) / s1).astype(int)

out_cols1 = np.floor(1 + (in_cols + pc11 + pc12 - fc1) / s1).astype(int)

# compute the dimensions of the convolution2 output

s2 = self.stride2

fr2, fc2 = self.kernel_shape2

pr21, pr22, pc21, pc22 = self.pad2

out_rows2 = np.floor(1 + (out_rows1 + pr21 + pr22 - fr2) / s2).astype(int)

out_cols2 = np.floor(1 + (out_cols1 + pc21 + pc22 - fc2) / s2).astype(int)

# finally, compute the appropriate padding dims for the skip convolution

desired_dims = (out_rows2, out_cols2)

self.pad_skip = calc_pad_dims_2D(

X.shape,

def _init_conv_skip(self, X):

self.conv_skip = Conv2D(

init=self.init,

act_fn=Affine(slope=1, intercept=0),

optimizer=self.optimizer,

@property

def parameters(self):

"""A dictionary of the module parameters."""

return {

"components": {

"add3": self.add3.parameters,

"conv1": self.conv1.parameters,

"conv2": self.conv2.parameters,

"conv_skip": self.conv_skip.parameters

if hasattr(self, "conv_skip")

else None,

"batchnorm1": self.batchnorm1.parameters,

"batchnorm2": self.batchnorm2.parameters,

"batchnorm_skip": self.batchnorm_skip.parameters,

@property

def hyperparameters(self):

"""A dictionary of the module hyperparameters."""

return {

"layer": "SkipConnectionConvModule",

"init": self.init,

"pad1": self.pad1,

"pad2": self.pad2,

"in_ch": self.in_ch,

"out_ch1": self.out_ch1,

"out_ch2": self.out_ch2,

"epsilon": self.epsilon,

"stride1": self.stride1,

"stride2": self.stride2,

"momentum": self.momentum,

"act_fn": str(self.act_fn),

"stride_skip": self.stride_skip,

"kernel_shape1": self.kernel_shape1,

"kernel_shape2": self.kernel_shape2,

"kernel_shape_skip": self.kernel_shape_skip,

"pad_skip": self.pad_skip if hasattr(self, "pad_skip") else None,

"component_ids": [

"components": {

"add3": self.add3.hyperparameters,

"conv1": self.conv1.hyperparameters,

"conv2": self.conv2.hyperparameters,

"conv_skip": self.conv_skip.hyperparameters

if hasattr(self, "conv_skip")

else None,

"batchnorm1": self.batchnorm1.hyperparameters,

"batchnorm2": self.batchnorm2.hyperparameters,

"batchnorm_skip": self.batchnorm_skip.hyperparameters,

@property

def derived_variables(self):

"""A dictionary of intermediate values computed during the

forward/backward passes."""

dv = {

"conv1_out": None,

"conv2_out": None,

"conv_skip_out": None,

"batchnorm1_out": None,

"batchnorm2_out": None,

"batchnorm_skip_out": None,

"components": {

"add3": self.add3.derived_variables,

"conv1": self.conv1.derived_variables,

"conv2": self.conv2.derived_variables,

"conv_skip": self.conv_skip.derived_variables

if hasattr(self, "conv_skip")

else None,

"batchnorm1": self.batchnorm1.derived_variables,

"batchnorm2": self.batchnorm2.derived_variables,

"batchnorm_skip": self.batchnorm_skip.derived_variables,

return dv

@property

def gradients(self):

"""A dictionary of the accumulated module parameter gradients."""

return {

"components": {

"add3": self.add3.gradients,

"conv1": self.conv1.gradients,

"conv2": self.conv2.gradients,

"conv_skip": self.conv_skip.gradients

if hasattr(self, "conv_skip")

else None,

"batchnorm1": self.batchnorm1.gradients,

"batchnorm2": self.batchnorm2.gradients,

"batchnorm_skip": self.batchnorm_skip.gradients,

def forward(self, X, retain_derived=True):

Compute the layer output given input volume `X`.

Parameters

The input volume consisting of `n_ex` examples, each with dimension

(`in_rows`, `in_cols`, `in_ch`).

retain_derived : bool

Whether to retain the variables calculated during the forward pass

for use later during backprop. If False, this suggests the layer

will not be expected to backprop through wrt. this input. Default

is True.

Returns

Y : :py:class:`ndarray <numpy.ndarray>` of shape `(n_ex, out_rows, out_cols, out_ch)`

The module output volume.

# now that we have the input dims for X we can initialize the proper

# padding in the `conv_skip` layer

if not hasattr(self, "conv_skip"):

self.in_ch = X.shape[3]

conv1_out = self.conv1.forward(X, retain_derived)

bn1_out = self.batchnorm1.forward(conv1_out, retain_derived)

conv2_out = self.conv2.forward(bn1_out, retain_derived)

bn2_out = self.batchnorm2.forward(conv2_out, retain_derived)

conv_skip_out = self.conv_skip.forward(X, retain_derived)

bn_skip_out = self.batchnorm_skip.forward(conv_skip_out, retain_derived)

Y = self.add3.forward([bn_skip_out, bn2_out], retain_derived)

if retain_derived:

self._dv["conv1_out"] = conv1_out

self._dv["conv2_out"] = conv2_out

self._dv["batchnorm1_out"] = bn1_out

self._dv["batchnorm2_out"] = bn2_out

self._dv["conv_skip_out"] = conv_skip_out

self._dv["batchnorm_skip_out"] = bn_skip_out

return Y

def backward(self, dLdY, retain_grads=True):

Compute the gradient of the loss with respect to the module parameters.

Parameters

dLdy : :py:class:`ndarray <numpy.ndarray>` of shape `(n_ex, out_rows, out_cols, out_ch)`

or list of arrays

The gradient(s) of the loss with respect to the module output(s).

retain_grads : bool

Whether to include the intermediate parameter gradients computed

during the backward pass in the final parameter update. Default is

True.

Returns

dX : :py:class:`ndarray <numpy.ndarray>` of shape `(n_ex, in_rows, in_cols, in_ch)`

The gradient of the loss with respect to the module input volume.

dBnskip_out, dBn2_out = self.add3.backward(dLdY)

dConvskip_out = self.batchnorm_skip.backward(dBnskip_out)

dX = self.conv_skip.backward(dConvskip_out)

dConv2_out = self.batchnorm2.backward(dBn2_out)

dBn1_out = self.conv2.backward(dConv2_out)

dConv1_out = self.batchnorm1.backward(dBn1_out)

dX += self.conv1.backward(dConv1_out)

if retain_grads:

self._dv["dLdAdd3_X"] = dX

self._dv["dLdBn1"] = dBn1_out

self._dv["dLdBn2"] = dBn2_out

self._dv["dLdConv1"] = dConv1_out

self._dv["dLdConv2"] = dConv2_out

self._dv["dLdBnSkip"] = dBnskip_out

self._dv["dLdConvSkip"] = dConvskip_out

return dX

class BidirectionalLSTM(ModuleBase):

def __init__(

self,

optimizer=None,

A single bidirectional long short-term memory (LSTM) layer.

Parameters

n_out : int

The dimension of a single hidden state / output on a given timestep

The activation function for computing ``A[t]``. If not specified,

use :class:`~numpy_ml.neural_nets.activations.Tanh` by default.

gate_fn : :doc:`Activation <numpy_ml.neural_nets.activations>` object or None

The gate function for computing the update, forget, and output

gates. If not specified, use

:class:`~numpy_ml.neural_nets.activations.Sigmoid` by default.

merge_mode : {"sum", "multiply", "concat", "average"}

Mode by which outputs of the forward and backward LSTMs will be

combined. Default is 'concat'.

The optimization strategy to use when performing gradient updates

within the `update` method.  If None, use the

:class:`~numpy_ml.neural_nets.optimizers.SGD` optimizer with

default parameters. Default is None.

The weight initialization strategy. Default is 'glorot_uniform'.

self.init = init

self.n_in = None

self.n_out = n_out

self.optimizer = optimizer

self.merge_mode = merge_mode

self.act_fn = Tanh() if act_fn is None else act_fn

self.gate_fn = Sigmoid() if gate_fn is None else gate_fn

def _init_params(self):

self.cell_fwd = LSTMCell(

init=self.init,

optimizer=self.optimizer,

self.cell_bwd = LSTMCell(

init=self.init,

optimizer=self.optimizer,

def forward(self, X):

Run a forward pass across all timesteps in the input.

Parameters

X : :py:class:`ndarray <numpy.ndarray>` of shape `(n_ex, n_in, n_t)`

Input consisting of `n_ex` examples each of dimensionality `n_in`

and extending for `n_t` timesteps.

Returns

Y : :py:class:`ndarray <numpy.ndarray>` of shape `(n_ex, n_out, n_t)`

The value of the hidden state for each of the `n_ex` examples

across each of the `n_t` timesteps.

Y_fwd, Y_bwd, Y = [], [], []

n_ex, self.n_in, n_t = X.shape

# forward LSTM

for t in range(n_t):

yt, ct = self.cell_fwd.forward(X[:, :, t])

# backward LSTM

for t in reversed(range(n_t)):

yt, ct = self.cell_bwd.forward(X[:, :, t])

Y_bwd.insert(0, yt)

# merge forward and backward states

for t in range(n_t):

if self.merge_mode == "concat":

Y.append(np.concatenate([Y_fwd[t], Y_bwd[t]], axis=1))

elif self.merge_mode == "sum":

Y.append(Y_fwd[t] + Y_bwd[t])

elif self.merge_mode == "average":

Y.append((Y_fwd[t] + Y_bwd[t]) / 2)

elif self.merge_mode == "multiply":

Y.append(Y_fwd[t] * Y_bwd[t])

self.Y_fwd, self.Y_bwd = Y_fwd, Y_bwd

return np.dstack(Y)

def backward(self, dLdA):

Run a backward pass across all timesteps in the input.

Parameters

dLdA : :py:class:`ndarray <numpy.ndarray>` of shape `(n_ex, n_out, n_t)`

The gradient of the loss with respect to the layer output for each

of the `n_ex` examples across all `n_t` timesteps.

Returns

dLdX : :py:class:`ndarray <numpy.ndarray>` of shape `(n_ex, n_in, n_t)`

across each of the `n_t` timesteps.

assert self.trainable, "Layer is frozen"

n_ex, n_out, n_t = dLdA.shape

dLdX_f, dLdX_b, dLdX = [], [], []

# forward LSTM

for t in reversed(range(n_t)):

if self.merge_mode == "concat":

dLdXt_f = self.cell_fwd.backward(dLdA[:, : self.n_out, t])

elif self.merge_mode == "sum":

dLdXt_f = self.cell_fwd.backward(dLdA[:, :, t])

elif self.merge_mode == "multiplty":

dLdXt_f = self.cell_fwd.backward(dLdA[:, :, t] * self.Y_bwd[t])

elif self.merge_mode == "average":

dLdXt_f = self.cell_fwd.backward(dLdA[:, :, t] * 0.5)

dLdX_f.insert(0, dLdXt_f)

# backward LSTM

for t in range(n_t):

if self.merge_mode == "concat":

dLdXt_b = self.cell_bwd.backward(dLdA[:, self.n_out :, t])

elif self.merge_mode == "sum":

dLdXt_b = self.cell_bwd.backward(dLdA[:, :, t])

elif self.merge_mode == "multiplty":

dLdXt_b = self.cell_bwd.backward(dLdA[:, :, t] * self.Y_fwd[t])

elif self.merge_mode == "average":

dLdXt_b = self.cell_bwd.backward(dLdA[:, :, t] * 0.5)

for t in range(n_t):

dLdX.append(dLdX_f[t] + dLdX_b[t])

return np.dstack(dLdX)

@property

def derived_variables(self):

"""A dictionary of intermediate values computed during the

forward/backward passes."""

return {

"components": {

"cell_fwd": self.cell_fwd.derived_variables,

"cell_bwd": self.cell_bwd.derived_variables,

@property

def gradients(self):

"""A dictionary of the accumulated module parameter gradients."""

return {

"components": {

"cell_fwd": self.cell_fwd.gradients,

"cell_bwd": self.cell_bwd.gradients,

@property

def parameters(self):

"""A dictionary of the module parameters."""

return {

"components": {

"cell_fwd": self.cell_fwd.parameters,

"cell_bwd": self.cell_bwd.parameters,

@property

def hyperparameters(self):

"""A dictionary of the module hyperparameters."""

return {

"layer": "BidirectionalLSTM",

"init": self.init,

"n_in": self.n_in,

"n_out": self.n_out,

"act_fn": str(self.act_fn),

"optimizer": self.optimizer,

"merge_mode": self.merge_mode,

"component_ids": ["cell_fwd", "cell_bwd"],

"components": {

"cell_fwd": self.cell_fwd.hyperparameters,

"cell_bwd": self.cell_bwd.hyperparameters,

class MultiHeadedAttentionModule(ModuleBase):

def __init__(self, n_heads=8, dropout_p=0, init="glorot_uniform", optimizer=None):

A mutli-headed attention module.

Notes

Multi-head attention allows a model to jointly attend to information from

different representation subspaces at different positions. With a

single head, this information would get averaged away when the

attention weights are combined with the value

.. math::

\\text{MultiHead}(\mathbf{Q}, \mathbf{K}, \mathbf{V})

= [\\text{head}_1; ...; \\text{head}_h] \\mathbf{W}^{(O)}

where

.. math::

\\text{head}_i = \\text{SDP_attention}(

\mathbf{Q W}_i^{(Q)}, \mathbf{K W}_i^{(K)}, \mathbf{V W}_i^{(V)})

and the projection weights are parameter matrices:

.. math::

\mathbf{W}_i^{(Q)}  &\in

\mathbb{R}^{(\\text{kqv_dim} \ \\times \ \\text{latent_dim})} \\\\

\mathbf{W}_i^{(K)}  &\in

\mathbf{W}_i^{(V)}  &\in

\mathbf{W}^{(O)}  &\in

\mathbb{R}^{(\\text{n_heads} \cdot \\text{latent_dim} \ \\times \ \\text{kqv_dim})}

Importantly, the current module explicitly assumes that

.. math::

\\text{kqv_dim} = \\text{dim(query)} = \\text{dim(keys)} = \\text{dim(values)}

and that

.. math::

\\text{latent_dim} = \\text{kqv_dim / n_heads}

**[MH Attention Head h]**:

.. code-block:: text

K --> W_h^(K) ------\\

V --> W_h^(V) ------- > DP_Attention --> head_h

Q --> W_h^(Q) ------/

The full **[MultiHeadedAttentionModule]** then becomes

.. code-block:: text

K --> | [Attn Head 1] | --> head_1 --\\

V --> | [Attn Head 2] | --> head_2 --\\

Q --> |      ...      |      ...       --> Concat --> W^(O) --> MH_out

| [Attn Head Z] | --> head_Z --/

Due to the reduced dimension of each head, the total computational cost

is similar to that of a single attention head with full (i.e., kqv_dim)

dimensionality.

Parameters

n_heads : int

The number of simultaneous attention heads to use. Note that the

larger `n_heads`, the smaller the dimensionality of any single

head, since ``latent_dim = kqv_dim / n_heads``. Default is 8.

dropout_p : float in [0, 1)

The dropout propbability during training, applied to the output of

the softmax in each dot-product attention head. If 0, no dropout is

applied. Default is 0.

init : {'glorot_normal', 'glorot_uniform', 'he_normal', 'he_uniform'}

The weight initialization strategy. Default is 'glorot_uniform'.

optimizer : str, :doc:`Optimizer <numpy_ml.neural_nets.optimizers>` object, or None

The optimization strategy to use when performing gradient updates

within the :meth:`update` method.  If None, use the

:class:`~numpy_ml.neural_nets.optimizers.SGD` optimizer with default

parameters. Default is None.

self.init = init

self.kqv_dim = None

self.projections = {}

self.n_heads = n_heads

self.optimizer = optimizer

self.dropout_p = dropout_p

self.is_initialized = False

def _init_params(self):

self._dv = {}

# assume dim(keys) = dim(query) = dim(values)

assert self.kqv_dim % self.n_heads == 0

self.latent_dim = self.kqv_dim // self.n_heads

self.attention = DotProductAttention(scale=True, dropout_p=self.dropout_p)

self.projections = {

k: Dropout(

FullyConnected(

init=self.init,

optimizer=self.optimizer,

act_fn="Affine(slope=1, intercept=0)",

for k in ["Q", "K", "V", "O"]

self.is_initialized = True

def forward(self, Q, K, V):

if not self.is_initialized:

self.kqv_dim = Q.shape[-1]

# project queries, keys, and values into the `latent_dim`-dimensional subspace

n_ex = Q.shape[0]

for k, x in zip(["Q", "K", "V"], [Q, K, V]):

proj = self.projections[k].forward(x)

proj = proj.reshape(n_ex, -1, self.n_heads, self.latent_dim).swapaxes(1, 2)

self._dv["{}_proj".format(k)] = proj

dv = self.derived_variables

Q_proj, K_proj, V_proj = dv["Q_proj"], dv["K_proj"], dv["V_proj"]

# apply scaled dot-product attention to the projected vectors

attn = self.attention

attn_out = attn.forward(Q_proj, K_proj, V_proj)

self._dv["attention_weights"] = attn.derived_variables["attention_weights"]

# concatenate the different heads using `reshape` to create an

# `kqv_dim`-dim vector

attn_out = attn_out.swapaxes(1, 2).reshape(n_ex, self.kqv_dim)

self._dv["attention_out"] = attn_out.reshape(n_ex, -1, self.kqv_dim)

# apply the final output projection

Y = self.projections["O"].forward(attn_out)

Y = Y.reshape(n_ex, -1, self.kqv_dim)

return Y

def backward(self, dLdy):

n_ex = dLdy.shape[0]

dLdy = dLdy.reshape(n_ex, self.kqv_dim)

dLdX = self.projections["O"].backward(dLdy)

dLdX = dLdX.reshape(n_ex, self.n_heads, -1, self.latent_dim)

dLdQ_proj, dLdK_proj, dLdV_proj = self.attention.backward(dLdX)

self._dv["dQ_proj"] = dLdQ_proj

self._dv["dK_proj"] = dLdK_proj

self._dv["dV_proj"] = dLdV_proj

dLdQ_proj = dLdQ_proj.reshape(n_ex, self.kqv_dim)

dLdK_proj = dLdK_proj.reshape(n_ex, self.kqv_dim)

dLdV_proj = dLdV_proj.reshape(n_ex, self.kqv_dim)

dLdQ = self.projections["Q"].backward(dLdQ_proj)

dLdK = self.projections["K"].backward(dLdK_proj)

dLdV = self.projections["V"].backward(dLdV_proj)

return dLdQ, dLdK, dLdV

@property

def derived_variables(self):

"""A dictionary of intermediate values computed during the

forward/backward passes."""

dv = {

"Q_proj": None,

"K_proj": None,

"V_proj": None,

"components": {

"Q": self.projections["Q"].derived_variables,

"K": self.projections["K"].derived_variables,

"V": self.projections["V"].derived_variables,

"O": self.projections["O"].derived_variables,

"attention": self.attention.derived_variables,

return dv

@property

def gradients(self):

"""A dictionary of the accumulated module parameter gradients."""

return {

"components": {

"Q": self.projections["Q"].gradients,

"K": self.projections["K"].gradients,

"V": self.projections["V"].gradients,

"O": self.projections["O"].gradients,

"attention": self.attention.gradients,

@property

def parameters(self):

"""A dictionary of the module parameters."""

return {

"components": {

"Q": self.projections["Q"].parameters,

"K": self.projections["K"].parameters,

"V": self.projections["V"].parameters,

"O": self.projections["O"].parameters,

"attention": self.attention.parameters,

@property

def hyperparameters(self):

"""A dictionary of the module hyperparameters."""

return {

"layer": "MultiHeadedAttentionModule",

"init": self.init,

"kqv_dim": self.kqv_dim,

"latent_dim": self.latent_dim,

"n_heads": self.n_heads,

"dropout_p": self.dropout_p,

"component_ids": ["attention", "Q", "K", "V", "O"],

"components": {

"Q": self.projections["Q"].hyperparameters,

"K": self.projections["K"].hyperparameters,

"V": self.projections["V"].hyperparameters,

"O": self.projections["O"].hyperparameters,

"attention": self.attention.hyperparameters,


from .schedulers import *


from copy import deepcopy

from abc import ABC, abstractmethod

import numpy as np

from math import erf

def gaussian_cdf(x, mean, var):

Compute the probability that a random draw from a 1D Gaussian with mean

`mean` and variance `var` is less than or equal to `x`.

eps = np.finfo(float).eps

x_scaled = (x - mean) / np.sqrt(var + eps)

return (1 + erf(x_scaled / np.sqrt(2))) / 2

class SchedulerBase(ABC):

def __init__(self):

"""Abstract base class for all Scheduler objects."""

self.hyperparameters = {}

def __call__(self, step=None, cur_loss=None):

return self.learning_rate(step=step, cur_loss=cur_loss)

def copy(self):

"""Return a copy of the current object."""

return deepcopy(self)

def set_params(self, hparam_dict):

"""Set the scheduler hyperparameters from a dictionary."""

if hparam_dict is not None:

for k, v in hparam_dict.items():

if k in self.hyperparameters:

self.hyperparameters[k] = v

@abstractmethod

def learning_rate(self, step=None):

raise NotImplementedError

class ConstantScheduler(SchedulerBase):

def __init__(self, lr=0.01, **kwargs):

Returns a fixed learning rate, regardless of the current step.

Parameters

initial_lr : float

The learning rate. Default is 0.01

self.lr = lr

self.hyperparameters = {"id": "ConstantScheduler", "lr": self.lr}

def __str__(self):

return "ConstantScheduler(lr={})".format(self.lr)

def learning_rate(self, **kwargs):

Return the current learning rate.

Returns

lr : float

The learning rate

return self.lr

class ExponentialScheduler(SchedulerBase):

def __init__(

self, initial_lr=0.01, stage_length=500, staircase=False, decay=0.1, **kwargs

An exponential learning rate scheduler.

Notes

The exponential scheduler decays the learning rate by `decay` every

`stage_length` steps, starting from `initial_lr`::

learning_rate = initial_lr * decay ** curr_stage

where::

curr_stage = step / stage_length          if staircase = False

curr_stage = floor(step / stage_length)   if staircase = True

Parameters

initial_lr : float

The learning rate at the first step. Default is 0.01.

stage_length : int

The length of each stage, in steps. Default is 500.

staircase : bool

If True, only adjusts the learning rate at the stage transitions,

producing a step-like decay schedule. If False, adjusts the

learning rate after each step, creating a smooth decay schedule.

Default is False.

decay : float

The amount to decay the learning rate at each new stage. Default is

self.decay = decay

self.staircase = staircase

self.initial_lr = initial_lr

self.stage_length = stage_length

self.hyperparameters = {

"id": "StepScheduler",

"decay": self.decay,

"staircase": self.staircase,

"initial_lr": self.initial_lr,

"stage_length": self.stage_length,

def __str__(self):

return "ExponentialScheduler(initial_lr={}, stage_length={}, staircase={}, decay={})".format(

self.initial_lr, self.stage_length, self.staircase, self.decay

def learning_rate(self, step, **kwargs):

Return the current learning rate as a function of `step`.

Parameters

step : int

The current step number.

Returns

lr : float

The learning rate for the current step.

cur_stage = step / self.stage_length

if self.staircase:

cur_stage = np.floor(cur_stage)

return self.initial_lr * self.decay ** cur_stage

class NoamScheduler(SchedulerBase):

def __init__(self, model_dim=512, scale_factor=1, warmup_steps=4000, **kwargs):

The Noam learning rate scheduler, originally used in conjunction with

the Adam optimizer in [1].

Notes

The Noam scheduler increases the learning rate linearly for the first

`warmup_steps` steps, and decreases it thereafter proportionally to the

inverse square root of the step number::

lr = scale_factor * ( (model_dim ** (-0.5)) * adj_step )

adj_step = min(step_num ** (-0.5), step_num * warmup_steps ** (-1.5))

References

.. [1] Vaswani et al. (2017) "Attention is all you need". *31st

Conference on Neural Information Processing Systems*,

Parameters

model_dim : int

The number of units in the layer output. Default is 512.

scale_factor : float

A fixed coefficient for rescaling the final learning rate. Default

is 1.

warmup_steps : int

The number of steps in the warmup stage of training. Default is

self.model_dim = model_dim

self.scale_factor = scale_factor

self.warmup_steps = warmup_steps

self.hyperparameters = {

"id": "NoamScheduler",

"model_dim": self.model_dim,

"scale_factor": self.scale_factor,

"warmup_steps": self.warmup_steps,

def __str__(self):

return "NoamScheduler(model_dim={}, scale_factor={}, warmup_steps={})".format(

self.model_dim, self.scale_factor, self.warmup_steps

def learning_rate(self, step, **kwargs):

warmup, d_model = self.warmup_steps, self.model_dim

new_lr = d_model ** (-0.5) * min(step ** (-0.5), step * warmup ** (-1.5))

return self.scale_factor * new_lr

class KingScheduler(SchedulerBase):

def __init__(self, initial_lr=0.01, patience=1000, decay=0.99, **kwargs):

The Davis King / DLib learning rate scheduler.

Notes

The KingScheduler computes the probability that the slope of the OLS

fit to the loss history is negative. If the probability that it is

negative is less than 51% over the last `patience` steps, the scheduler

exponentially decreases the current learning rate by `decay`.

References

.. [1] King, D. (2018). "Automatic learning rate scheduling that really

works". http://blog.dlib.net/2018/02/automatic-learning-rate-scheduling-that.html

Parameters

initial_lr : float

The learning rate to begin at. Default is 0.01.

patience : int

Amount of time to maintain the current learning rate without a

decrease in loss before adjustment. Default is 1000.

decay : float

self.decay = decay

self.patience = patience

self.initial_lr = initial_lr

self.current_lr = initial_lr

self.max_history = np.ceil(1.1 * (patience + 1)).astype(int)

self.loss_history = []

self.hyperparameters = {

"id": "KingScheduler",

"decay": self.decay,

"patience": self.patience,

"initial_lr": self.initial_lr,

def __str__(self):

return "KingScheduler(initial_lr={}, patience={}, decay={})".format(

self.initial_lr, self.patience, self.decay

def _steps_without_decrease(self, robust=False, check_all=False):

Returns the maximum number of timesteps for which `P(loss is decreasing)

< 0.51`.

Parameters

robust : bool

If `robust=True`, first filter out the largest 10% of the loss

values to remove transient spikes in the loss due to, e.g., a few

bad minibatches. Default is False.

check_all : bool

If False, returns the maximum number of timesteps for which P(loss

is decreasing) < 0.51. If True, only checks whether the number of

timesteps for which P(loss is decreasing) < 0.51 is equal to

``self.patience``. The former provides more information but is

significantly more computationally expensive.  Default is False.

Returns

steps_without_decrease: int

The maximum number of steps back in loss_history for which P(loss

is decreasing) < 0.51.

lh = np.array(self.loss_history)

# drop top 10% of loss values to filter out large loss spikes

if robust:

thresh = np.quantile(lh, 0.9)

lh = np.array([i for i in lh if i <= thresh])

N = len(lh)

steps_without_decrease = 0

if check_all:

for i in reversed(range(N - 2)):

if self._p_decreasing(lh, i) < 0.51:

steps_without_decrease = N - i

else:

i = max(0, N - self.patience - 1)

if self._p_decreasing(lh, i) < 0.51:

steps_without_decrease = N - i

return steps_without_decrease

def _p_decreasing(self, loss_history, i):

Compute the probability that the slope of the OLS fit to the loss

history is negative.

Parameters

loss_history : numpy array of shape (N,)

The sequence of loss values for the previous `N` minibatches.

i : int

Compute P(Slope < 0) beginning at index i in `history`.

Returns

p_decreasing : float

The probability that the slope of the OLS fit to loss_history is

less than or equal to 0.

loss = loss_history[i:]

N = len(loss)

# perform OLS on the loss entries to calc the slope mean

X = np.c_[np.ones(N), np.arange(i, len(loss_history))]

intercept, s_mean = np.linalg.inv(X.T @ X) @ X.T @ loss

loss_pred = s_mean * X[:, 1] + intercept

# compute the variance of our loss predictions and use this to compute

# the (unbiased) estimate of the slope variance

loss_var = 1 / (N - 2) * np.sum((loss - loss_pred) ** 2)

s_var = (12 * loss_var) / (N ** 3 - N)

# compute the probability that a random sample from a Gaussian

# parameterized by s_mean and s_var is less than or equal to 0

p_decreasing = gaussian_cdf(0, s_mean, s_var)

return p_decreasing

def learning_rate(self, step, cur_loss):

Compute the updated learning rate for the current step and loss.

Parameters

step : int

The current step number. Unused.

cur_loss : float

The loss at the current step.

Returns

lr : float

The learning rate for the current step.

if cur_loss is None:

raise ValueError("cur_loss must be a float, but got {}".format(cur_loss))

# this happens if we initialize the scheduler from a string / dict

if not hasattr(self, "max_history"):

self.max_history = np.ceil(1.1 * (self.patience + 1)).astype(int)

patience, max_history = self.patience, self.max_history

if len(self.loss_history) < patience:

return self.current_lr

self.loss_history = self.loss_history[-max_history:]

# if the loss has not decreased for `patience` timesteps, drop the

# learning rate

if (

self._steps_without_decrease() > patience

and self._steps_without_decrease(robust=True) > patience

self.current_lr *= self.decay

return self.current_lr


from .initializers import *


"""A module containing objects to instantiate various neural network components."""

import re

from functools import partial

from ast import literal_eval as _eval

import numpy as np

from ..optimizers import OptimizerBase, SGD, AdaGrad, RMSProp, Adam

from ..activations import (

ELU,

GELU,

SELU,

ReLU,

Tanh,

Affine,

Sigmoid,

Identity,

SoftPlus,

LeakyReLU,

Exponential,

HardSigmoid,

ActivationBase,

from ..schedulers import (

SchedulerBase,

ConstantScheduler,

ExponentialScheduler,

NoamScheduler,

KingScheduler,

from ..utils import (

class ActivationInitializer(object):

def __init__(self, param=None):

A class for initializing activation functions. Valid `param` values

are:

(a) ``__str__`` representations of an `ActivationBase` instance

(b) `ActivationBase` instance

If `param` is `None`, return the identity function: f(X) = X

self.param = param

def __call__(self):

"""Initialize activation function"""

param = self.param

if param is None:

act = Identity()

elif isinstance(param, ActivationBase):

act = param

elif isinstance(param, str):

act = self.init_from_str(param)

else:

raise ValueError("Unknown activation: {}".format(param))

return act

def init_from_str(self, act_str):

"""Initialize activation function from the `param` string"""

act_str = act_str.lower()

if act_str == "relu":

act_fn = ReLU()

elif act_str == "tanh":

act_fn = Tanh()

elif act_str == "selu":

act_fn = SELU()

elif act_str == "sigmoid":

act_fn = Sigmoid()

elif act_str == "identity":

act_fn = Identity()

elif act_str == "hardsigmoid":

act_fn = HardSigmoid()

elif act_str == "softplus":

act_fn = SoftPlus()

elif act_str == "exponential":

act_fn = Exponential()

elif "affine" in act_str:

r = r"affine\(slope=(.*), intercept=(.*)\)"

slope, intercept = re.match(r, act_str).groups()

act_fn = Affine(float(slope), float(intercept))

elif "leaky relu" in act_str:

r = r"leaky relu\(alpha=(.*)\)"

alpha = re.match(r, act_str).groups()[0]

act_fn = LeakyReLU(float(alpha))

elif "gelu" in act_str:

r = r"gelu\(approximate=(.*)\)"

approx = re.match(r, act_str).groups()[0] == "true"

act_fn = GELU(approximation=approx)

elif "elu" in act_str:

r = r"elu\(alpha=(.*)\)"

act_fn = ELU(alpha=float(alpha))

else:

raise ValueError("Unknown activation: {}".format(act_str))

return act_fn

class SchedulerInitializer(object):

def __init__(self, param=None, lr=None):

A class for initializing learning rate schedulers. Valid `param` values

are:

(a) __str__ representations of `SchedulerBase` instances

(b) `SchedulerBase` instances

(c) Parameter dicts (e.g., as produced via the `summary` method in

`LayerBase` instances)

If `param` is `None`, return the ConstantScheduler with learning rate

equal to `lr`.

if all([lr is None, param is None]):

raise ValueError("lr and param cannot both be `None`")

self.lr = lr

self.param = param

def __call__(self):

"""Initialize scheduler"""

param = self.param

if param is None:

scheduler = ConstantScheduler(self.lr)

elif isinstance(param, SchedulerBase):

scheduler = param

elif isinstance(param, str):

scheduler = self.init_from_str()

elif isinstance(param, dict):

scheduler = self.init_from_dict()

return scheduler

def init_from_str(self):

"""Initialize scheduler from the param string"""

r = r"([a-zA-Z]*)=([^,)]*)"

sch_str = self.param.lower()

kwargs = {i: _eval(j) for i, j in re.findall(r, sch_str)}

if "constant" in sch_str:

scheduler = ConstantScheduler(**kwargs)

elif "exponential" in sch_str:

scheduler = ExponentialScheduler(**kwargs)

elif "noam" in sch_str:

scheduler = NoamScheduler(**kwargs)

elif "king" in sch_str:

scheduler = KingScheduler(**kwargs)

else:

raise NotImplementedError("{}".format(sch_str))

return scheduler

def init_from_dict(self):

"""Initialize scheduler from the param dictionary"""

S = self.param

sc = S["hyperparameters"] if "hyperparameters" in S else None

if sc is None:

raise ValueError("Must have `hyperparameters` key: {}".format(S))

if sc and sc["id"] == "ConstantScheduler":

scheduler = ConstantScheduler()

elif sc and sc["id"] == "ExponentialScheduler":

scheduler = ExponentialScheduler()

elif sc and sc["id"] == "NoamScheduler":

scheduler = NoamScheduler()

elif sc:

raise NotImplementedError("{}".format(sc["id"]))

return scheduler

class OptimizerInitializer(object):

def __init__(self, param=None):

A class for initializing optimizers. Valid `param` values are:

(a) __str__ representations of `OptimizerBase` instances

(b) `OptimizerBase` instances

`LayerBase` instances)

If `param` is `None`, return the SGD optimizer with default parameters.

self.param = param

def __call__(self):

"""Initialize the optimizer"""

param = self.param

if param is None:

opt = SGD()

elif isinstance(param, OptimizerBase):

opt = param

elif isinstance(param, str):

opt = self.init_from_str()

elif isinstance(param, dict):

opt = self.init_from_dict()

return opt

def init_from_str(self):

"""Initialize optimizer from the `param` string"""

r = r"([a-zA-Z]*)=([^,)]*)"

opt_str = self.param.lower()

kwargs = {i: _eval(j) for i, j in re.findall(r, opt_str)}

if "sgd" in opt_str:

optimizer = SGD(**kwargs)

elif "adagrad" in opt_str:

optimizer = AdaGrad(**kwargs)

elif "rmsprop" in opt_str:

optimizer = RMSProp(**kwargs)

elif "adam" in opt_str:

optimizer = Adam(**kwargs)

else:

raise NotImplementedError("{}".format(opt_str))

return optimizer

def init_from_dict(self):

"""Initialize optimizer from the `param` dictonary"""

D = self.param

cc = D["cache"] if "cache" in D else None

op = D["hyperparameters"] if "hyperparameters" in D else None

if op is None:

raise ValueError("`param` dictionary has no `hyperparemeters` key")

if op and op["id"] == "SGD":

optimizer = SGD()

elif op and op["id"] == "RMSProp":

optimizer = RMSProp()

elif op and op["id"] == "AdaGrad":

optimizer = AdaGrad()

elif op and op["id"] == "Adam":

optimizer = Adam()

elif op:

raise NotImplementedError("{}".format(op["id"]))

optimizer.set_params(op, cc)

return optimizer

class WeightInitializer(object):

def __init__(self, act_fn_str, mode="glorot_uniform"):

A factory for weight initializers.

Parameters

act_fn_str : str

The string representation for the layer activation function

mode : str (default: 'glorot_uniform')

The weight initialization strategy. Valid entries are {"he_normal",

"he_uniform", "glorot_normal", glorot_uniform", "std_normal",

if mode not in [

raise ValueError("Unrecognize initialization mode: {}".format(mode))

self.mode = mode

self.act_fn = act_fn_str

if mode == "glorot_uniform":

self._fn = glorot_uniform

elif mode == "glorot_normal":

self._fn = glorot_normal

elif mode == "he_uniform":

self._fn = he_uniform

elif mode == "he_normal":

self._fn = he_normal

elif mode == "std_normal":

self._fn = np.random.randn

elif mode == "trunc_normal":

self._fn = partial(truncated_normal, mean=0, std=1)

def __call__(self, weight_shape):

"""Initialize weights according to the specified strategy"""

if "glorot" in self.mode:

gain = self._calc_glorot_gain()

W = self._fn(weight_shape, gain)

elif self.mode == "std_normal":

W = self._fn(*weight_shape)

else:

W = self._fn(weight_shape)

return W

def _calc_glorot_gain(self):

Values from:

gain = 1.0

act_str = self.act_fn.lower()

if act_str == "tanh":

gain = 5.0 / 3.0

elif act_str == "relu":

gain = np.sqrt(2)

elif "leaky relu" in act_str:

r = r"leaky relu\(alpha=(.*)\)"

gain = np.sqrt(2 / 1 + float(alpha) ** 2)

return gain


import numpy as np

from .dt import DecisionTree

def bootstrap_sample(X, Y):

N, M = X.shape

idxs = np.random.choice(N, N, replace=True)

return X[idxs], Y[idxs]

class RandomForest:

def __init__(

self, n_trees, max_depth, n_feats, classifier=True, criterion="entropy"

An ensemble (forest) of decision trees where each split is calculated

using a random subset of the features in the input.

Parameters

n_trees : int

The number of individual decision trees to use within the ensemble.

max_depth: int or None

The depth at which to stop growing each decision tree. If None,

grow each tree until the leaf nodes are pure.

n_feats : int

The number of features to sample on each split.

classifier : bool

Whether `Y` contains class labels or real-valued targets. Default

is True.

criterion : {'entropy', 'gini', 'mse'}

The error criterion to use when calculating splits for each weak

learner. When ``classifier = False``, valid entries are {'mse'}.

When ``classifier = True``, valid entries are {'entropy', 'gini'}.

Default is 'entropy'.

self.trees = []

self.n_trees = n_trees

self.n_feats = n_feats

self.max_depth = max_depth

self.criterion = criterion

self.classifier = classifier

def fit(self, X, Y):

Create `n_trees`-worth of bootstrapped samples from the training data

and use each to fit a separate decision tree.

self.trees = []

for _ in range(self.n_trees):

X_samp, Y_samp = bootstrap_sample(X, Y)

tree = DecisionTree(

criterion=self.criterion,

classifier=self.classifier,

tree.fit(X_samp, Y_samp)

self.trees.append(tree)

def predict(self, X):

Predict the target value for each entry in `X`.

Parameters

X : :py:class:`ndarray <numpy.ndarray>` of shape `(N, M)`

The training data of `N` examples, each with `M` features.

Returns

y_pred : :py:class:`ndarray <numpy.ndarray>` of shape `(N,)`

Model predictions for each entry in `X`.

tree_preds = np.array([[t._traverse(x, t.root) for x in X] for t in self.trees])

return self._vote(tree_preds)

def _vote(self, predictions):

Return the aggregated prediction across all trees in the RF for each problem.

Parameters

predictions : :py:class:`ndarray <numpy.ndarray>` of shape `(n_trees, N)`

The array of predictions from each decision tree in the RF for each

of the `N` problems in `X`.

Returns

If classifier is True, the class label predicted by the majority of

the decision trees for each problem in `X`. If classifier is False,

the average prediction across decision trees on each problem.

if self.classifier:

out = [np.bincount(x).argmax() for x in predictions.T]

else:

out = [np.mean(x) for x in predictions.T]

return np.array(out)


import numpy as np

class Node:

def __init__(self, left, right, rule):

self.left = left

self.right = right

self.feature = rule[0]

self.threshold = rule[1]

class Leaf:

def __init__(self, value):

`value` is an array of class probabilities if classifier is True, else

the mean of the region

self.value = value

class DecisionTree:

def __init__(

self,

classifier=True,

criterion="entropy",

seed=None,

A decision tree model for regression and classification problems.

Parameters

classifier : bool

Whether to treat target values as categorical (classifier =

True) or continuous (classifier = False). Default is True.

max_depth: int or None

The depth at which to stop growing the tree. If None, grow the tree

until all leaves are pure. Default is None.

n_feats : int

Specifies the number of features to sample on each split. If None,

use all features on each split. Default is None.

criterion : {'mse', 'entropy', 'gini'}

The error criterion to use when calculating splits. When

`classifier` is False, valid entries are {'mse'}. When `classifier`

is True, valid entries are {'entropy', 'gini'}. Default is

'entropy'.

seed : int or None

Seed for the random number generator. Default is None.

if seed:

np.random.seed(seed)

self.depth = 0

self.root = None

self.n_feats = n_feats

self.criterion = criterion

self.classifier = classifier

self.max_depth = max_depth if max_depth else np.inf

if not classifier and criterion in ["gini", "entropy"]:

raise ValueError(

"{} is a valid criterion only when classifier = True.".format(criterion)

if classifier and criterion == "mse":

raise ValueError("`mse` is a valid criterion only when classifier = False.")

def fit(self, X, Y):

Fit a binary decision tree to a dataset.

Parameters

X : :py:class:`ndarray <numpy.ndarray>` of shape `(N, M)`

The training data of `N` examples, each with `M` features

Y : :py:class:`ndarray <numpy.ndarray>` of shape `(N,)`

An array of integer class labels for each example in `X` if

self.classifier = True, otherwise the set of target values for

each example in `X`.

self.n_classes = max(Y) + 1 if self.classifier else None

self.n_feats = X.shape[1] if not self.n_feats else min(self.n_feats, X.shape[1])

self.root = self._grow(X, Y)

def predict(self, X):

Use the trained decision tree to classify or predict the examples in `X`.

Parameters

Returns

preds : :py:class:`ndarray <numpy.ndarray>` of shape `(N,)`

The integer class labels predicted for each example in `X` if

self.classifier = True, otherwise the predicted target values.

return np.array([self._traverse(x, self.root) for x in X])

def predict_class_probs(self, X):

Use the trained decision tree to return the class probabilities for the

examples in `X`.

Parameters

Returns

preds : :py:class:`ndarray <numpy.ndarray>` of shape `(N, n_classes)`

The class probabilities predicted for each example in `X`.

assert self.classifier, "`predict_class_probs` undefined for classifier = False"

return np.array([self._traverse(x, self.root, prob=True) for x in X])

def _grow(self, X, Y, cur_depth=0):

# if all labels are the same, return a leaf

if len(set(Y)) == 1:

if self.classifier:

prob = np.zeros(self.n_classes)

prob[Y[0]] = 1.0

return Leaf(prob) if self.classifier else Leaf(Y[0])

# if we have reached max_depth, return a leaf

if cur_depth >= self.max_depth:

v = np.mean(Y, axis=0)

if self.classifier:

v = np.bincount(Y, minlength=self.n_classes) / len(Y)

return Leaf(v)

cur_depth += 1

self.depth = max(self.depth, cur_depth)

N, M = X.shape

feat_idxs = np.random.choice(M, self.n_feats, replace=False)

# greedily select the best split according to `criterion`

feat, thresh = self._segment(X, Y, feat_idxs)

l = np.argwhere(X[:, feat] <= thresh).flatten()

r = np.argwhere(X[:, feat] > thresh).flatten()

# grow the children that result from the split

left = self._grow(X[l, :], Y[l], cur_depth)

right = self._grow(X[r, :], Y[r], cur_depth)

return Node(left, right, (feat, thresh))

def _segment(self, X, Y, feat_idxs):

Find the optimal split rule (feature index and split threshold) for the

data according to `self.criterion`.

best_gain = -np.inf

split_idx, split_thresh = None, None

for i in feat_idxs:

vals = X[:, i]

levels = np.unique(vals)

thresholds = (levels[:-1] + levels[1:]) / 2 if len(levels) > 1 else levels

gains = np.array([self._impurity_gain(Y, t, vals) for t in thresholds])

if gains.max() > best_gain:

split_idx = i

best_gain = gains.max()

split_thresh = thresholds[gains.argmax()]

return split_idx, split_thresh

def _impurity_gain(self, Y, split_thresh, feat_values):

Compute the impurity gain associated with a given split.

IG(split) = loss(parent) - weighted_avg[loss(left_child), loss(right_child)]

if self.criterion == "entropy":

loss = entropy

elif self.criterion == "gini":

loss = gini

elif self.criterion == "mse":

loss = mse

parent_loss = loss(Y)

# generate split

left = np.argwhere(feat_values <= split_thresh).flatten()

right = np.argwhere(feat_values > split_thresh).flatten()

if len(left) == 0 or len(right) == 0:

return 0

# compute the weighted avg. of the loss for the children

n = len(Y)

n_l, n_r = len(left), len(right)

e_l, e_r = loss(Y[left]), loss(Y[right])

child_loss = (n_l / n) * e_l + (n_r / n) * e_r

# impurity gain is difference in loss before vs. after split

ig = parent_loss - child_loss

return ig

def _traverse(self, X, node, prob=False):

if isinstance(node, Leaf):

if self.classifier:

return node.value if prob else node.value.argmax()

return node.value

if X[node.feature] <= node.threshold:

return self._traverse(X, node.left, prob)

return self._traverse(X, node.right, prob)

def mse(y):

Mean squared error for decision tree (ie., mean) predictions

return np.mean((y - np.mean(y)) ** 2)

def entropy(y):

Entropy of a label sequence

hist = np.bincount(y)

ps = hist / np.sum(hist)

return -np.sum([p * np.log2(p) for p in ps if p > 0])

def gini(y):

Gini impurity (local entropy) of a label sequence

hist = np.bincount(y)

N = np.sum(hist)

return 1 - sum([(i / N) ** 2 for i in hist])


from . import losses
from .dt import *
from .rf import *
from .gbdt import *


import numpy as np

#                           Base Estimators                           #

class ClassProbEstimator:

def fit(self, X, y):

self.class_prob = y.sum() / len(y)

def predict(self, X):

pred = np.empty(X.shape[0], dtype=np.float64)

return pred

class MeanBaseEstimator:

def fit(self, X, y):

self.avg = np.mean(y)

def predict(self, X):

pred.fill(self.avg)

return pred

#                           Loss Functions                            #

class MSELoss:

def __call__(self, y, y_pred):

return np.mean((y - y_pred) ** 2)

def base_estimator(self):

return MeanBaseEstimator()

def grad(self, y, y_pred):

return -2 / len(y) * (y - y_pred)

def line_search(self, y, y_pred, h_pred):

# TODO: revise this

Lp = np.sum((y - y_pred) * h_pred)

Lpp = np.sum(h_pred * h_pred)

# if we perfectly fit the residuals, use max step size

return 1 if np.sum(Lpp) == 0 else Lp / Lpp

class CrossEntropyLoss:

def __call__(self, y, y_pred):

eps = np.finfo(float).eps

return -np.sum(y * np.log(y_pred + eps))

def base_estimator(self):

return ClassProbEstimator()

def grad(self, y, y_pred):

eps = np.finfo(float).eps

return -y * 1 / (y_pred + eps)

raise NotImplementedError


import numpy as np

from .dt import DecisionTree

from .losses import MSELoss, CrossEntropyLoss

def to_one_hot(labels, n_classes=None):

if labels.ndim > 1:

raise ValueError("labels must have dimension 1, but got {}".format(labels.ndim))

N = labels.size

n_cols = np.max(labels) + 1 if n_classes is None else n_classes

one_hot = np.zeros((N, n_cols))

one_hot[np.arange(N), labels] = 1.0

return one_hot

class GradientBoostedDecisionTree:

def __init__(

self,

classifier=True,

loss="crossentropy",

A gradient boosted ensemble of decision trees.

Notes

Gradient boosted machines (GBMs) fit an ensemble of `m` weak learners such that:

.. math::

f_m(X) = b(X) + \eta w_1 g_1 + \ldots + \eta w_m g_m

where `b` is a fixed initial estimate for the targets, :math:`\eta` is

a learning rate parameter, and :math:`w_{\cdot}` and :math:`g_{\cdot}`

denote the weights and learner predictions for subsequent fits.

We fit each `w` and `g` iteratively using a greedy strategy so that at each

iteration `i`,

.. math::

w_i, g_i = \\arg \min_{w_i, g_i} L(Y, f_{i-1}(X) + w_i g_i)

On each iteration we fit a new weak learner to predict the negative

gradient of the loss with respect to the previous prediction, :math:`f_{i-1}(X)`.

We then use the element-wise product of the predictions of this weak

learner, :math:`g_i`, with a weight, :math:`w_i`, to compute the amount to

adjust the predictions of our model at the previous iteration, :math:`f_{i-1}(X)`:

.. math::

f_i(X) := f_{i-1}(X) + w_i g_i

Parameters

n_iter : int

The number of iterations / weak estimators to use when fitting each

dimension / class of `Y`.

max_depth : int

The maximum depth of each decision tree weak estimator. Default is

None.

classifier : bool

Whether `Y` contains class labels or real-valued targets. Default

is True.

learning_rate : float

Value in [0, 1] controlling the amount each weak estimator

contributes to the overall model prediction. Sometimes known as the

`shrinkage parameter` in the GBM literature. Default is 1.

loss : {'crossentropy', 'mse'}

The loss to optimize for the GBM. Default is 'crossentropy'.

step_size : {"constant", "adaptive"}

How to choose the weight for each weak learner. If "constant", use

a fixed weight of 1 for each learner. If "adaptive", use a step

size computed via line-search on the current iteration's loss.

Default is 'constant'.

self.loss = loss

self.weights = None

self.learners = None

self.out_dims = None

self.n_iter = n_iter

self.base_estimator = None

self.max_depth = max_depth

self.step_size = step_size

self.classifier = classifier

self.learning_rate = learning_rate

def fit(self, X, Y):

Fit the gradient boosted decision trees on a dataset.

Parameters

X : :py:class:`ndarray <numpy.ndarray>` of shape (N, M)

The training data of `N` examples, each with `M` features

Y : :py:class:`ndarray <numpy.ndarray>` of shape (N,)

An array of integer class labels for each example in `X` if

``self.classifier = True``, otherwise the set of target values for

each example in `X`.

if self.loss == "mse":

loss = MSELoss()

elif self.loss == "crossentropy":

loss = CrossEntropyLoss()

# convert Y to one_hot if not already

if self.classifier:

Y = to_one_hot(Y.flatten())

else:

Y = Y.reshape(-1, 1) if len(Y.shape) == 1 else Y

N, M = X.shape

self.out_dims = Y.shape[1]

self.learners = np.empty((self.n_iter, self.out_dims), dtype=object)

self.weights = np.ones((self.n_iter, self.out_dims))

self.weights[1:, :] *= self.learning_rate

# fit the base estimator

Y_pred = np.zeros((N, self.out_dims))

for k in range(self.out_dims):

t = loss.base_estimator()

t.fit(X, Y[:, k])

Y_pred[:, k] += t.predict(X)

self.learners[0, k] = t

# incrementally fit each learner on the negative gradient of the loss

# wrt the previous fit (pseudo-residuals)

for i in range(1, self.n_iter):

for k in range(self.out_dims):

y, y_pred = Y[:, k], Y_pred[:, k]

neg_grad = -1 * loss.grad(y, y_pred)

# use MSE as the surrogate loss when fitting to negative gradients

t = DecisionTree(

classifier=False, max_depth=self.max_depth, criterion="mse"

# fit current learner to negative gradients

t.fit(X, neg_grad)

self.learners[i, k] = t

# compute step size and weight for the current learner

step = 1.0

h_pred = t.predict(X)

if self.step_size == "adaptive":

step = loss.line_search(y, y_pred, h_pred)

# update weights and our overall prediction for Y

self.weights[i, k] *= step

Y_pred[:, k] += self.weights[i, k] * h_pred

def predict(self, X):

Use the trained model to classify or predict the examples in `X`.

Parameters

X : :py:class:`ndarray <numpy.ndarray>` of shape `(N, M)`

Returns

preds : :py:class:`ndarray <numpy.ndarray>` of shape `(N,)`

The integer class labels predicted for each example in `X` if

``self.classifier = True``, otherwise the predicted target values.

Y_pred = np.zeros((X.shape[0], self.out_dims))

for i in range(self.n_iter):

for k in range(self.out_dims):

Y_pred[:, k] += self.weights[i, k] * self.learners[i, k].predict(X)

if self.classifier:

Y_pred = Y_pred.argmax(axis=1)

return Y_pred


from .bandits import *
from . import policies
from . import trainer


"""A module containing different variations on multi-armed bandit environments."""

from abc import ABC, abstractmethod

import numpy as np

from numpy_ml.utils.testing import random_one_hot_matrix, is_number

class Bandit(ABC):

def __init__(self, rewards, reward_probs, context=None):

assert len(rewards) == len(reward_probs)

self.step = 0

self.n_arms = len(rewards)

def __repr__(self):

"""A string representation for the bandit"""

HP = self.hyperparameters

params = ", ".join(["{}={}".format(k, v) for (k, v) in HP.items() if k != "id"])

return "{}({})".format(HP["id"], params)

@property

def hyperparameters(self):

"""A dictionary of the bandit hyperparameters"""

return {}

@abstractmethod

def oracle_payoff(self, context=None):

Return the expected reward for an optimal agent.

Parameters

context : :py:class:`ndarray <numpy.ndarray>` of shape `(D, K)` or None

The current context matrix for each of the bandit arms, if

applicable. Default is None.

Returns

optimal_rwd : float

The expected reward under an optimal policy.

pass

def pull(self, arm_id, context=None):

"Pull" (i.e., sample from) a given arm's payoff distribution.

Parameters

arm_id : int

The integer ID of the arm to sample from

context : :py:class:`ndarray <numpy.ndarray>` of shape `(D,)` or None

The context vector for the current timestep if this is a contextual

bandit. Otherwise, this argument is unused and defaults to None.

Returns

reward : float

The reward sampled from the given arm's payoff distribution

assert arm_id < self.n_arms

self.step += 1

return self._pull(arm_id, context)

def reset(self):

"""Reset the bandit step and action counters to zero."""

self.step = 0

@abstractmethod

def _pull(self, arm_id):

pass

class MultinomialBandit(Bandit):

def __init__(self, payoffs, payoff_probs):

A multi-armed bandit where each arm is associated with a different

multinomial payoff distribution.

Parameters

payoffs : ragged list of length `K`

The payoff values for each of the `n` bandits. ``payoffs[k][i]``

holds the `i` th payoff value for arm `k`.

payoff_probs : ragged list of length `K`

A list of the probabilities associated with each of the payoff

values in ``payoffs``. ``payoff_probs[k][i]`` holds the probability

of payoff index `i` for arm `k`.

super().__init__(payoffs, payoff_probs)

for r, rp in zip(payoffs, payoff_probs):

assert len(r) == len(rp)

np.testing.assert_almost_equal(sum(rp), 1.0)

payoffs = np.array([np.array(x) for x in payoffs])

payoff_probs = np.array([np.array(x) for x in payoff_probs])

self.payoffs = payoffs

self.payoff_probs = payoff_probs

self.arm_evs = np.array([sum(p * v) for p, v in zip(payoff_probs, payoffs)])

self.best_ev = np.max(self.arm_evs)

self.best_arm = np.argmax(self.arm_evs)

@property

def hyperparameters(self):

"""A dictionary of the bandit hyperparameters"""

return {

"id": "MultinomialBandit",

"payoffs": self.payoffs,

"payoff_probs": self.payoff_probs,

def oracle_payoff(self, context=None):

Return the expected reward for an optimal agent.

Parameters

Unused. Default is None.

Returns

optimal_rwd : float

The expected reward under an optimal policy.

optimal_arm : float

The arm ID with the largest expected reward.

return self.best_ev, self.best_arm

def _pull(self, arm_id, context):

payoffs = self.payoffs[arm_id]

probs = self.payoff_probs[arm_id]

return np.random.choice(payoffs, p=probs)

class BernoulliBandit(Bandit):

def __init__(self, payoff_probs):

A multi-armed bandit where each arm is associated with an independent

Bernoulli payoff distribution.

Parameters

payoff_probs : list of length `K`

A list of the payoff probability for each arm. ``payoff_probs[k]``

holds the probability of payoff for arm `k`.

payoffs = [1] * len(payoff_probs)

super().__init__(payoffs, payoff_probs)

for p in payoff_probs:

assert p >= 0 and p <= 1

self.payoffs = np.array(payoffs)

self.payoff_probs = np.array(payoff_probs)

self.arm_evs = self.payoff_probs

self.best_ev = np.max(self.arm_evs)

self.best_arm = np.argmax(self.arm_evs)

@property

def hyperparameters(self):

return {

"id": "BernoulliBandit",

"payoff_probs": self.payoff_probs,

def oracle_payoff(self, context=None):

Parameters

Unused. Default is None.

Returns

optimal_rwd : float

optimal_arm : float

return self.best_ev, self.best_arm

def _pull(self, arm_id, context):

return int(np.random.rand() <= self.payoff_probs[arm_id])

class GaussianBandit(Bandit):

def __init__(self, payoff_dists, payoff_probs):

A multi-armed bandit that is similar to

:class:`BernoulliBandit`, but instead of each arm having

a fixed payout of 1, the payoff values are sampled from independent

Gaussian RVs.

Parameters

payoff_dists : list of 2-tuples of length `K`

The parameters the distributions over payoff values for each of the

`n` arms. Specifically, ``payoffs[k]`` is a tuple of (mean, variance)

for the Gaussian distribution over payoffs associated with arm `k`.

payoff_probs : list of length `n`

values in ``payoffs``. ``payoff_probs[k]`` holds the probability of

payoff for arm `k`.

super().__init__(payoff_dists, payoff_probs)

for (mean, var), rp in zip(payoff_dists, payoff_probs):

assert var > 0

assert np.testing.assert_almost_equal(sum(rp), 1.0)

self.payoff_dists = payoff_dists

self.payoff_probs = payoff_probs

self.arm_evs = np.array([mu for (mu, var) in payoff_dists])

self.best_ev = np.max(self.arm_evs)

self.best_arm = np.argmax(self.arm_evs)

@property

def hyperparameters(self):

"""A dictionary of the bandit hyperparameters"""

return {

"id": "GaussianBandit",

"payoff_dists": self.payoff_dists,

"payoff_probs": self.payoff_probs,

def _pull(self, arm_id, context):

mean, var = self.payoff_dists[arm_id]

reward = 0

if np.random.rand() < self.payoff_probs[arm_id]:

reward = np.random.normal(mean, var)

return reward

def oracle_payoff(self, context=None):

Return the expected reward for an optimal agent.

Parameters

Unused. Default is None.

Returns

optimal_rwd : float

The expected reward under an optimal policy.

optimal_arm : float

The arm ID with the largest expected reward.

return self.best_ev, self.best_arm

class ShortestPathBandit(Bandit):

def __init__(self, G, start_vertex, end_vertex):

A weighted graph shortest path problem formulated as a multi-armed

bandit.

Notes

Each arm corresponds to a valid path through the graph from start to

end vertex. The agent's goal is to find the path that minimizes the

expected sum of the weights on the edges it traverses.

Parameters

G : :class:`Graph <numpy_ml.utils.graphs.Graph>` instance

A weighted graph object. Weights can be fixed or probabilistic.

start_vertex : int

The index of the path's start vertex in the graph

end_vertex : int

The index of the path's end vertex in the graph

self.G = G

self.end_vertex = end_vertex

self.adj_dict = G.to_adj_dict()

self.start_vertex = start_vertex

self.paths = G.all_paths(start_vertex, end_vertex)

self.arm_evs = self._calc_arm_evs()

self.best_ev = np.max(self.arm_evs)

self.best_arm = np.argmax(self.arm_evs)

placeholder = [None] * len(self.paths)

super().__init__(placeholder, placeholder)

@property

def hyperparameters(self):

"""A dictionary of the bandit hyperparameters"""

return {

"id": "ShortestPathBandit",

"G": self.G,

"end_vertex": self.end_vertex,

"start_vertex": self.start_vertex,

def oracle_payoff(self, context=None):

Parameters

Unused. Default is None.

Returns

optimal_rwd : float

optimal_arm : float

return self.best_ev, self.best_arm

def _calc_arm_evs(self):

I2V = self.G.get_vertex

evs = np.zeros(len(self.paths))

for p_ix, path in enumerate(self.paths):

for ix, v_i in enumerate(path[:-1]):

e = [e for e in self.adj_dict[v_i] if e.to == I2V(path[ix + 1])][0]

evs[p_ix] -= e.weight

return evs

def _pull(self, arm_id, context):

reward = 0

I2V = self.G.get_vertex

path = self.paths[arm_id]

for ix, v_i in enumerate(path[:-1]):

reward -= e.weight

return reward

class ContextualBernoulliBandit(Bandit):

def __init__(self, context_probs):

A contextual version of :class:`BernoulliBandit` where each binary

context feature is associated with an independent Bernoulli payoff

distribution.

Parameters

context_probs : :py:class:`ndarray <numpy.ndarray>` of shape `(D, K)`

A matrix of the payoff probabilities associated with each of the

`D` context features, for each of the `K` arms. Index `(i, j)`

contains the probability of payoff for arm `j` under context `i`.

D, K = context_probs.shape

# use a dummy placeholder variable to initialize the Bandit superclass

placeholder = [None] * K

self.context_probs = context_probs

self.arm_evs = self.context_probs

self.best_evs = self.arm_evs.max(axis=1)

self.best_arms = self.arm_evs.argmax(axis=1)

@property

def hyperparameters(self):

return {

"id": "ContextualBernoulliBandit",

"context_probs": self.context_probs,

def get_context(self):

Sample a random one-hot context vector. This vector will be the same

for all arms.

Returns

A random `D`-dimensional one-hot context vector repeated for each

of the `K` bandit arms.

D, K = self.context_probs.shape

context = np.zeros((D, K))

context[np.random.choice(D), :] = 1

return random_one_hot_matrix(1, D).ravel()

def oracle_payoff(self, context):

Return the expected reward for an optimal agent.

Parameters

The current context matrix for each of the bandit arms.

Returns

optimal_rwd : float

The expected reward under an optimal policy.

optimal_arm : float

The arm ID with the largest expected reward.

context_id = context[:, 0].argmax()

return self.best_evs[context_id], self.best_arms[context_id]

def _pull(self, arm_id, context):

D, K = self.context_probs.shape

arm_probs = context[:, arm_id] @ self.context_probs

arm_rwds = (np.random.rand(K) <= arm_probs).astype(int)

return arm_rwds[arm_id]

class ContextualLinearBandit(Bandit):

def __init__(self, K, D, payoff_variance=1):

r"""

A contextual linear multi-armed bandit.

Notes

In a contextual linear bandit the expected payoff of an arm :math:`a

\in \mathcal{A}` at time `t` is a linear combination of its context

vector :math:`\mathbf{x}_{t,a}` with a coefficient vector

.. math::

\mathbb{E}[r_{t, a} \mid \mathbf{x}_{t, a}] = \mathbf{x}_{t,a}^\top \theta_a

In this implementation, the arm coefficient vectors :math:`\theta` are

initialized independently from a uniform distribution on the interval

[-1, 1], and the specific reward at timestep `t` is normally

distributed:

.. math::

r_{t, a} \mid \mathbf{x}_{t, a} \sim

\mathcal{N}(\mathbf{x}_{t,a}^\top \theta_a, \sigma_a^2)

Parameters

K : int

The number of bandit arms

D : int

The dimensionality of the context vectors

payoff_variance : float or :py:class:`ndarray <numpy.ndarray>` of shape `(K,)`

The variance of the random noise in the arm payoffs. If a float,

the variance is assumed to be equal for each arm. Default is 1.

if is_number(payoff_variance):

payoff_variance = [payoff_variance] * K

assert len(payoff_variance) == K

assert all(v > 0 for v in payoff_variance)

self.K = K

self.D = D

self.payoff_variance = payoff_variance

placeholder = [None] * K

super().__init__(placeholder, placeholder)

# initialize the theta matrix

self.thetas = np.random.uniform(-1, 1, size=(D, K))

self.thetas /= np.linalg.norm(self.thetas, 2)

@property

def hyperparameters(self):

"""A dictionary of the bandit hyperparameters"""

return {

"id": "ContextualLinearBandit",

"K": self.K,

"D": self.D,

"payoff_variance": self.payoff_variance,

@property

def parameters(self):

"""A dictionary of the current bandit parameters"""

return {"thetas": self.thetas}

def get_context(self):

Sample the context vectors for each arm from a multivariate standard

normal distribution.

Returns

A `D`-dimensional context vector sampled from a standard normal

distribution for each of the `K` bandit arms.

return np.random.normal(size=(self.D, self.K))

def oracle_payoff(self, context):

Return the expected reward for an optimal agent.

Parameters

applicable. Default is None.

Returns

optimal_rwd : float

The expected reward under an optimal policy.

optimal_arm : float

The arm ID with the largest expected reward.

best_arm = np.argmax(self.arm_evs)

return self.arm_evs[best_arm], best_arm

def _pull(self, arm_id, context):

K, thetas = self.K, self.thetas

self._noise = np.random.normal(scale=self.payoff_variance, size=self.K)

self.arm_evs = np.array([context[:, k] @ thetas[:, k] for k in range(K)])

return (self.arm_evs + self._noise)[arm_id]


"""A trainer/runner object for executing and comparing MAB policies."""

import warnings

import os.path as op

from collections import defaultdict

import numpy as np

from numpy_ml.utils.testing import DependencyWarning

try:

import matplotlib.pyplot as plt

_PLOTTING = True

except ImportError:

fstr = "Cannot import matplotlib. Plotting functionality disabled."

warnings.warn(fstr, DependencyWarning)

_PLOTTING = False

def get_scriptdir():

"""Return the directory containing the `trainer.py` script"""

return op.dirname(op.realpath(__file__))

def mse(bandit, policy):

Computes the mean squared error between a policy's estimates of the

expected arm payouts and the true expected payouts.

if not hasattr(policy, "ev_estimates") or len(policy.ev_estimates) == 0:

return np.nan

se = []

evs = bandit.arm_evs

ests = sorted(policy.ev_estimates.items(), key=lambda x: x[0])

for ix, (est, ev) in enumerate(zip(ests, evs)):

se.append((est[1] - ev) ** 2)

return np.mean(se)

def smooth(prev, cur, weight):

r"""

Compute a simple weighted average of the previous and current value.

Notes

The smoothed value at timestep `t`, :math:`\tilde{X}_t` is calculated as

.. math::

\tilde{X}_t = \epsilon \tilde{X}_{t-1} + (1 - \epsilon) X_t

where :math:`X_t` is the value at timestep `t`, :math:`\tilde{X}_{t-1}` is

the value of the smoothed signal at timestep `t-1`, and :math:`\epsilon` is

the smoothing weight.

Parameters

prev : float or :py:class:`ndarray <numpy.ndarray>` of shape `(N,)`

The value of the smoothed signal at the immediately preceding

timestep.

cur : float or :py:class:`ndarray <numpy.ndarray>` of shape `(N,)`

The value of the signal at the current timestep

weight : float or :py:class:`ndarray <numpy.ndarray>` of shape `(N,)`

The smoothing weight. Values closer to 0 result in less smoothing,

values closer to 1 produce more aggressive smoothing. If weight is an

array, each dimension will be interpreted as a separate smoothing

weight the corresponding dimension in `cur`.

Returns

smoothed : float or :py:class:`ndarray <numpy.ndarray>` of shape `(N,)`

The smoothed signal

return weight * prev + (1 - weight) * cur

class BanditTrainer:

def __init__(self):

An object to facilitate multi-armed bandit training, comparison, and

evaluation.

self.logs = {}

def compare(

self,

policies,

bandit,

plot=True,

seed=None,

Compare the performance of multiple policies on the same bandit

environment, generating a plot for each.

Parameters

policies : list of :class:`BanditPolicyBase <numpy_ml.bandits.policies.BanditPolicyBase>` instances

The multi-armed bandit policies to compare.

bandit : :class:`Bandit <numpy_ml.bandits.bandits.Bandit>` instance

The environment to train the policies on.

n_trials : int

The number of trials per run.

n_duplicates: int

The number of times to evaluate each policy on the bandit

environment. Larger values permit a better estimate of the

variance in payoff / cumulative regret for each policy.

plot : bool

Whether to generate a plot of the policy's average reward and

regret across the episodes. Default is True.

seed : int

The seed for the random number generator. Default is None.

smooth_weight : float in [0, 1]

values closer to 1 produce more aggressive smoothing. Default is

out_dir : str or None

Plots will be saved to this directory if `plot` is True. If

`out_dir` is None, plots will not be saved. Default is None.

"""  # noqa: E501

all_axes = [None] * len(policies)

if plot and _PLOTTING:

fig, all_axes = plt.subplots(len(policies), 2, sharex=True)

fig.set_size_inches(10.5, len(policies) * 5.25)

for policy, axes in zip(policies, all_axes):

if seed:

np.random.seed(seed)

bandit.reset()

policy.reset()

self.train(

policy,

bandit,

axes=axes,

plot=plot,

verbose=False,

# enforce the same y-ranges across plots for straightforward comparison

a1_r, a2_r = zip(*[(a1.get_ylim(), a2.get_ylim()) for (a1, a2) in all_axes])

a1_min = min(a1_r, key=lambda x: x[0])[0]

a1_max = max(a1_r, key=lambda x: x[1])[1]

a2_min = min(a2_r, key=lambda x: x[0])[0]

a2_max = max(a2_r, key=lambda x: x[1])[1]

for (a1, a2) in all_axes:

a1.set_ylim(a1_min, a1_max)

a2.set_ylim(a2_min, a2_max)

if plot and _PLOTTING:

if out_dir is not None:

plt.savefig(op.join(out_dir, "bandit_comparison.png"), dpi=300)

plt.show()

def train(

self,

policy,

bandit,

plot=True,

axes=None,

verbose=True,

Train a MAB policies on a multi-armed bandit problem, logging training

statistics along the way.

Parameters

policy : :class:`BanditPolicyBase <numpy_ml.bandits.policies.BanditPolicyBase>` instance

The multi-armed bandit policy to train.

The environment to run the policy on.

n_trials : int

The number of trials per run.

n_duplicates: int

The number of runs to evaluate

plot : bool

regret across the episodes. Default is True.

axes : list of :py:class:`Axis <matplotlib.axes.Axis>` instances or None

If not None and ``plot = True``, these are the axes that will be

used to plot the cumulative reward and regret, respectively.

Default is None.

verbose : boolean

Whether to print run statistics during training. Default is True.

print_every : int

The number of episodes to run before printing loss values to

stdout. This is ignored if ``verbose`` is false. Default is 100.

smooth_weight : float in [0, 1]

out_dir : str or None

Returns

The policy trained during the last (i.e. most recent) duplicate

run.

"""  # noqa: E501

if not str(policy) in self.logs:

p = str(policy)

D, L = n_duplicates, self.logs

for d in range(D):

if verbose:

print("\nDUPLICATE {}/{}\n".format(d + 1, D))

bandit.reset()

policy.reset()

avg_oracle_reward, cregret = 0, 0

for trial_id in range(n_trials):

rwd, arm, orwd, oarm = self._train_step(bandit, policy)

loss = mse(bandit, policy)

regret = orwd - rwd

avg_oracle_reward += orwd

cregret += regret

L[p]["mse"][trial_id + 1].append(loss)

L[p]["reward"][trial_id + 1].append(rwd)

L[p]["regret"][trial_id + 1].append(regret)

L[p]["cregret"][trial_id + 1].append(cregret)

L[p]["optimal_arm"][trial_id + 1].append(oarm)

L[p]["selected_arm"][trial_id + 1].append(arm)

L[p]["optimal_reward"][trial_id + 1].append(orwd)

if (trial_id + 1) % print_every == 0 and verbose:

fstr = "Trial {}/{}, {}/{}, Regret: {:.4f}"

print(fstr.format(trial_id + 1, n_trials, d + 1, D, regret))

avg_oracle_reward /= n_trials

if verbose:

self._print_run_summary(bandit, policy, regret)

if plot and _PLOTTING:

self._plot_reward(avg_oracle_reward, policy, smooth_weight, axes, out_dir)

return policy

def _train_step(self, bandit, policy):

P, B = policy, bandit

C = B.get_context() if hasattr(B, "get_context") else None

rwd, arm = P.act(B, C)

oracle_rwd, oracle_arm = B.oracle_payoff(C)

return rwd, arm, oracle_rwd, oracle_arm

def init_logs(self, policies):

Initialize the episode logs.

Notes

Training logs are represented as a nested set of dictionaries with the

following structure:

log[model_id][metric][trial_number][duplicate_number]

For example, ``logs['model1']['regret'][3][1]`` holds the regret value

accrued on the 3rd trial of the 2nd duplicate run for model1.

Available fields are 'regret', 'cregret' (cumulative regret), 'reward',

'mse' (mean-squared error between estimated arm EVs and the true EVs),

'optimal_arm', 'selected_arm', and 'optimal_reward'.

if not isinstance(policies, list):

policies = [policies]

self.logs = {

str(p): {

"mse": defaultdict(lambda: []),

"regret": defaultdict(lambda: []),

"reward": defaultdict(lambda: []),

"cregret": defaultdict(lambda: []),

"optimal_arm": defaultdict(lambda: []),

"selected_arm": defaultdict(lambda: []),

"optimal_reward": defaultdict(lambda: []),

for p in policies

def _print_run_summary(self, bandit, policy, regret):

return None

evs, se = bandit.arm_evs, []

fstr = "Arm {}: {:.4f} v. {:.4f}"

print("\n\nEstimated vs. Real EV\n" + "-" * 21)

for ix, (est, ev) in enumerate(zip(ests, evs)):

print(fstr.format(ix + 1, est[1], ev))

se.append((est[1] - ev) ** 2)

fstr = "\nFinal MSE: {:.4f}\nFinal Regret: {:.4f}\n\n"

print(fstr.format(np.mean(se), regret))

def _plot_reward(self, optimal_rwd, policy, smooth_weight, axes=None, out_dir=None):

L = self.logs[str(policy)]

smds = self._smoothed_metrics(policy, optimal_rwd, smooth_weight)

if axes is None:

fig, [ax1, ax2] = plt.subplots(1, 2)

else:

assert len(axes) == 2

ax1, ax2 = axes

e_ids = range(1, len(L["reward"]) + 1)

plot_params = [[ax1, ax2], ["reward", "cregret"], ["b", "r"], [optimal_rwd, 0]]

for (ax, m, c, opt) in zip(*plot_params):

avg, std = "sm_{}_avg sm_{}_std".format(m, m).split()

ax.plot(e_ids, smds[avg], color=c)

ax.axhline(opt, 0, 1, color=c, ls="--")

smds[avg] + smds[std],

smds[avg] - smds[std],

color=c,

m = "Cumulative Regret" if m == "cregret" else m

ax.set_ylabel("Smoothed Avg. {}".format(m.title()))

if axes is None:

ax.set_aspect(np.diff(ax.get_xlim()) / np.diff(ax.get_ylim()))

if axes is not None:

if axes is None:

fig.suptitle(str(policy))

if out_dir is not None:

bid = policy.hyperparameters["id"]

plt.savefig(op.join(out_dir, f"{bid}.png"), dpi=300)

plt.show()

return ax1, ax2

def _smoothed_metrics(self, policy, optimal_rwd, smooth_weight):

L = self.logs[str(policy)]

# pre-allocate smoothed data structure

smds = {}

for m in L.keys():

if m == "selections":

continue

smds["sm_{}_avg".format(m)] = np.zeros(len(L["reward"]))

smds["sm_{}_avg".format(m)][0] = np.mean(L[m][1])

smds["sm_{}_std".format(m)] = np.zeros(len(L["reward"]))

smds["sm_{}_std".format(m)][0] = np.std(L[m][1])

smoothed = {m: L[m][1] for m in L.keys()}

for e_id in range(2, len(L["reward"]) + 1):

for m in L.keys():

if m == "selections":

continue

prev, cur = smoothed[m], L[m][e_id]

smoothed[m] = [smooth(p, c, smooth_weight) for p, c in zip(prev, cur)]

smds["sm_{}_avg".format(m)][e_id - 1] = np.mean(smoothed[m])

smds["sm_{}_std".format(m)][e_id - 1] = np.std(smoothed[m])

return smds