# -*- coding: utf-8 -*-
# MegEngine is Licensed under the Apache License, Version 2.0 (the "License")
#
# Copyright (c) 2014-2021 Megvii Inc. All rights reserved.
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT ARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
import json
from contextlib import contextmanager
from typing import List
from ..core._imperative_rt.core2 import (
pop_scope,
push_scope,
start_profile,
stop_profile,
sync,
)
[文档]class Profiler:
r"""
Profile graph execution in imperative mode.
:type path: Optional[str]
:param path: default path prefix for profiler to dump.
Examples:
.. code-block::
import megengine as mge
import megengine.module as M
from megengine.utils.profiler import Profiler
# With Learnable Parameters
for iter in range(0, 10):
# Only profile record of last iter would be saved
with Profiler("profile"):
# your code here
# Then open the profile file in chrome timeline window
"""
CHROME_TIMELINE = "chrome_timeline.json"
COMMAND = 1 << 0
OPERATOR = 1 << 1
TENSOR_LIFETIME = 1 << 2
TENSOR_PROP = 1 << 3
SYNC = 1 << 4
SCOPE = 1 << 5
ALL = (1 << 6) - 1
def __init__(
self,
path: str = "profile",
format: str = CHROME_TIMELINE,
*,
topic=OPERATOR | SCOPE,
align_time=True,
show_operator_name=True
) -> None:
self._path = path
self._format = format
self._options = {
"topic": int(topic),
"align_time": int(align_time),
"show_operator_name": int(show_operator_name),
}
def __enter__(self):
start_profile(self._options)
return self
def __exit__(self, val, tp, trace):
stop_profile(self._path, self._format)
# dump is async, so it's necessary to sync interpreter
sync()
def __call__(self, func):
def wrapper(*args, **kwargs):
with self:
return func(*args, **kwargs)
return wrapper
@contextmanager
def scope(name):
push_scope(name)
yield
pop_scope(name)
profile = Profiler
def merge_trace_events(sources: List[str], target: str):
names = list(map(lambda x: x + ".chrome_timeline.json", sources))
result = []
for name in names:
with open(name, "r", encoding="utf-8") as f:
content = json.load(f)
for entry in content:
result.append(entry)
with open(target + ".chrome_timeline.json", "w") as f:
json.dump(result, f, ensure_ascii=False, indent=4)