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(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_pytorch_profiler_trainer_ddp(tmpdir, use_output_filename):
"""Ensure that the profiler can be given to the training and default step are properly recorded. """
if use_output_filename:
output_filename = os.path.join(tmpdir, "profiler.txt")
else:
output_filename = None
profiler = PyTorchProfiler(output_filename=output_filename)
model = BoringModel()
trainer = Trainer(
fast_dev_run=True,
profiler=profiler,
accelerator="ddp",
gpus=2,
)
trainer.fit(model)
enabled = use_output_filename or not use_output_filename and profiler.local_rank == 0
if enabled:
assert len(profiler.summary()) > 0
assert set(profiler.profiled_actions.keys()) == {
'training_step_and_backward', 'validation_step'
}
else:
assert profiler.summary() is None
assert set(profiler.profiled_actions.keys()) == set()
if use_output_filename:
profiler.describe()
data = Path(profiler.output_fname).read_text()
assert len(data) > 0
def test_pytorch_profiler_deepcopy(tmpdir):
pytorch_profiler = PyTorchProfiler(dirpath=tmpdir,
filename="profiler",
schedule=None)
pytorch_profiler.start("on_train_start")
torch.tensor(1)
pytorch_profiler.describe()
assert deepcopy(pytorch_profiler)
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())
