mbrs.modules package

mbrs.modules package#

Submodules#

Module contents#

class mbrs.modules.Kmeans(cfg: Config)[source]#

Bases: object

k-means clustering implemented in PyTorch.

Parameters:

cfg (Kmeans.Config) – Configuration for k-means.

class Config(ncentroids: int = 8, niter: int = 5, kmeanspp: bool = True, seed: int = 0)[source]#

Bases: object

Configuration for k-means.

  • ncentroids (int): Number of centroids.

  • niter (int): Number of k-means iteration

  • kmeanspp (bool): Initialize the centroids using k-means++.

  • seed (bool): Random seed.

kmeanspp: bool = True#
ncentroids: int = 8#
niter: int = 5#
seed: int = 0#
assign(x: Tensor, centroids: Tensor) → Tensor[source]#

Assigns the nearest neighbor centroid ID.

Parameters:

  • x (torch.Tensor) – Assigned vectors of shape (n, dim).

  • centroids (torch.Tensor) – Centroids tensor of shape (ncentroids, dim).

Returns:

Assigned IDs of shape (n,).

Return type:

torch.Tensor

init_kmeanspp(x: Tensor, rng: Generator, ncentroids: int) → Tensor[source]#

Initializes the centroids via k-means++.

Parameters:

  • x (Tensor) – Input vectors of shape (n, dim).

  • rng (Generator) – Random number generator.

  • ncentroids (int) – Number of centroids.

Returns:

Centroid vectors obtained using k-means++.

Return type:

Tensor

train(x: Tensor) → Tuple[Tensor, Tensor][source]#

Trains k-means.

Parameters:

x (torch.Tensor) – Input vectors of shape (n, dim).

Returns:

Centroids tensor of shape (ncentroids, dim). Tensor: Assigend IDs of shape (n,).

Return type:

Tensor

class mbrs.modules.MatrixFactorizationALS(regularization_weight: float = 0.1, rank: int = 8)[source]#

Bases: object

Alternating least squares (ALS) implemented in PyTorch.

Parameters:

  • regularization_weight (float) – Weight of L2 regularization.

  • rank (int) – Rank of the factarized matrices.

compute_loss(matrix: Tensor, x: Tensor, y: Tensor, observed_mask: Tensor | None = None) → float[source]#

Compute the objective loss function.

Parameters:

  • matrix (Tensor) – Target matrix of shape (N, M).

  • x (Tensor) – Left-side low-rank matrix of shape (N, r).

  • y (Tensor) – Right-side low-rank matrix of shape (M, r).

  • observed_mask (Tensor, optional) – Valid indices boolean mask of shape (N, M).

Returns:

Objective loss.

Return type:

float

factorize(matrix: Tensor, observed_mask: Tensor | None = None, niter: int = 30, tolerance: float = 0.0001, seed: int = 0) → Tuple[Tensor, Tensor][source]#

Factorize the given matrix.

The input matrix of shape (N, M) is decomposed into X @ Y.T, where X and Y shape (N, r) and (M, r), respectively.

This implementation does not compute the inverse matrix directly in X = A^-1 @ b. Instead, AX = b is solved.

Parameters:

  • matrix (Tensor) – Input matrix of shape (N, M).

  • observed_mask (Tensor, optional) – Boolean mask of valid indices of shape (N, M).

  • niter (int) – The number of alternating steps performed.

  • tolerance (float) – If the difference between the previous and current loss is smaller this value, ALS is regarded as converged.

  • seed (int) – A seed for the random number generator.

Returns:

Low-rank matrix X of shape (N, r). Tensor: Low-rank matrix Y of shape (M, r).

Return type:

Tensor

Contents