Batchnorm
This module implements batchnormalization on SPD Manifolds.
- anotherspdnet.batchnorm.riemannian_mean_spd(X: torch.Tensor, initial_stepsize: float = 1.0, max_iter: int = 5, tol: float = 1e-06, weights: torch.Tensor | None = None) torch.Tensor[source]
Riemannian mean for SPD matrices.
Inspired by pyRiemann’s implementation.
- Parameters:
X (torch.Tensor of shape (..., N, N)) – Batch of SPD matrices. The mean is computed along the … axes.
initial_stepsize (float) – Initial step size to move on the tangent space.
max_iter (int) – Maximum number of iterations.
tol (float) – Tolerance to stop the algorithm.
weights (torch.Tensor) – Weights to apply to the samples. Must be of same shape as X until ndim - 2. The weights are applied on the … axes. by default, all are put to one.
- Returns:
mean – Riemannian mean.
- Return type:
torch.Tensor of shape (N, N)
- class anotherspdnet.batchnorm.BatchNormSPD(*args: Any, **kwargs: Any)[source]
- __init__(n_features: int, momentum: float = 0.1, tol_mean: float = 1e-06, max_iter_mean: int = 5, initial_stepsize_mean: float = 1.0, device: torch.device = torch.device, dtype: torch.dtype = torch.float64) None[source]
Batch normalization for SPD matrices. From paper: Brooks et al., “Riemannian Batch Normalization for SPD Neural Networks.” NEURIPS 2019.
- Parameters:
n_features (int) – Number of features.
momentum (float) – Momentum for the running covariance mean
tol_mean (float) – Tolerance for the mean computation.
max_iter_mean (int) – Maximum number of iterations for the mean computation.
initial_stepsize_mean (float) – Initial step size for the mean computation.
device (torch.device) – Device to use.