Merge branch 'main' of github.com:lxuechen/ml-swissknife into main

README.md

@@ -25,12 +25,11 @@ pip install git+https://github.com/lxuechen/ml-swissknife.git
     - [x] simultaneous iteration
     - [x] Lanczos
     - [x] spectral density plotting
-    - [ ] distributed simultaneous iteration
+    - [x] distributed simultaneous iteration
 - NLP
     - [x] New token embedding for GPT-2
     - [x] Evaluation pipeline for generation
     - [x] Turk templates
-    - [ ] Decoding utils, fast dataset loading (w/ auto-download)
 - Misc
     - [ ] Fast two-sample test utils
     - [ ] Helper code for calibration metrics (reliability diagram, ECE, Brier score)
@@ -39,5 +38,6 @@ pip install git+https://github.com/lxuechen/ml-swissknife.git
     - [x] Confidence interval utils
 - Data
     - [ ] UTKFaces
+    - [ ] data2text
 - wandb
     - [x] Project-based helper for downloading files

ml_swissknife/numerical_distributed.py

@@ -21,7 +21,6 @@ def orthogonal_iteration(
     k: int,
     num_power_iteration=1,
     disable_tqdm=False,
-    dtype=torch.float,
     orthogonalization_chunk_size=10,
     dim0_chunk_size: Optional[int] = None,
     callback: Optional[Callable] = None,
@@ -33,14 +32,15 @@ def orthogonal_iteration(
     WARNING:
         Good reconstruction of the data does not imply converged eigenvalues!
 
+    This function can be run on CPU, single GPU, or multiple GPUs on a single machine.
+
     Args:
         input_mat: DataLoader or Dataset or torch.Tensor as data.
             The underlying tensor is of size (n, p), and we want the eigenvalues and eigenvectors of a p x p matrix.
             That is we always assume batch first.
         k: Number of principal components to return.
         num_power_iteration: Number of power iterations.
         disable_tqdm: If True, disable progress bar.
-        dtype: Precision in string format.
         dim0_chunk_size: Size of chunks for dim0 -- the batch dimension of input_mat.
             Reduce to save memory in matmul and rayleigh quotient computation.
         orthogonalization_chunk_size: Size of chunks for orthogonalization.
@@ -61,7 +61,7 @@ def orthogonal_iteration(
     k = min(k, p, n)
 
     eigenvalues = None
-    eigenvectors = torch.randn(size=(p, k), dtype=dtype)  # This step will be very slow for large models.
+    eigenvectors = torch.randn(size=(p, k), dtype=batch.dtype)  # This step will be very slow for large models.
 
     if callback is not None:
         callback(0, eigenvalues, eigenvectors)

ml_swissknife/utils.py

@@ -67,6 +67,13 @@ def int2str(x, leading_zeros=8):
     return f"{x:0{leading_zeros}d}"
 
 
+def average_over_seed(seq_of_seq):  # Here purely due to backward compatibility.
+    min_len = min(len(seq) for seq in seq_of_seq)
+    seq_of_seq = [seq[:min_len] for seq in seq_of_seq]
+    seq_of_seq = np.array(seq_of_seq)
+    return seq_of_seq.mean(0), seq_of_seq.std(0)
+
+
 def average_metric_over_seeds(*seqs: Union[Sequence[Numeric], Sequence[Dict[str, Numeric]]]):
     # seqs is an iterable. Each seq is a sequence of numbers or dicts to average over.
     single_input = len(seqs) == 1

tests/test_numerical_distributed.py

@@ -0,0 +1,26 @@
+import torch
+from torch.utils.data import DataLoader, TensorDataset
+
+from ml_swissknife import numerical_distributed
+
+
+def test_orthogonal_iteration(n=100, d=30, k=10, bsz=10):
+    torch.set_default_dtype(torch.float64)
+
+    features = torch.randn(n, d)
+    dataset = TensorDataset(features)
+    loader = DataLoader(dataset, batch_size=bsz)
+
+    eigenvalues, eigenvectors = numerical_distributed.orthogonal_iteration(
+        input_mat=loader,
+        k=k,
+        num_power_iteration=400,
+        dtype=torch.get_default_dtype(),
+    )
+    eigenvalues_expected, eigenvectors_expected = torch.linalg.eigh(features.T @ features)
+    print(eigenvalues)
+    print(eigenvalues_expected.flip(dims=(0,)))
+    print('---')
+    print(eigenvectors)
+    print(eigenvectors_expected.flip(dims=(1,)))
+    torch.testing.assert_allclose(eigenvalues, eigenvalues_expected.flip(dims=(0,))[:k], atol=1e-4, rtol=1e-4)