Source code for pytorch_lightning.profiler.profiler
# Copyright The PyTorch Lightning team.## 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."""Profiler to check if there are any bottlenecks in your code."""importloggingimportosfromabcimportABC,abstractmethodfromcontextlibimportcontextmanagerfrompathlibimportPathfromtypingimportAny,Callable,Dict,Generator,Iterable,Optional,TextIO,Unionfrompytorch_lightning.utilities.cloud_ioimportget_filesystemfrompytorch_lightning.utilities.rank_zeroimportrank_zero_deprecationlog=logging.getLogger(__name__)
[docs]classProfiler(ABC):"""If you wish to write a custom profiler, you should inherit from this class."""def__init__(self,dirpath:Optional[Union[str,Path]]=None,filename:Optional[str]=None,)->None:self.dirpath=dirpathself.filename=filenameself._output_file:Optional[TextIO]=Noneself._write_stream:Optional[Callable]=Noneself._local_rank:Optional[int]=Noneself._stage:Optional[str]=None
[docs]@abstractmethoddefstart(self,action_name:str)->None:"""Defines how to start recording an action."""
[docs]@abstractmethoddefstop(self,action_name:str)->None:"""Defines how to record the duration once an action is complete."""
defsummary(self)->str:return""
[docs]@contextmanagerdefprofile(self,action_name:str)->Generator:"""Yields a context manager to encapsulate the scope of a profiled action. Example:: with self.profile('load training data'): # load training data code The profiler will start once you've entered the context and will automatically stop once you exit the code block. """try:self.start(action_name)yieldaction_namefinally:self.stop(action_name)
[docs]defprofile_iterable(self,iterable:Iterable,action_name:str)->Generator:"""Profiles over each value of an iterable. See deprecation message below. .. deprecated:: v1.6 `Profiler.profile_iterable` is deprecated in v1.6 and will be removed in v1.8. """rank_zero_deprecation(f"`{self.__class__.__name__}.profile_iterable` is deprecated in v1.6 and will be removed in v1.8.")iterator=iter(iterable)whileTrue:try:self.start(action_name)value=next(iterator)self.stop(action_name)yieldvalueexceptStopIteration:self.stop(action_name)break
[docs]defdescribe(self)->None:"""Logs a profile report after the conclusion of run."""# users might call `describe` directly as the profilers can be used by themselves.# to allow this, we open and close the files within this function by calling `_prepare_streams` and `teardown`# manually instead of letting the `Trainer` do it through `setup` and `teardown`self._prepare_streams()summary=self.summary()ifsummaryandself._write_streamisnotNone:self._write_stream(summary)ifself._output_fileisnotNone:self._output_file.flush()self.teardown(stage=self._stage)
[docs]defteardown(self,stage:Optional[str]=None)->None:"""Execute arbitrary post-profiling tear-down steps. Closes the currently open file and stream. """self._write_stream=Noneifself._output_fileisnotNone:self._output_file.close()self._output_file=None# can't pickle TextIOWrapper
To analyze traffic and optimize your experience, we serve cookies on this site. By clicking or navigating, you agree to allow our usage of cookies. Read PyTorch Lightning's Privacy Policy.