Sngp #28
Sngp #28
Conversation
fortuna/model/sngp.py
Outdated
| block_cls=self.block_cls, | ||
| num_filters=self.num_filters, | ||
| dtype=self.dtype, | ||
| activation=self.activation, |
There was a problem hiding this comment.
In the MLP model above, there was a [:-1] here. Just to raise attention in case of bug.
There was a problem hiding this comment.
Which MLP model above? This is a ResNet model.
There was a problem hiding this comment.
Here. Isn't the deep feature extractor coming from an MLP?
fortuna/model/sngp.py
Outdated
| pass | ||
|
|
||
|
|
||
| class ResNetSNGP(SNGPMixin, ResNet): |
There was a problem hiding this comment.
Should't ResNets be actually the only model SNGPs are allowed to work with? In such case, why in the notebook example above are we doing this we an MLP?
There was a problem hiding this comment.
On a second note, do you think ResNetSNGP is better than SNGPResNet? Same comment for the WideResNet models below.
There was a problem hiding this comment.
In the notebook we used an MLP with residual connections.
There was a problem hiding this comment.
Maybe SNGPResNet is better, I'll change that
| # Default field value for kwargs, to be used for data class declaration. | ||
| default_kwarg_dict = lambda: dataclasses.field(default_factory=dict) | ||
|
|
||
| SUPPORTED_LIKELIHOOD = ("binary_logistic", "poisson", "gaussian") |
There was a problem hiding this comment.
As mentioned above, can we make this better?
| self.covmat_layer = LaplaceRandomFeatureCovariance( | ||
| hidden_features=self.hidden_features, **self.covmat_kwargs | ||
| ) | ||
| # pylint:enable=invalid-name,not-a-mapping |
There was a problem hiding this comment.
I'd remove these #pylint comments.
There was a problem hiding this comment.
Nice PR - very clean implementation :D As with @gianlucadetommaso, most of my comments are minor ones.
A general comments is to make sure that any type signatures of jnp.ndarray are replaced by jax.Array. You may like to look into JaxTyping where the array's shape can dynamically be specified. This will allow for nice type checking in the future.
| # Fortuna helps you converting data and data loaders into a data loader that Fortuna can digest. | ||
|
|
||
| # %% | ||
| from fortuna.data import DataLoader |
There was a problem hiding this comment.
A general comment that all imports should be at the top of the file to adhere to PEP rules.
There was a problem hiding this comment.
Agreed, but is this true also for the example notebooks?
There was a problem hiding this comment.
This is on me. I'm using imports in each cell of example notebooks to make it easier to understand where objects are called from. I guess it is debatable what's best practice in this case.
fortuna/model/sngp.py
Outdated
|
|
||
| def __call__( | ||
| self, x: Array, train: bool = False, **kwargs | ||
| ) -> Union[jnp.ndarray, Tuple[jnp.ndarray, jnp.ndarray]]: |
| kernel_init: Callable[[PRNGKeyArray, Shape, Type], Array] = default_rbf_kernel_init | ||
| bias_init: Callable[[PRNGKeyArray, Shape, Type], Array] = default_rbf_bias_init | ||
| seed: int = 0 | ||
| dtype: Type = jnp.float32 |
There was a problem hiding this comment.
Will those pose any problems when we compute Cholesky factors?
There was a problem hiding this comment.
This is a great point. I do not know if this will work as expected in half precision for example. How do you suggest to proceed? Should we force the execution of this part in jnp.float32?
- Moved SNGP callback to fortuna.posterior.sngp - Improved SNGP docs - Moved ClassificationModelManagers to own py file - Typos in two_moons_classification_sngp.pct.py
| - When calling the train method for a :class:`~fortuna.prob_model.base.ProbModel` instance, add a list of callbacks containing the ones previously defined when initializing :class:`~fortuna.prob_model.fit_config.base.FitConfig`. | ||
|
|
||
| The following example outlines the usage of :class:`~fortuna.prob_model.callbacks.base.Callback`. | ||
| It assumes that the user already obtained an insatnce of :class:`~fortuna.prob_model.base.ProbModel`: |
There was a problem hiding this comment.
insatnce -> instance
- dryrun tests for sngp - added possibility to do a sanity check of the posterior state in the posterior fit method - for SNGP check that the deep feature extractor has spectral norm
Implemenation of Spectral-normalized Neural Gaussian Process.
Pull request type
Please check the type of change your PR introduces:
What is the current behavior?
What is the new behavior?
New uncertainty quantification method that can be used with ResNet-like models.
A large part of the SNGP implementation (
spectral_norm.pyandrandom_features.py) has been adapted from the Jax (and Tensorflow) implementations in Edward2 and uncertainty-baselines.Other information
prob_modelinstead of being withinprob_model.fit_config