docs: update readme and benchmarks

Author: percevalw

README.md

@@ -12,6 +12,7 @@
 [![Coverage](https://raw.githubusercontent.com/aphp/foldedtensor/coverage/coverage.svg)](https://raw.githubusercontent.com/aphp/foldedtensor/coverage/coverage.txt)
 [![License](https://img.shields.io/github/license/aphp/foldedtensor?color=x&style=flat-square)](https://github.com/aphp/foldedtensor/blob/main/LICENSE)
 ![PyPI - Downloads](https://img.shields.io/pypi/dm/foldedtensor?style=flat-square&color=purple)
+![Python versions](https://img.shields.io/pypi/pyversions/foldedtensor?style=flat-square)
 
 # FoldedTensor: PyTorch extension for handling deeply nested sequences of variable length
 

docs/benchmark.md

@@ -7,9 +7,11 @@ This file was generated from [`scripts/benchmark.py`](../scripts/benchmark.py).
 It compares the performance of `foldedtensor` with various alternatives for padding
 and working with nested lists and tensors.
 
-Versions:
-- `torch.__version__ == '2.0.1'`
-- `foldedtensor.__version__ == '0.3.3'`
+Environment:
+- `torch.__version__ == '2.6.0'`
+- `foldedtensor.__version__ == '0.4.0'`
+- `python == 3.9.20`
+- `sys.platform == 'darwin'`
 
 
 ## Case 1 (pad variable lengths nested list)
@@ -20,79 +22,79 @@ nested_list = make_nested_list(32, (50, 100), (25, 30), value=1)
 
 Comparisons:
 %timeit python_padding(nested_list)
-# 100 loops, best of 5: 13.24 ms per loop
+# 100 loops, best of 5: 15.09 ms per loop
 
 %timeit foldedtensor.as_folded_tensor(nested_list)
-# 100 loops, best of 5: 0.63 ms per loop
+# 100 loops, best of 5: 0.73 ms per loop
 
 ```
-
+Speedup against best alternative: **20.67x** :rocket:
 
 ## Case 2 (same lengths nested lists)
 
 ```python
 nested_list = make_nested_list(32, 100, 30, value=1)
 
 %timeit torch.tensor(nested_list)
-# 100 loops, best of 5: 6.44 ms per loop
+# 100 loops, best of 5: 6.51 ms per loop
 
 %timeit torch.LongTensor(nested_list)
-# 100 loops, best of 5: 2.64 ms per loop
+# 100 loops, best of 5: 2.78 ms per loop
 
 %timeit python_padding(nested_list)
-# 100 loops, best of 5: 16.68 ms per loop
+# 100 loops, best of 5: 18.38 ms per loop
 
 %timeit torch.nested.nested_tensor([torch.LongTensor(sub) for sub in nested_list]).to_padded_tensor(0)
-# 100 loops, best of 5: 2.90 ms per loop
+# 100 loops, best of 5: 3.00 ms per loop
 
 %timeit foldedtensor.as_folded_tensor(nested_list)
-# 100 loops, best of 5: 0.96 ms per loop
+# 100 loops, best of 5: 1.08 ms per loop
 
 ```
-
+Speedup against best alternative: **2.58x** :rocket:
 
 ## Case 3 (simple list)
 
 ```python
 simple_list = make_nested_list(10000, value=1)
 
 %timeit torch.tensor(simple_list)
-# 100 loops, best of 5: 0.65 ms per loop
+# 100 loops, best of 5: 0.63 ms per loop
 
 %timeit torch.LongTensor(simple_list)
 # 100 loops, best of 5: 0.27 ms per loop
 
 %timeit python_padding(simple_list)
-# 100 loops, best of 5: 0.27 ms per loop
+# 100 loops, best of 5: 0.28 ms per loop
 
 %timeit foldedtensor.as_folded_tensor(simple_list)
 # 100 loops, best of 5: 0.08 ms per loop
 
 ```
-
+Speedup against best alternative: **3.32x** :rocket:
 
 ## Case 4 (same lengths nested lists to flat tensor)
 
 ```python
 nested_list = make_nested_list(32, 100, 30, value=1)
 
 %timeit torch.tensor(nested_list).view(-1)
-# 100 loops, best of 5: 6.67 ms per loop
+# 100 loops, best of 5: 6.52 ms per loop
 
 %timeit torch.LongTensor(nested_list).view(-1)
-# 100 loops, best of 5: 2.74 ms per loop
+# 100 loops, best of 5: 2.76 ms per loop
 
 %timeit python_padding(nested_list).view(-1)
-# 100 loops, best of 5: 17.16 ms per loop
+# 100 loops, best of 5: 18.62 ms per loop
 
 %timeit foldedtensor.as_folded_tensor(nested_list).view(-1)
-# 100 loops, best of 5: 1.02 ms per loop
+# 100 loops, best of 5: 1.12 ms per loop
 
 %timeit foldedtensor.as_folded_tensor(nested_list, data_dims=(2,))
-# 100 loops, best of 5: 0.95 ms per loop
+# 100 loops, best of 5: 1.08 ms per loop
 
 ```
-
+Speedup against best alternative: **2.47x** :rocket:
 ## Case 5 (variable lengths nested lists) to padded embeddings
 
 Nested lists with different lengths (second level lists have lengths between 50 and 150). We compare `foldedtensor` with `torch.nested`.
@@ -102,37 +104,41 @@ nested_list = make_nested_list(32, (50, 150), 30, value=1)
 # Padding with 0
 
 %timeit torch.nested.nested_tensor([torch.LongTensor(sub) for sub in nested_list]).to_padded_tensor(0)
-# 100 loops, best of 5: 3.11 ms per loop
+# 100 loops, best of 5: 3.02 ms per loop
 
 %timeit foldedtensor.as_folded_tensor(nested_list).as_tensor()
-# 100 loops, best of 5: 0.90 ms per loop
+# 100 loops, best of 5: 1.03 ms per loop
 
+```
+Speedup against best alternative: **2.95x** :rocket:
+```python
 # Padding with 1
 
 %timeit torch.nested.nested_tensor([torch.FloatTensor(sub) for sub in nested_list]).to_padded_tensor(1)
-# 100 loops, best of 5: 3.57 ms per loop
+# 100 loops, best of 5: 3.72 ms per loop
 
 %timeit x = foldedtensor.as_folded_tensor(nested_list); x.masked_fill_(x.mask, 1)
-# 100 loops, best of 5: 1.33 ms per loop
+# 100 loops, best of 5: 1.62 ms per loop
 
 ```
-
+Speedup against best alternative: **2.30x** :rocket:
 
 ## Case 6 (2d padding)
 
 ```python
 nested_list = make_nested_list(160, (50, 150), value=1)
 
 %timeit python_padding(nested_list)
-# 100 loops, best of 5: 1.24 ms per loop
+# 100 loops, best of 5: 1.33 ms per loop
 
 %timeit torch.nested.nested_tensor([torch.LongTensor(sub) for sub in nested_list]).to_padded_tensor(0)
-# 100 loops, best of 5: 1.09 ms per loop
+# 100 loops, best of 5: 1.14 ms per loop
 
 %timeit torch.nn.utils.rnn.pad_sequence([torch.LongTensor(sub) for sub in nested_list], batch_first=True, padding_value=0)
-# 100 loops, best of 5: 0.78 ms per loop
+# 100 loops, best of 5: 0.86 ms per loop
 
 %timeit foldedtensor.as_folded_tensor(nested_list)
-# 100 loops, best of 5: 0.13 ms per loop
+# 100 loops, best of 5: 0.15 ms per loop
 
 ```
+Speedup against best alternative: **5.88x** :rocket:

scripts/benchmark.py

@@ -1,15 +1,16 @@
 # ruff: noqa: F401, E501
 import contextlib
 import random
+import subprocess
+import sys
 import warnings
 from timeit import Timer
 
+import foldedtensor  # noqa: F401
 import torch
 import torch.nested
 import torch.nn.utils.rnn
 
-import foldedtensor  # noqa: F401
-
 warnings.filterwarnings("ignore")
 
 torch.set_default_device("cpu")
@@ -108,6 +109,7 @@ def format_time(dt):
         "# %d loop%s, best of %d: %s per loop\n"
         % (number, "s" if number != 1 else "", repeat, format_time(best))
     )
+    return best
 
 
 print(
@@ -120,9 +122,11 @@ def format_time(dt):
 It compares the performance of `foldedtensor` with various alternatives for padding
 and working with nested lists and tensors.
 
-Versions:
+Environment:
 - `torch.__version__ == {torch.__version__!r}`
 - `foldedtensor.__version__ == {foldedtensor.__version__!r}`
+- `python == {sys.version_info.major}.{sys.version_info.minor}.{sys.version_info.micro}`
+- `sys.platform == {sys.platform!r}`
 """
 )
 
@@ -139,40 +143,53 @@ def format_time(dt):
             exec_and_print("nested_list = make_nested_list(32, (50, 100), (25, 30), value=1)")
 
             print("Comparisons:")
-            timeit("python_padding(nested_list)")
-            timeit("foldedtensor.as_folded_tensor(nested_list)")
+            alt = []
+            alt.append(timeit("python_padding(nested_list)"))
+            ft_time = timeit("foldedtensor.as_folded_tensor(nested_list)")
+
+        print(f"Speedup against best alternative: **{alt[0] / ft_time:.2f}x** :rocket:")
 
     if 2 in cases:
         print("\n## Case 2 (same lengths nested lists)\n")
 
         with block_code():
             exec_and_print("nested_list = make_nested_list(32, 100, 30, value=1)")
-            timeit("torch.tensor(nested_list)")
-            timeit("torch.LongTensor(nested_list)")
-            timeit("python_padding(nested_list)")
-            timeit("torch.nested.nested_tensor([torch.LongTensor(sub) for sub in nested_list]).to_padded_tensor(0)")
-            timeit("foldedtensor.as_folded_tensor(nested_list)")
+            alt = []
+            alt.append(timeit("torch.tensor(nested_list)"))
+            alt.append(timeit("torch.LongTensor(nested_list)"))
+            alt.append(timeit("python_padding(nested_list)"))
+            alt.append(timeit("torch.nested.nested_tensor([torch.LongTensor(sub) for sub in nested_list]).to_padded_tensor(0)"))
+            ft_time = timeit("foldedtensor.as_folded_tensor(nested_list)")
+
+        print(f"Speedup against best alternative: **{min(alt) / ft_time:.2f}x** :rocket:")
 
     if 3 in cases:
         print("\n## Case 3 (simple list)\n")
 
         with block_code():
             exec_and_print("simple_list = make_nested_list(10000, value=1)")
-            timeit("torch.tensor(simple_list)")
-            timeit("torch.LongTensor(simple_list)")
-            timeit("python_padding(simple_list)")
-            timeit("foldedtensor.as_folded_tensor(simple_list)")
+            alt = []
+            alt.append(timeit("torch.tensor(simple_list)"))
+            alt.append(timeit("torch.LongTensor(simple_list)"))
+            alt.append(timeit("python_padding(simple_list)"))
+            ft_time = timeit("foldedtensor.as_folded_tensor(simple_list)")
+
+        print(f"Speedup against best alternative: **{min(alt) / ft_time:.2f}x** :rocket:")
 
     if 4 in cases:
         print("\n## Case 4 (same lengths nested lists to flat tensor)\n")
 
         with block_code():
             exec_and_print("nested_list = make_nested_list(32, 100, 30, value=1)")
-            timeit("torch.tensor(nested_list).view(-1)")
-            timeit("torch.LongTensor(nested_list).view(-1)")
-            timeit("python_padding(nested_list).view(-1)")
-            timeit("foldedtensor.as_folded_tensor(nested_list).view(-1)")
-            timeit("foldedtensor.as_folded_tensor(nested_list, data_dims=(2,))")
+            alt = []
+            ft_times = []
+            alt.append(timeit("torch.tensor(nested_list).view(-1)"))
+            alt.append(timeit("torch.LongTensor(nested_list).view(-1)"))
+            alt.append(timeit("python_padding(nested_list).view(-1)"))
+            ft_times.append(timeit("foldedtensor.as_folded_tensor(nested_list).view(-1)"))
+            ft_times.append(timeit("foldedtensor.as_folded_tensor(nested_list, data_dims=(2,))"))
+
+        print(f"Speedup against best alternative: **{min(alt) / max(ft_times):.2f}x** :rocket:")
 
     if 5 in cases:
         print("## Case 5 (variable lengths nested lists) to padded embeddings\n")
@@ -184,26 +201,34 @@ def format_time(dt):
 
             print("# Padding with 0\n")
 
-            timeit("torch.nested.nested_tensor([torch.LongTensor(sub) for sub in nested_list]).to_padded_tensor(0)")
-            timeit("foldedtensor.as_folded_tensor(nested_list).as_tensor()")
+            nt_time = timeit("torch.nested.nested_tensor([torch.LongTensor(sub) for sub in nested_list]).to_padded_tensor(0)")
+            ft_time = timeit("foldedtensor.as_folded_tensor(nested_list).as_tensor()")
+
+        print(f"Speedup against best alternative: **{nt_time / ft_time:.2f}x** :rocket:")
 
+        with block_code():
             print("# Padding with 1\n")
-            timeit("torch.nested.nested_tensor([torch.FloatTensor(sub) for sub in nested_list]).to_padded_tensor(1)")
-            timeit("x = foldedtensor.as_folded_tensor(nested_list); x.masked_fill_(x.mask, 1)")
+            nt_time = timeit("torch.nested.nested_tensor([torch.FloatTensor(sub) for sub in nested_list]).to_padded_tensor(1)")
+            ft_time = timeit("x = foldedtensor.as_folded_tensor(nested_list); x.masked_fill_(x.mask, 1)")
+
+        print(f"Speedup against best alternative: **{nt_time / ft_time:.2f}x** :rocket:")
 
     if 6 in cases:
         print("\n## Case 6 (2d padding)\n")
 
         with block_code():
             exec_and_print("nested_list = make_nested_list(160, (50, 150), value=1)")
 
-            timeit("python_padding(nested_list)")
-            timeit("torch.nested.nested_tensor([torch.LongTensor(sub) for sub in nested_list]).to_padded_tensor(0)")
-            timeit(
-                "torch.nn.utils.rnn.pad_sequence([torch.LongTensor(sub) for sub in nested_list], batch_first=True, padding_value=0)")
-            timeit("foldedtensor.as_folded_tensor(nested_list)")
+            alt = []
+            alt.append(timeit("python_padding(nested_list)"))
+            alt.append(timeit("torch.nested.nested_tensor([torch.LongTensor(sub) for sub in nested_list]).to_padded_tensor(0)"))
+            alt.append(timeit("torch.nn.utils.rnn.pad_sequence([torch.LongTensor(sub) for sub in nested_list], batch_first=True, padding_value=0)"))
+            ft_time = timeit("foldedtensor.as_folded_tensor(nested_list)")
+
+        print(f"Speedup against best alternative: **{min(alt) / ft_time:.2f}x** :rocket:")
 
     if 7 in cases:
+        # Test case not working yet
 
         def sum_all_words_per_sample(ft):
             lengths = ft.lengths
@@ -235,8 +260,10 @@ def sum_all_words_per_sample(ft):
                 "nt = embedder(nt)\n"
             )
 
-            timeit("nt.sum(dim=1)")
-            timeit("sum_all_words_per_sample(ft)")
+            nt_time = timeit("nt.sum(dim=1)")
+            ft_time = timeit("sum_all_words_per_sample(ft)")
+
+        print(f"Speedup against best alternative: **{nt_time / ft_time:.2f}x** :rocket:")
 
         # timeit("embedder(ft)")
         # timeit("embedder(ft).refold(0, 1)")