mirror of https://github.com/hpcaitech/ColossalAI
41 lines
983 B
Python
41 lines
983 B
Python
import torch
|
|
import torch.fx
|
|
from torch.fx import GraphModule
|
|
|
|
from .passes import ShapeProp, graph_profile_pass, shape_prop_pass
|
|
from .passes.graph_profile import FlopProfiler
|
|
|
|
|
|
def register_flop_count_impl(func):
|
|
|
|
def wrapper(impl):
|
|
FlopProfiler._custom_flop_count_impl[func] = impl
|
|
return impl
|
|
|
|
return wrapper
|
|
|
|
|
|
def register_shape_impl(func):
|
|
|
|
def wrapper(impl):
|
|
ShapeProp._custom_dispatch_func[func] = impl
|
|
return impl
|
|
|
|
return wrapper
|
|
|
|
|
|
def symbolic_profile(module: GraphModule, *args, verbose=False) -> GraphModule:
|
|
"""Symbolically profile a model with sample inputs.
|
|
|
|
Args:
|
|
module (GraphModule): The module to be profiled
|
|
args (Tuple): The sample inputs
|
|
verbose (bool): Whether to print the profiling result
|
|
|
|
Returns:
|
|
GraphModule: The profiled module
|
|
"""
|
|
module = shape_prop_pass(module, *args)
|
|
module = graph_profile_pass(module, *args, verbose=verbose)
|
|
return module
|