BayesianNeuralFieldMAP

Bases: BayesianNeuralFieldEstimator

Fits models using stochastic ensembles of maximum-a-posteriori estimates.

Implementation of BayesianNeuralFieldEstimator.

Source code in bayesnf/spatiotemporal.py

class BayesianNeuralFieldMAP(BayesianNeuralFieldEstimator):
  """Fits models using stochastic ensembles of maximum-a-posteriori estimates.

  Implementation of
  [BayesianNeuralFieldEstimator](BayesianNeuralFieldEstimator.md).
  """

  _ensemble_dims = 2

  def fit(
      self,
      table,
      seed,
      ensemble_size=16,
      learning_rate=0.005,
      num_epochs=5_000,
      batch_size=None,
      num_splits=1,
      ) -> BayesianNeuralFieldEstimator:
    """Run inference using stochastic MAP ensembles.

    Args:
      table (pandas.DataFrame):
        See documentation of
        [`table`][bayesnf.spatiotemporal.BayesianNeuralFieldEstimator.fit]
        in the base class.

      seed (jax.random.PRNGKey): The jax random key.

      ensemble_size (int): Number of particles in the ensemble. It currently
        an error if `ensemble_size < jax.device_count`, but will be fixed
        in https://github.com/google/bayesnf/issues/28.


      learning_rate (float): Learning rate for SGD.

      num_epochs (int): Number of full epochs through the training data.

      batch_size (None | int): Batch size for SGD. Default is `None`,
        meaning full-batch. Each epoch will perform `len(table) // batch_size`
        SGD updates.

      num_splits (int): Number of splits over the data to run training.
        Defaults to 1, meaning there are no splits.

    Returns:
      Instance of `self`.
    """
    if ensemble_size < jax.device_count():
      raise ValueError('ensemble_size cannot be smaller than device_count. '
                       'https://github.com/google/bayesnf/issues/28.')
    train_data = self.data_handler.get_train(table)
    train_target = self.data_handler.get_target(table)
    if batch_size is None:
      batch_size = train_data.shape[0]
    if self._scale_epochs_by_batch_size:
      num_epochs = num_epochs * (train_data.shape[0] // batch_size)
    model_args = self._model_args((batch_size, train_data.shape[-1]))
    self.params_, self.losses_ = inference.fit_map(
        train_data,
        train_target,
        seed=seed,
        observation_model=self.observation_model,
        model_args=model_args,
        num_particles=ensemble_size,
        learning_rate=learning_rate,
        num_epochs=num_epochs,
        prior_weight=self._prior_weight,
        batch_size=batch_size,
        num_splits=num_splits)
    return self

fit

fit(table, seed, ensemble_size=16, learning_rate=0.005, num_epochs=5000, batch_size=None, num_splits=1)

Run inference using stochastic MAP ensembles.

PARAMETER	DESCRIPTION
table	See documentation of table in the base class. TYPE: DataFrame
seed	The jax random key. TYPE: PRNGKey
ensemble_size	Number of particles in the ensemble. It currently an error if ensemble_size < jax.device_count, but will be fixed in https://github.com/google/bayesnf/issues/28. TYPE: int DEFAULT: 16
learning_rate	Learning rate for SGD. TYPE: float DEFAULT: 0.005
num_epochs	Number of full epochs through the training data. TYPE: int DEFAULT: 5000
batch_size	Batch size for SGD. Default is None, meaning full-batch. Each epoch will perform len(table) // batch_size SGD updates. TYPE: None | int DEFAULT: None
num_splits	Number of splits over the data to run training. Defaults to 1, meaning there are no splits. TYPE: int DEFAULT: 1

RETURNS	DESCRIPTION
BayesianNeuralFieldEstimator	Instance of self.

Source code in bayesnf/spatiotemporal.py

def fit(
    self,
    table,
    seed,
    ensemble_size=16,
    learning_rate=0.005,
    num_epochs=5_000,
    batch_size=None,
    num_splits=1,
    ) -> BayesianNeuralFieldEstimator:
  """Run inference using stochastic MAP ensembles.

  Args:
    table (pandas.DataFrame):
      See documentation of
      [`table`][bayesnf.spatiotemporal.BayesianNeuralFieldEstimator.fit]
      in the base class.

    seed (jax.random.PRNGKey): The jax random key.

    ensemble_size (int): Number of particles in the ensemble. It currently
      an error if `ensemble_size < jax.device_count`, but will be fixed
      in https://github.com/google/bayesnf/issues/28.


    learning_rate (float): Learning rate for SGD.

    num_epochs (int): Number of full epochs through the training data.

    batch_size (None | int): Batch size for SGD. Default is `None`,
      meaning full-batch. Each epoch will perform `len(table) // batch_size`
      SGD updates.

    num_splits (int): Number of splits over the data to run training.
      Defaults to 1, meaning there are no splits.

  Returns:
    Instance of `self`.
  """
  if ensemble_size < jax.device_count():
    raise ValueError('ensemble_size cannot be smaller than device_count. '
                     'https://github.com/google/bayesnf/issues/28.')
  train_data = self.data_handler.get_train(table)
  train_target = self.data_handler.get_target(table)
  if batch_size is None:
    batch_size = train_data.shape[0]
  if self._scale_epochs_by_batch_size:
    num_epochs = num_epochs * (train_data.shape[0] // batch_size)
  model_args = self._model_args((batch_size, train_data.shape[-1]))
  self.params_, self.losses_ = inference.fit_map(
      train_data,
      train_target,
      seed=seed,
      observation_model=self.observation_model,
      model_args=model_args,
      num_particles=ensemble_size,
      learning_rate=learning_rate,
      num_epochs=num_epochs,
      prior_weight=self._prior_weight,
      batch_size=batch_size,
      num_splits=num_splits)
  return self