Repository: pfnet/pytorch-pfn-extras

Issue:
Profiler: Automatically fill `tag` in `record`?
Maybe we can use the caller's function name (`inspect.stack()`) if tag is not given.

TODO: Need to measure overhead.


Discussion:

Resolving patch:
diff --git a/pytorch_pfn_extras/profiler/_record.py b/pytorch_pfn_extras/profiler/_record.py
--- a/pytorch_pfn_extras/profiler/_record.py
+++ b/pytorch_pfn_extras/profiler/_record.py
@@ -1,17 +1,35 @@
 from contextlib import contextmanager
+import inspect
 from typing import Any, Callable, Generator, Iterable, Optional, TypeVar
+import types
 
 import torch
 
 from pytorch_pfn_extras.profiler._time_summary import time_summary, _ReportNotification
 
 
+def _infer_tag_name(frame: Optional[types.FrameType], depth: int) -> str:
+    for _ in range(depth):
+        assert frame is not None
+        frame = frame.f_back
+    assert frame is not None
+    frame_info = inspect.getframeinfo(frame, context=0)
+    return '{}:{}:{}'.format(
+        inspect.getmodulename(frame_info.filename),
+        frame_info.lineno,
+        frame_info.function,
+    )
+
+
 @contextmanager
 def record(
-        tag: str,
+        tag: Optional[str],
         metric: Optional[str] = None,
         use_cuda: bool = False,
 ) -> Generator[_ReportNotification, None, None]:
+    if tag is None:
+        tag = _infer_tag_name(inspect.currentframe(), depth=2)
+
     if metric is None:
         metric = tag
 
@@ -30,12 +48,12 @@ def record(
 
 
 def record_function(
-        tag: str,
+        tag: Optional[str],
         use_cuda: bool = False,
 ) -> Callable[[Callable[..., _T]], Callable[..., _T]]:
     def wrapper(f: Callable[..., _T]) -> Callable[..., _T]:
         def wrapped(*args: Any, **kwargs: Any) -> _T:
-            with record(tag, use_cuda=use_cuda):
+            with record(tag or f.__name__, use_cuda=use_cuda):
                 return f(*args, **kwargs)
 
         return wrapped
@@ -44,11 +62,14 @@ def wrapped(*args: Any, **kwargs: Any) -> _T:
 
 
 def record_iterable(
-    tag: str,
-    iter: Iterable[_T],
-    divide_metric: bool = False,
-    use_cuda: bool = False,
+        tag: Optional[str],
+        iter: Iterable[_T],
+        divide_metric: bool = False,
+        use_cuda: bool = False,
 ) -> Iterable[_T]:
+    if tag is None:
+        tag = _infer_tag_name(inspect.currentframe(), depth=1)
+
     def wrapped() -> Iterable[_T]:
         for i, x in enumerate(iter):
             name = f"{tag}-{i}"


Test patch:
diff --git a/tests/pytorch_pfn_extras_tests/profiler_tests/test_record.py b/tests/pytorch_pfn_extras_tests/profiler_tests/test_record.py
new file mode 100644
--- /dev/null
+++ b/tests/pytorch_pfn_extras_tests/profiler_tests/test_record.py
@@ -0,0 +1,123 @@
+import os
+from packaging import version
+
+import pytest
+import torch
+
+import pytorch_pfn_extras as ppe
+
+
+_profiler_available = (
+    os.name != 'nt'
+    or version.Version(torch.__version__) >= version.Version("1.9")
+)
+
+
+@pytest.mark.skipif(not _profiler_available, reason="profiler is not available")
+@pytest.mark.parametrize('device', ['cpu', 'cuda'])
+def test_record(device):
+    model = torch.nn.Linear(30, 40)
+    model.to(device)
+    x = torch.arange(30, dtype=torch.float32).to(device)
+
+    with torch.profiler.profile() as prof:
+        with ppe.profiler.record('my_tag_1'):
+            model(x)
+
+    keys = [event.key for event in prof.key_averages()]
+    assert 'my_tag_1' in keys
+    assert 'aten::linear' in keys
+
+
+@pytest.mark.skipif(not _profiler_available, reason="profiler is not available")
+@pytest.mark.parametrize('device', ['cpu', 'cuda'])
+def test_record_without_tag(device):
+    model = torch.nn.Linear(30, 40)
+    model.to(device)
+    x = torch.arange(30, dtype=torch.float32).to(device)
+
+    with torch.profiler.profile() as prof:
+        with ppe.profiler.record(None):
+            model(x)
+
+    keys = [event.key for event in prof.key_averages()]
+    assert 'aten::linear' in keys
+    assert any(k.endswith('test_record_without_tag') for k in keys)
+
+
+@pytest.mark.skipif(not _profiler_available, reason="profiler is not available")
+@pytest.mark.parametrize('device', ['cpu', 'cuda'])
+def test_record_function(device):
+    model = torch.nn.Linear(30, 40)
+    model.to(device)
+
+    @ppe.profiler.record_function('my_tag_2')
+    def my_run(x):
+        model(x)
+
+    with torch.profiler.profile() as prof:
+        x = torch.arange(30, dtype=torch.float32).to(device)
+        my_run(x)
+
+    keys = [event.key for event in prof.key_averages()]
+    assert 'aten::linear' in keys
+    assert 'my_tag_2' in keys
+
+
+@pytest.mark.skipif(not _profiler_available, reason="profiler is not available")
+@pytest.mark.parametrize('device', ['cpu', 'cuda'])
+def test_record_function_without_tag(device):
+    model = torch.nn.Linear(30, 40)
+    model.to(device)
+    x = torch.arange(30, dtype=torch.float32).to(device)
+
+    @ppe.profiler.record_function(None)
+    def my_run(x):
+        model(x)
+
+    with torch.profiler.profile() as prof:
+        my_run(x)
+
+    keys = [event.key for event in prof.key_averages()]
+    assert 'aten::linear' in keys
+    assert 'my_run' in keys
+
+
+@pytest.mark.skipif(not _profiler_available, reason="profiler is not available")
+@pytest.mark.parametrize('device', ['cpu', 'cuda'])
+def test_record_iterable(device):
+    model = torch.nn.Linear(30, 40)
+    model.to(device)
+
+    x = torch.arange(30, dtype=torch.float32).to(device)
+    iters = [x, x, x]
+
+    with torch.profiler.profile() as prof:
+        for x in ppe.profiler.record_iterable('my_tag_3', iters):
+            model(x)
+
+    keys = [event.key for event in prof.key_averages()]
+    assert 'aten::linear' in keys
+    assert 'my_tag_3-0' in keys
+    assert 'my_tag_3-1' in keys
+    assert 'my_tag_3-2' in keys
+
+
+@pytest.mark.skipif(not _profiler_available, reason="profiler is not available")
+@pytest.mark.parametrize('device', ['cpu', 'cuda'])
+def test_record_iterable_without_tag(device):
+    model = torch.nn.Linear(30, 40)
+    model.to(device)
+
+    x = torch.arange(30, dtype=torch.float32).to(device)
+    iters = [x, x, x]
+
+    with torch.profiler.profile() as prof:
+        for x in ppe.profiler.record_iterable(None, iters):
+            model(x)
+
+    keys = [event.key for event in prof.key_averages()]
+    assert 'aten::linear' in keys
+    assert any(k.endswith('test_record_iterable_without_tag-0') for k in keys)
+    assert any(k.endswith('test_record_iterable_without_tag-1') for k in keys)
+    assert any(k.endswith('test_record_iterable_without_tag-2') for k in keys)