def test_pytorch_profiler_nested(tmpdir):
"""Ensure that the profiler handles nested context"""
pytorch_profiler = PyTorchProfiler(record_functions={"a", "b", "c"},
use_cuda=False,
dirpath=tmpdir,
filename="profiler",
schedule=None)
with pytorch_profiler.profile("a"):
a = torch.ones(42)
with pytorch_profiler.profile("b"):
b = torch.zeros(42)
with pytorch_profiler.profile("c"):
_ = a + b
pytorch_profiler.describe()
events_name = {e.name for e in pytorch_profiler.function_events}
names = {"a", "b", "c"}
ops = {"add", "empty", "fill_", "ones", "zero_", "zeros"}
if _TORCH_GREATER_EQUAL_1_7:
ops = {"aten::" + op for op in ops}
expected = names.union(ops)
assert events_name == expected, (events_name, torch.__version__,
platform.system())
def test_register_record_function(tmpdir):
use_cuda = torch.cuda.is_available()
pytorch_profiler = PyTorchProfiler(
export_to_chrome=False,
use_cuda=use_cuda,
dirpath=tmpdir,
filename="profiler",
schedule=None,
on_trace_ready=None,
)
class TestModel(BoringModel):
def __init__(self):
super().__init__()
self.layer = torch.nn.Sequential(torch.nn.Linear(1, 1), torch.nn.ReLU(), torch.nn.Linear(1, 1))
model = TestModel()
input = torch.rand((1, 1))
if use_cuda:
model = model.cuda()
input = input.cuda()
with pytorch_profiler.profile("a"):
with RegisterRecordFunction(model):
model(input)
pytorch_profiler.describe()
event_names = [e.name for e in pytorch_profiler.function_events]
assert "[pl][module]torch.nn.modules.container.Sequential: layer" in event_names
assert "[pl][module]torch.nn.modules.linear.Linear: layer.0" in event_names
assert "[pl][module]torch.nn.modules.activation.ReLU: layer.1" in event_names
assert "[pl][module]torch.nn.modules.linear.Linear: layer.2" in event_names
def test_pytorch_profiler_nested(tmpdir):
"""Ensure that the profiler handles nested context"""
pytorch_profiler = PyTorchProfiler(profiled_functions=["a", "b", "c"],
use_cuda=False,
output_filename=os.path.join(
tmpdir, "profiler.txt"))
with pytorch_profiler.profile("a"):
a = torch.ones(42)
with pytorch_profiler.profile("b"):
b = torch.zeros(42)
with pytorch_profiler.profile("c"):
_ = a + b
pa = pytorch_profiler.profiled_actions
# From PyTorch 1.8.0, less operation are being traced.
if LooseVersion(torch.__version__) >= LooseVersion("1.8.0"):
expected_ = {
'a': ['ones', 'empty', 'fill_', 'zeros', 'empty', 'zero_', 'add'],
'b': ['zeros', 'empty', 'zero_'],
'c': ['add'],
}
# From PyTorch 1.6.0, more operation are being traced.
elif LooseVersion(torch.__version__) >= LooseVersion("1.6.0"):
expected_ = {
'a': [
'ones', 'empty', 'fill_', 'zeros', 'empty', 'zero_', 'fill_',
'add', 'empty'
],
'b': ['zeros', 'empty', 'zero_', 'fill_'],
'c': ['add', 'empty'],
}
else:
expected_ = {
'a': ['add'],
'b': [],
'c': ['add'],
}
for n in ('a', 'b', 'c'):
pa[n] = [e.name for e in pa[n]]
if LooseVersion(torch.__version__) >= LooseVersion("1.7.1"):
pa[n] = [e.replace("aten::", "") for e in pa[n]]
assert pa[n] == expected_[n]
def test_pytorch_profiler_nested(tmpdir):
"""Ensure that the profiler handles nested context"""
pytorch_profiler = PyTorchProfiler(record_functions={"a", "b", "c"},
use_cuda=False,
dirpath=tmpdir,
filename="profiler",
schedule=None)
with pytorch_profiler.profile("a"):
a = torch.ones(42)
with pytorch_profiler.profile("b"):
b = torch.zeros(42)
with pytorch_profiler.profile("c"):
_ = a + b
pytorch_profiler.describe()
events_name = {e.name for e in pytorch_profiler.function_events}
if platform.system() == "Windows":
expected = {
'a', 'add', 'b', 'c', 'profiler::_record_function_enter',
'profiler::_record_function_exit'
}
else:
expected = {
'signed char', 'add', 'profiler::_record_function_exit', 'bool',
'char', 'profiler::_record_function_enter'
}
if Version(torch.__version__) >= Version("1.6.0"):
expected = {
'add', 'zeros', 'ones', 'zero_', 'b', 'fill_', 'c', 'a', 'empty'
}
if Version(torch.__version__) >= Version("1.7.0"):
expected = {
'aten::zeros', 'aten::add', 'aten::zero_', 'c', 'b', 'a',
'aten::fill_', 'aten::empty', 'aten::ones'
}
assert events_name == expected, (events_name, torch.__version__,
platform.system())
