diff --git a/monai/utils/__init__.py b/monai/utils/__init__.py index 8f27a7faf3..d2d3e41d67 100644 --- a/monai/utils/__init__.py +++ b/monai/utils/__init__.py @@ -15,3 +15,4 @@ from .enums import * from .misc import * from .module import * +from .profiling import * diff --git a/monai/utils/decorators.py b/monai/utils/decorators.py index 0d14eb0dd3..35a594d077 100644 --- a/monai/utils/decorators.py +++ b/monai/utils/decorators.py @@ -9,28 +9,9 @@ # See the License for the specific language governing permissions and # limitations under the License. -import time from functools import wraps -def timing(func): - """ - This simple timing function decorator prints to stdout/logfile (it uses printFlush) how many seconds a call to the - original function took to execute, as well as the name before and after the call. - """ - - @wraps(func) - def timingwrap(*args, **kwargs): - print(func.__name__, flush=True) - start = time.perf_counter() - res = func(*args, **kwargs) - end = time.perf_counter() - print(func.__name__, "dT (s) =", (end - start), flush=True) - return res - - return timingwrap - - class RestartGenerator: """ Wraps a generator callable which will be called whenever this class is iterated and its result returned. This is diff --git a/monai/utils/profiling.py b/monai/utils/profiling.py new file mode 100644 index 0000000000..a8d6e2c630 --- /dev/null +++ b/monai/utils/profiling.py @@ -0,0 +1,90 @@ +# Copyright 2020 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. + +import time +from functools import wraps + +import torch + +__all__ = ["torch_profiler_full", "torch_profiler_time_cpu_gpu", "torch_profiler_time_end_to_end"] + + +def torch_profiler_full(func): + """ + A decorator which will run the torch profiler for the decorated function, + printing the results in full. + Note: Enforces a gpu sync point which could slow down pipelines. + """ + + @wraps(func) + def wrapper(*args, **kwargs): + + with torch.autograd.profiler.profile(use_cuda=True) as prof: + result = func(*args, **kwargs) + + print(prof, flush=True) + + return result + + return wrapper + + +def torch_profiler_time_cpu_gpu(func): + """ + A decorator which measures the execution time of both the CPU and GPU components + of the decorated function, printing both results. + Note: Enforces a gpu sync point which could slow down pipelines. + """ + + @wraps(func) + def wrapper(*args, **kwargs): + + with torch.autograd.profiler.profile(use_cuda=True) as prof: + result = func(*args, **kwargs) + + cpu_time = prof.self_cpu_time_total + gpu_time = sum(evt.self_cuda_time_total for evt in prof.function_events) + + cpu_time = torch.autograd.profiler.format_time(cpu_time) + gpu_time = torch.autograd.profiler.format_time(gpu_time) + + print("cpu time: {}, gpu time: {}".format(cpu_time, gpu_time), flush=True) + + return result + + return wrapper + + +def torch_profiler_time_end_to_end(func): + """ + A decorator which measures the total execution time from when the decorated + function is called to when the last cuda operation finishes, printing the result. + Note: Enforces a gpu sync point which could slow down pipelines. + """ + + @wraps(func) + def wrapper(*args, **kwargs): + + torch.cuda.synchronize() + start = time.perf_counter() + + result = func(*args, **kwargs) + + torch.cuda.synchronize() + end = time.perf_counter() + + total_time = (end - start) * 1e6 + total_time_str = torch.autograd.profiler.format_time(total_time) + print("end to end time: {}".format(total_time_str), flush=True) + + return result + + return wrapper