fastersom 0.2.0__tar.gz → 0.3.0__tar.gz

{fastersom-0.2.0 → fastersom-0.3.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: fastersom
-Version: 0.2.0
+Version: 0.3.0
 Summary: Fast batch self-organising maps (SOM) in PyTorch for GPU and CPU
 Keywords: self-organizing map,som,kohonen,clustering,pytorch,gpu,spatial omics
 Author: Darius Schaub
@@ -30,9 +30,9 @@ Description-Content-Type: text/markdown
 
 A small, fast batch self-organising map (SOM) in PyTorch: a reimplementation of Kohonen's batch map with a
 rectangular grid, gaussian neighbourhood and exponentially decaying neighbourhood width. Its defaults reproduce the
-maps of [spatiomic](https://github.com/complextissue/spatiomic)'s `so.dimension.som`. Data stays on the device (CUDA,
-MPS or CPU) and the per-epoch update uses per-node sums, so an epoch costs one BMU search plus a (nodes x nodes) matrix
-product.
+batch-trained maps of [XPySom](https://github.com/Manciukic/xpysom) with a gaussian neighbourhood of initial width
+half the grid. Data stays on the device (CUDA, MPS or CPU) and the per-epoch update uses per-node sums, so an epoch
+costs one BMU search plus a (nodes x nodes) matrix product.
 
 ## Install
 
@@ -72,21 +72,33 @@ Options:
   (`SOM(init=som.weights_).fit(X)`); with `n_epochs=0`, `fit` keeps them unchanged. The sigma schedule restarts at
   `sigma_initial`, which with the default largely retrains the map; to continue training instead, pass a small
   `sigma_initial`, e.g. the previous map's `sigma_final`.
-- `sigma_initial` and `sigma_final` set the neighbourhood width schedule; the defaults reproduce spatiomic.
-- `batch_size` bounds memory (about `batch_size x n_nodes` floats) without changing the algorithm.
+- `sigma_initial` and `sigma_final` set the neighbourhood width schedule; the defaults reproduce XPySom's.
+- `batch_size` bounds memory (about `batch_size x n_nodes` floats) without changing the algorithm. On CUDA and MPS
+  it defaults to the largest batch that still helps speed and fits in `memory_fraction` (default 0.9 on CUDA, 0.5
+  on MPS) of the free GPU memory, so it rarely needs tuning.
 
 Fitted maps pickle without device tensors and load on the configured device if available, else on the default one,
 so a map fitted on a GPU loads on a CPU-only machine. `BatchSOM` is an alias for `SOM`.
 
 ## Benchmark
 
-480,000 x 34 pixels, 30 x 30 nodes, 10 epochs, fit time only (mean of seeds 0-2):
-
-| implementation               | hardware          | fit time | quantization error |
-| ---------------------------- | ----------------- | -------- | ------------------ |
-| fastersom                    | NVIDIA A100       | 0.085 s  | 0.169              |
-| spatiomic (CuPy)             | NVIDIA A100       | 0.46 s   | 0.169              |
-| spatiomic (NumPy)            | 256-core CPU      | 40 s     | 0.169              |
+480,000 x 34 pixels, 30 x 30 nodes, 10 epochs, mean of 3 seeds; peak GPU memory includes the data
+([`benchmarks/memory.py`](https://github.com/dschaub95/fastersom/blob/main/benchmarks/memory.py)).
+
+| implementation                                       | hardware     | batch size        | fit time | peak GPU memory | quantization error |
+| ---------------------------------------------------- | ------------ | ----------------- | -------- | --------------- | ------------------ |
+| fastersom                                            | NVIDIA A100  | default (298,261) | 0.12 s   | 1150 MiB        | 0.168              |
+| fastersom                                            | NVIDIA A100  | 32,768            | 0.12 s   | 204 MiB         | 0.168              |
+| fastersom                                            | NVIDIA A100  | 8,096             | 0.20 s   | 118 MiB         | 0.168              |
+| [XPySom](https://github.com/Manciukic/xpysom) (CuPy) | NVIDIA A100  | 8,096             | 0.46 s   | 313 MiB         | 0.169              |
+| XPySom (CuPy)                                        | NVIDIA A100  | 65,536            | 0.41 s   | 2088 MiB        | 0.169              |
+| [torchsom](https://github.com/michelin/TorchSOM)     | NVIDIA A100  | 1,024             | 21.7 s   | 44 MiB          | 0.175              |
+| fastersom                                            | 256-core CPU | default (8,192)   | 1.8 s    | -               | 0.169              |
+| XPySom (NumPy)                                       | 256-core CPU | 8,096             | 40.6 s   | -               | 0.168              |
+| torchsom                                             | 256-core CPU | 1,024             | 39.2 s   | -               | 0.176              |
+
+At equal GPU memory fastersom is about 4x faster than XPySom, and about 20x faster on the CPU. torchsom updates the map
+per mini-batch, so it needs small batches and is far slower.
 
 ## Algorithm
 

{fastersom-0.2.0 → fastersom-0.3.0}/README.md

@@ -2,9 +2,9 @@
 
 A small, fast batch self-organising map (SOM) in PyTorch: a reimplementation of Kohonen's batch map with a
 rectangular grid, gaussian neighbourhood and exponentially decaying neighbourhood width. Its defaults reproduce the
-maps of [spatiomic](https://github.com/complextissue/spatiomic)'s `so.dimension.som`. Data stays on the device (CUDA,
-MPS or CPU) and the per-epoch update uses per-node sums, so an epoch costs one BMU search plus a (nodes x nodes) matrix
-product.
+batch-trained maps of [XPySom](https://github.com/Manciukic/xpysom) with a gaussian neighbourhood of initial width
+half the grid. Data stays on the device (CUDA, MPS or CPU) and the per-epoch update uses per-node sums, so an epoch
+costs one BMU search plus a (nodes x nodes) matrix product.
 
 ## Install
 
@@ -44,21 +44,33 @@ Options:
   (`SOM(init=som.weights_).fit(X)`); with `n_epochs=0`, `fit` keeps them unchanged. The sigma schedule restarts at
   `sigma_initial`, which with the default largely retrains the map; to continue training instead, pass a small
   `sigma_initial`, e.g. the previous map's `sigma_final`.
-- `sigma_initial` and `sigma_final` set the neighbourhood width schedule; the defaults reproduce spatiomic.
-- `batch_size` bounds memory (about `batch_size x n_nodes` floats) without changing the algorithm.
+- `sigma_initial` and `sigma_final` set the neighbourhood width schedule; the defaults reproduce XPySom's.
+- `batch_size` bounds memory (about `batch_size x n_nodes` floats) without changing the algorithm. On CUDA and MPS
+  it defaults to the largest batch that still helps speed and fits in `memory_fraction` (default 0.9 on CUDA, 0.5
+  on MPS) of the free GPU memory, so it rarely needs tuning.
 
 Fitted maps pickle without device tensors and load on the configured device if available, else on the default one,
 so a map fitted on a GPU loads on a CPU-only machine. `BatchSOM` is an alias for `SOM`.
 
 ## Benchmark
 
-480,000 x 34 pixels, 30 x 30 nodes, 10 epochs, fit time only (mean of seeds 0-2):
-
-| implementation               | hardware          | fit time | quantization error |
-| ---------------------------- | ----------------- | -------- | ------------------ |
-| fastersom                    | NVIDIA A100       | 0.085 s  | 0.169              |
-| spatiomic (CuPy)             | NVIDIA A100       | 0.46 s   | 0.169              |
-| spatiomic (NumPy)            | 256-core CPU      | 40 s     | 0.169              |
+480,000 x 34 pixels, 30 x 30 nodes, 10 epochs, mean of 3 seeds; peak GPU memory includes the data
+([`benchmarks/memory.py`](https://github.com/dschaub95/fastersom/blob/main/benchmarks/memory.py)).
+
+| implementation                                       | hardware     | batch size        | fit time | peak GPU memory | quantization error |
+| ---------------------------------------------------- | ------------ | ----------------- | -------- | --------------- | ------------------ |
+| fastersom                                            | NVIDIA A100  | default (298,261) | 0.12 s   | 1150 MiB        | 0.168              |
+| fastersom                                            | NVIDIA A100  | 32,768            | 0.12 s   | 204 MiB         | 0.168              |
+| fastersom                                            | NVIDIA A100  | 8,096             | 0.20 s   | 118 MiB         | 0.168              |
+| [XPySom](https://github.com/Manciukic/xpysom) (CuPy) | NVIDIA A100  | 8,096             | 0.46 s   | 313 MiB         | 0.169              |
+| XPySom (CuPy)                                        | NVIDIA A100  | 65,536            | 0.41 s   | 2088 MiB        | 0.169              |
+| [torchsom](https://github.com/michelin/TorchSOM)     | NVIDIA A100  | 1,024             | 21.7 s   | 44 MiB          | 0.175              |
+| fastersom                                            | 256-core CPU | default (8,192)   | 1.8 s    | -               | 0.169              |
+| XPySom (NumPy)                                       | 256-core CPU | 8,096             | 40.6 s   | -               | 0.168              |
+| torchsom                                             | 256-core CPU | 1,024             | 39.2 s   | -               | 0.176              |
+
+At equal GPU memory fastersom is about 4x faster than XPySom, and about 20x faster on the CPU. torchsom updates the map
+per mini-batch, so it needs small batches and is far slower.
 
 ## Algorithm
 

{fastersom-0.2.0 → fastersom-0.3.0}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "fastersom"
-version = "0.2.0"
+version = "0.3.0"
 description = "Fast batch self-organising maps (SOM) in PyTorch for GPU and CPU"
 readme = "README.md"
 license = "Apache-2.0"

{fastersom-0.2.0 → fastersom-0.3.0}/pyproject.toml.orig

@@ -1,6 +1,6 @@
 [project]
 name = "fastersom"
-version = "0.2.0"
+version = "0.3.0"
 description = "Fast batch self-organising maps (SOM) in PyTorch for GPU and CPU"
 readme = "README.md"
 authors = [{ name = "Darius Schaub", email = "schaub.darius@gmail.com" }]

{fastersom-0.2.0 → fastersom-0.3.0}/src/fastersom/_som.py

@@ -51,8 +51,6 @@ def _is_available(device: torch.device) -> bool:
 class SOM:
     """Batch self-organising map on a rectangular grid with a gaussian neighbourhood.
 
-    The defaults reproduce the maps of spatiomic's ``so.dimension.som``.
-
     Args:
         n_rows: Number of grid rows.
         n_cols: Number of grid columns.
@@ -79,8 +77,11 @@ class SOM:
         batch_size: Samples processed at once in BMU search, aggregation, predict and transform. Affects memory
             (about batch_size x n_nodes floats) and speed only. Each epoch's update is the same up to float
             summation order, but over several epochs such rounding differences (as between devices) can flip
-            near-tied BMUs, so maps agree in quality rather than bit for bit. Defaults to 65536 on GPU and 8192
-            on CPU.
+            near-tied BMUs, so maps agree in quality rather than bit for bit. Defaults to 8192 on CPU and on GPUs
+            to 2^28 / n_nodes (where speed levels off), fewer if that does not fit in memory_fraction of the free GPU
+            memory (on MPS, of the memory torch recommends; 65536 before torch 2.5).
+        memory_fraction: Fraction of the free GPU memory the default batch_size may use, in (0, 1]. Defaults to 0.9
+            on CUDA and 0.5 on MPS, whose memory the system and other apps share.
 
     Attributes:
         weights_: Node weights, shape (n_rows * n_cols, n_features), row-major (node = row * n_cols + col).
@@ -100,6 +101,7 @@ class SOM:
         random_state: int | None = None,
         device: str | torch.device | None = None,
         batch_size: int | None = None,
+        memory_fraction: float | None = None,
     ):
         self.n_rows = n_rows
         self.n_cols = n_cols
@@ -111,6 +113,7 @@ class SOM:
         self.random_state = random_state
         self.device = device
         self.batch_size = batch_size
+        self.memory_fraction = memory_fraction
 
     @classmethod
     def from_weights(cls, weights: ArrayLike | torch.Tensor, n_rows: int, n_cols: int, **kwargs) -> "SOM":
@@ -130,6 +133,7 @@ class SOM:
 
     def __setstate__(self, state: dict) -> None:
         """Restores a pickled map on the configured device if available (cpu, cuda, mps), else the default one."""
+        state.setdefault("memory_fraction", None)  # pickled before 0.3
         self.__dict__.update(state)
         if hasattr(self, "weights_"):
             self._setup(fall_back=True)
@@ -145,24 +149,29 @@ class SOM:
                 raise ValueError(f"{name} must be an integer >= {minimum}, got {getattr(self, name)!r}")
         if self.batch_size is not None and not _is_int(self.batch_size, 1):
             raise ValueError(f"batch_size must be None or an integer >= 1, got {self.batch_size!r}")
+        if self.memory_fraction is not None and not 0 < self.memory_fraction <= 1:
+            raise ValueError(f"memory_fraction must be None or in (0, 1], got {self.memory_fraction!r}")
         self._device = torch.device(self.device) if self.device is not None else _default_device()
         if fall_back and not _is_available(self._device):
             self._device = _default_device()
-        self._batch_size = self.batch_size or (8192 if self._device.type == "cpu" else 1 << 16)
         self._sigma0 = ceil(min(self.n_rows, self.n_cols) / 2) / 2 if self.sigma_initial is None else self.sigma_initial
         self._sigma1 = 0.1 * self._sigma0 if self.sigma_final is None else self.sigma_final
         if not all(isfinite(sigma) and sigma > 0 for sigma in (self._sigma0, self._sigma1)):
             raise ValueError("sigma_initial and sigma_final must be finite and > 0")
         self._coords = torch.cartesian_prod(torch.arange(self.n_rows), torch.arange(self.n_cols))
         self._coords = self._coords.to(self._device, torch.float32)
-        self._grid_sq_dist = torch.cdist(self._coords, self._coords) ** 2
 
     def _as_tensor(self, X: ArrayLike | torch.Tensor) -> torch.Tensor:
-        return torch.as_tensor(
-            X.detach() if isinstance(X, torch.Tensor) else np.ascontiguousarray(X),
-            dtype=torch.float32,
-            device=self._device,
-        )
+        if isinstance(X, torch.Tensor):
+            X = X.detach()
+        elif hasattr(X, "__dlpack__") and not isinstance(X, np.ndarray):  # e.g. CuPy, without leaving the GPU
+            try:
+                X = torch.from_dlpack(X)
+            except RuntimeError:  # torch cannot import negative strides
+                X = torch.from_dlpack(X.copy())
+        else:  # numpy refuses to export read-only arrays through DLPack
+            X = np.ascontiguousarray(X)
+        return torch.as_tensor(X, dtype=torch.float32, device=self._device)
 
     def _check_X(self, X: ArrayLike | torch.Tensor) -> torch.Tensor:
         data = self._as_tensor(X)
@@ -171,10 +180,25 @@ class SOM:
         data = data.reshape(-1, data.shape[-1])
         if hasattr(self, "n_features_in_") and data.shape[1] != self.n_features_in_:
             raise ValueError(f"X has {data.shape[1]} features, but SOM was fitted with {self.n_features_in_}")
-        if not torch.isfinite(data).all():
+        if not torch.isfinite(torch.stack(torch.aminmax(data))).all():  # unlike isfinite(data), without a copy of X
             raise ValueError("X contains NaN or infinite values")
         return data
 
+    def _batch_size_for(self, data: torch.Tensor) -> int:
+        if self.batch_size is not None:
+            return self.batch_size
+        if self._device.type == "cuda":
+            cached = torch.cuda.memory_reserved(self._device) - torch.cuda.memory_allocated(self._device)
+            free, fraction = torch.cuda.mem_get_info(self._device)[0] + cached, 0.9
+        elif self._device.type == "mps" and hasattr(torch.mps, "recommended_max_memory"):  # torch >= 2.5
+            free, fraction = torch.mps.recommended_max_memory() - torch.mps.current_allocated_memory(), 0.5
+        else:
+            return 8192 if self._device.type == "cpu" else 1 << 16
+        fraction = fraction if self.memory_fraction is None else self.memory_fraction
+        n_nodes = self.n_rows * self.n_cols
+        # BMU search holds at most about 2 floats per sample and node, plus copies of the batch
+        return max(1, min((1 << 28) // n_nodes, int(fraction * free) // (8 * (n_nodes + data.shape[1]))))
+
     def _check_fitted(self) -> None:
         if not hasattr(self, "weights_"):
             raise RuntimeError("This SOM instance is not fitted yet, call fit() first")
@@ -233,11 +257,14 @@ class SOM:
     def _scores(self, batch: torch.Tensor, weights: torch.Tensor) -> torch.Tensor:
         """Per-node score to minimise.
 
-        ||w||^2 - 2 x.w (euclidean, ||x||^2 dropped), the L1 distance (manhattan) or -x.w / ||w|| of the (for
-        correlation centred) vectors (cosine, correlation).
+        ||w - c||^2 - 2 (x - c).(w - c) with c the mean node (euclidean, ||x - c||^2 dropped; centring keeps a large
+        offset from cancelling), the L1 distance (manhattan) or -x.w / ||w|| of the (for correlation centred) vectors
+        (cosine, correlation).
         """
         if self.metric == "euclidean":
-            return torch.addmm((weights**2).sum(1), batch, weights.T, alpha=-2)
+            centre = weights.mean(0)
+            weights = weights - centre
+            return torch.addmm((weights**2).sum(1), batch - centre, weights.T, alpha=-2)
         if self.metric == "manhattan":
             return torch.cdist(batch, weights, p=1)
         weights = self._centre(weights)
@@ -248,14 +275,15 @@ class SOM:
         """Distances from the scores of any number of nodes per sample, for every metric but euclidean."""
         if self.metric == "manhattan":
             return scores
-        return (1 + scores / self._centre(batch).norm(dim=1, keepdim=True)).clamp(0, 2)
+        # in place, as transform's scores hold a batch x nodes matrix; callers pass scores they no longer need
+        return scores.div_(self._centre(batch).norm(dim=1, keepdim=True)).add_(1).clamp_(0, 2)
 
     def _bmus(
         self, data: torch.Tensor, weights: torch.Tensor, return_distance: bool = True
     ) -> tuple[torch.Tensor, torch.Tensor | None]:
         """BMU and, if return_distance, distance to it per sample; -1 and NaN for samples without a BMU."""
         bmus, distances = [], []
-        for batch in data.split(self._batch_size):
+        for batch in data.split(self._batch_size_for(data)):
             score, bmu = self._scores(batch, weights).min(1)
             bmus.append(bmu)
             if not return_distance:
@@ -277,8 +305,9 @@ class SOM:
         """Fits the map, starting from the initialisation given by init.
 
         Args:
-            X: Samples, shape (n_samples, n_features) or (..., n_features); numpy array, array-like or torch tensor.
-                Any real values are fine; data is copied to the device once and stays there.
+            X: Samples, shape (n_samples, n_features) or (..., n_features); numpy array, array-like, torch tensor
+                or DLPack array such as CuPy. Any real values are fine; data is copied to the device once and stays
+                there.
             y: Ignored.
 
         Returns:
@@ -289,19 +318,21 @@ class SOM:
             self.__dict__.pop(attr, None)
         data = self._check_X(X)
         weights = self._init_weights(data)
-        n_nodes = weights.shape[0]
+        n_nodes, batch_size = weights.shape[0], self._batch_size_for(data)
         for t in range(self.n_epochs):
             sigma = self._sigma0 * (self._sigma1 / self._sigma0) ** (t / self.n_epochs)
             sums, counts = torch.zeros_like(weights), weights.new_zeros(n_nodes)
-            for batch in data.split(self._batch_size):
+            for batch in data.split(batch_size):
                 bmus = self._scores(batch, weights).argmin(1)
                 keep = self._has_bmu(batch)  # samples without a BMU do not contribute
                 if keep is not None:
                     batch, bmus = batch[keep], bmus[keep]
                 sums.index_add_(0, bmus, batch)
                 counts += torch.bincount(bmus, minlength=n_nodes)
-            h = torch.exp(-self._grid_sq_dist / (2 * sigma**2))
+            # in place, so large maps hold one nodes x nodes matrix at a time
+            h = torch.cdist(self._coords, self._coords).square_().mul_(-0.5 / sigma**2).exp_()
             numerator, denominator = h @ sums, h @ counts
+            del h  # before the next epoch's BMU search and neighbourhood matrix
             weights = torch.where(denominator[:, None] > 0, numerator / denominator[:, None], weights)
         self._set_weights(weights)
         self.labels_ = self._bmus(data, weights, return_distance=False)[0].cpu().numpy()
@@ -352,7 +383,12 @@ class SOM:
         """
         self._check_fitted()
         data = self._check_X(X)
-        return torch.cat([self._distances(batch) for batch in data.split(self._batch_size)]).cpu().numpy()
+        batch_size = self._batch_size_for(data)
+        # on the CPU, as it can exceed device memory
+        distances = torch.empty(len(data), len(self._weights), dtype=torch.float32, device="cpu")
+        for batch, rows in zip(data.split(batch_size), distances.split(batch_size), strict=True):
+            rows.copy_(self._distances(batch))
+        return distances.numpy()
 
     def fit_transform(self, X: ArrayLike | torch.Tensor, y: object = None) -> np.ndarray:
         """Fits the map and returns transform(X)."""
@@ -376,7 +412,7 @@ class SOM:
         nodes = torch.cat(
             [
                 self._scores(batch, self._weights).topk(2, dim=1, largest=False).indices
-                for batch in data.split(self._batch_size)
+                for batch in data.split(self._batch_size_for(data))
             ]
         )
         offset = (self._coords[nodes[:, 0]] - self._coords[nodes[:, 1]]).abs().amax(1)