meta tensor basic profiling

.gitignore

@@ -131,10 +131,14 @@ tests/testing_data/MedNIST*
 tests/testing_data/*Hippocampus*
 tests/testing_data/*.tiff
 tests/testing_data/schema.json
+*.svg
 
 # clang format tool
 .clang-format-bin/
 
 # VSCode
 .vscode/
 *.zip
+
+# profiling results
+*.prof

tests/profile_subclass/README.md

@@ -0,0 +1,43 @@
+# Profiling the performance of subclassing/`__torch_function__` in MONAI
+
+## Requirements
+```bash
+pip install py-spy
+pip install snakeviz  # for viewing the cProfile results
+```
+
+## Commands
+
+### Install MONAI
+```
+./runtests.sh --build   # from monai's root directory
+```
+or follow the installation guide (https://docs.monai.io/en/latest/installation.html)
+
+### Profiling the task of adding two MetaTensors
+```bash
+python profiling.py
+```
+
+### Profiling using `py-spy`
+```bash
+py-spy record -o Tensor.svg -- python pyspy_profiling.py Tensor
+py-spy record -o SubTensor.svg -- python pyspy_profiling.py SubTensor
+py-spy record -o SubWithTorchFunc.svg -- python pyspy_profiling.py SubWithTorchFunc
+py-spy record -o MetaTensor.svg -- python pyspy_profiling.py MetaTensor
+```
+
+### Profiling using `cProfile` and `SNAKEVIZ`
+
+```bash
+python cprofile_profiling.py
+snakeviz out_200.prof
+```
+
+---
+These tests are based on the following code:
+https://github.com/pytorch/pytorch/tree/v1.11.0/benchmarks/overrides_benchmark
+
+- Overhead for torch functions when run on `torch.Tensor` objects is on the order of 2 microseconds.
+- `__torch_function__` should add zero overhead for `torch.Tensor` inputs, a small overhead for subclasses of `torch.Tensor`, and an order of microseconds for `MeatTensor`.
+- Changing the dispatching mechanism may result in changes that are on the order of 100 ns, which are hard to detect due to noise, but important.

tests/profile_subclass/cprofile_profiling.py

@@ -0,0 +1,28 @@
+# Copyright (c) MONAI Consortium
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#     http://www.apache.org/licenses/LICENSE-2.0
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""
+Profiling MetaTensor
+"""
+
+import cProfile
+
+import torch
+
+from monai.data.meta_tensor import MetaTensor
+
+if __name__ == "__main__":
+    n_chan = 3
+    for hwd in (10, 200):
+        shape = (n_chan, hwd, hwd, hwd)
+        a = MetaTensor(torch.rand(shape), meta={"affine": torch.eye(4) * 2, "fname": "something1"})
+        b = MetaTensor(torch.rand(shape), meta={"affine": torch.eye(4) * 3, "fname": "something2"})
+        cProfile.run("c = a + b", filename=f"out_{hwd}.prof")

tests/profile_subclass/min_classes.py

@@ -0,0 +1,29 @@
+# Copyright (c) MONAI Consortium
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#     http://www.apache.org/licenses/LICENSE-2.0
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+"""
+Minimal subclassing as baselines
+Adapted from https://github.com/pytorch/pytorch/tree/v1.11.0/benchmarks/overrides_benchmark
+"""
+
+import torch
+
+__all__ = ["SubTensor", "SubWithTorchFunc"]
+
+
+class SubTensor(torch.Tensor):
+    pass
+
+
+class SubWithTorchFunc(torch.Tensor):
+    def __torch_function__(self, func, types, args=(), kwargs=None):
+        return super().__torch_function__(func, types, args, {} if kwargs is None else kwargs)

tests/profile_subclass/profiling.py

@@ -0,0 +1,72 @@
+# Copyright (c) MONAI Consortium
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#     http://www.apache.org/licenses/LICENSE-2.0
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""
+Comparing torch.Tensor, SubTensor, SubWithTorchFunc, MetaTensor
+Adapted from https://github.com/pytorch/pytorch/tree/v1.11.0/benchmarks/overrides_benchmark
+"""
+import argparse
+
+import torch
+from min_classes import SubTensor, SubWithTorchFunc
+
+from monai.data import MetaTensor
+from monai.utils.profiling import PerfContext
+
+NUM_REPEATS = 1000
+NUM_REPEAT_OF_REPEATS = 1000
+
+
+def bench(t1, t2):
+    bench_times = []
+    for _ in range(NUM_REPEAT_OF_REPEATS):
+        with PerfContext() as pc:
+            for _ in range(NUM_REPEATS):
+                torch.add(t1, t2)
+        bench_times.append(pc.total_time)
+
+    bench_time_min = float(torch.min(torch.Tensor(bench_times))) / NUM_REPEATS
+    bench_time_avg = float(torch.sum(torch.Tensor(bench_times))) / (NUM_REPEATS * NUM_REPEAT_OF_REPEATS)
+    bench_time_med = float(torch.median(torch.Tensor(bench_times))) / NUM_REPEATS
+    bench_std = float(torch.std(torch.Tensor(bench_times))) / NUM_REPEATS
+    return bench_time_min, bench_time_avg, bench_time_med, bench_std
+
+
+def main():
+    global NUM_REPEATS
+    global NUM_REPEAT_OF_REPEATS
+
+    parser = argparse.ArgumentParser(description="Run the __torch_function__ benchmarks.")
+    parser.add_argument(
+        "--nreps", "-n", type=int, default=NUM_REPEATS, help="The number of repeats for one measurement."
+    )
+    parser.add_argument("--nrepreps", "-m", type=int, default=NUM_REPEAT_OF_REPEATS, help="The number of measurements.")
+    args = parser.parse_args()
+
+    NUM_REPEATS = args.nreps
+    NUM_REPEAT_OF_REPEATS = args.nrepreps
+
+    types = torch.Tensor, SubTensor, SubWithTorchFunc, MetaTensor
+
+    for t in types:
+        tensor_1 = t(1)
+        tensor_2 = t(2)
+
+        b_min, b_avg, b_med, b_std = bench(tensor_1, tensor_2)
+        print(
+            "Type {} time (microseconds):  min: {}, avg: {}, median: {}, and std {}.".format(
+                t.__name__, (10**6 * b_min), (10**6) * b_avg, (10**6) * b_med, (10**6) * b_std
+            )
+        )
+
+
+if __name__ == "__main__":
+    main()

tests/profile_subclass/pyspy_profiling.py

@@ -0,0 +1,40 @@
+# Copyright (c) MONAI Consortium
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#     http://www.apache.org/licenses/LICENSE-2.0
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""
+To be used with py-spy, comparing torch.Tensor, SubTensor, SubWithTorchFunc, MetaTensor
+Adapted from https://github.com/pytorch/pytorch/tree/v1.11.0/benchmarks/overrides_benchmark
+"""
+import argparse
+
+import torch
+from min_classes import SubTensor, SubWithTorchFunc  # noqa: F401
+
+from monai.data import MetaTensor  # noqa: F401
+
+Tensor = torch.Tensor
+
+NUM_REPEATS = 1000000
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="Run the torch.add for a given class a given number of times.")
+    parser.add_argument("tensor_class", metavar="TensorClass", type=str, help="The class to benchmark.")
+    parser.add_argument("--nreps", "-n", type=int, default=NUM_REPEATS, help="The number of repeats.")
+    args = parser.parse_args()
+
+    TensorClass = globals()[args.tensor_class]
+    NUM_REPEATS = args.nreps
+
+    t1 = TensorClass(1)
+    t2 = TensorClass(2)
+
+    for _ in range(NUM_REPEATS):
+        torch.add(t1, t2)