适配PyTorch FX，OneFlow，让量化感知训练更简单

import oneflow
class MyModule(oneflow.nn.Module):    def __init__(self):        super().__init__()        self.linear = oneflow.nn.Linear(512, 512)
    def forward(self, x):        x = self.linear(x)        y = oneflow.ones([2, 3])
        x = oneflow.relu(x)        return y
m = MyModule()
traced = oneflow.fx.symbolic_trace(m)print(traced.code)"""def forward(self, x):    linear = self.linear(x);  x = None    relu = oneflow.relu(linear);  linear = None    _tensor_constant0 = self._tensor_constant0    return _tensor_constant0"""

def _create_wrapped_func(orig_fn):    @functools.wraps(orig_fn)    def wrapped(*args, **kwargs):        # 判断参数中是否存在proxy变量        proxy = _find_proxy(args, kwargs)        if proxy is not None:            # 如果参数中有Proxy变量，创建节点并返回Proxy包装后的节点            return_proxy = proxy.tracer.create_proxy(                "call_function", orig_fn, args, kwargs            )            return_proxy.node.meta["is_wrapped"] = True            return return_proxy        # 如果没有Proxy变量，直接调用原函数        return orig_fn(*args, **kwargs)
    return wrapped

def create_proxy(self, kind: str, target: Target, args: Tuple[Any, ...], kwargs: Dict[str, Any],                     name: Optional[str] = None, type_expr : Optional[Any] = None,                     proxy_factory_fn: Callable[[Node], 'Proxy'] = None):    args_ = self.create_arg(args)    kwargs_ = self.create_arg(kwargs)    assert isinstance(args_, tuple)    assert isinstance(kwargs_, dict)
    # 创建节点    node = self.create_node(kind, target, args_, kwargs_, name, type_expr)
    if not proxy_factory_fn:        proxy = self.proxy(node)    else:        proxy = proxy_factory_fn(node)
    return proxy

def create_node(self, op: str, target: 'Target',                    args: Optional[Tuple['Argument', ...]] = None,                    kwargs: Optional[Dict[str, 'Argument']] = None,                    name: Optional[str] = None,                    type_expr: Optional[Any] = None) -> Node:    # 此处有一些assert
    # 创建一个节点名称，避免重复    candidate = name if name is not None else self._target_to_str(target)    name = self._graph_namespace.create_name(candidate, None)    # 创建节点    n = Node(self, name, op, target, args, kwargs, type_expr)
    # 建立名称与节点的映射关系    self._graph_namespace.associate_name_with_obj(name, n)
    return n

def trace(        self,        root: Union[oneflow.nn.Module, Callable[..., Any]],        concrete_args: Optional[Dict[str, Any]] = None,    ) -> Graph:    # 确定模块主体以及forward函数，其中fn即forward函数    if isinstance(root, oneflow.nn.Module):        self.root = root
        assert hasattr(            type(root), self.traced_func_name        ), f"traced_func_name={self.traced_func_name} doesn't exist in {type(root).__name__}"
        fn = getattr(type(root), self.traced_func_name)        self.submodule_paths = {mod: name for name, mod in root.named_modules()}    else:        self.root = oneflow.nn.Module()        fn = root
    tracer_cls: Optional[Type["Tracer"]] = getattr(self, "__class__", None)    # 在Tracer中初始化一张图    self.graph = Graph(tracer_cls=tracer_cls)        self.tensor_attrs: Dict[oneflow.Tensor, str] = {}    # 这个子函数用于收集模型中所有Tensor类型的变量    def collect_tensor_attrs(m: oneflow.nn.Module, prefix_atoms: List[str]):        for k, v in m.__dict__.items():            if isinstance(v, oneflow.Tensor):                self.tensor_attrs[v] = ".".join(prefix_atoms + [k])        for k, v in m.named_children():            collect_tensor_attrs(v, prefix_atoms + [k])
    collect_tensor_attrs(self.root, [])
    assert isinstance(fn, FunctionType)
    # 获取fn所在模块的所有可读变量    fn_globals = fn.__globals__    # 创建包装后的参数    fn, args = self.create_args_for_root(        fn, isinstance(root, oneflow.nn.Module), concrete_args    )

# 下面代码仍然是`trace`函数的一部分
# 定义对于`nn.Module`的getattr方法的包装@functools.wraps(_orig_module_getattr)def module_getattr_wrapper(mod, attr):    attr_val = _orig_module_getattr(mod, attr)    return self.getattr(attr, attr_val, parameter_proxy_cache)
# 定义对于`nn.Module`的forward方法的包装@functools.wraps(_orig_module_call)def module_call_wrapper(mod, *args, **kwargs):    def forward(*args, **kwargs):        return _orig_module_call(mod, *args, **kwargs)
    _autowrap_check(        patcher,        getattr(getattr(mod, "forward", mod), "__globals__", {}),        self._autowrap_function_ids,    )    return self.call_module(mod, forward, args, kwargs)# 这里Patcher的作用是在退出这一环境的时候恢复现场，避免包装函数、方法的影响溢出到`trace`之外。with _Patcher() as patcher:    # 对`__getattr__`和`nn.Module.__call__`这两个方法默认进行包装    patcher.patch_method(        oneflow.nn.Module,        "__getattr__",        module_getattr_wrapper,        deduplicate=False,    )    patcher.patch_method(        oneflow.nn.Module, "__call__", module_call_wrapper, deduplicate=False    )    # 对预定好需要进行包装的函数进行包装    _patch_wrapped_functions(patcher)    _autowrap_check(patcher, fn_globals, self._autowrap_function_ids)    # 遍历所有需要对其中函数进行自动包装的package    for module in self._autowrap_search:        if module is oneflow:            dict = {}            # 当package为oneflow时，对此进行特殊处理，单独分出一个字典存放原本`oneflow.__dict__`中的内容            for name, value in module.__dict__.items():                if not isinstance(value, oneflow.nn.Module) and not value in _oneflow_no_wrapped_functions:                    dict[name] = value            _autowrap_check_oneflow(                patcher, dict, module.__dict__, self._autowrap_function_ids            )        else:            _autowrap_check(                patcher, module.__dict__, self._autowrap_function_ids            )    # 创建节点，这里的`create_node`调用实际上只是创建了最后一个节点，即输出节点。    # 但是这里`fn`就是forward函数，在运行这一函数的时候，就会如前面所说依次创建节点。    self.create_node(        "output",        "output",        (self.create_arg(fn(*args)),),        {},        type_expr=fn.__annotations__.get("return", None),    )

def _patch_wrapped_functions(patcher: _Patcher):    # `_wrapped_fns_to_patch`中包含了所有需要自动包装的函数    for frame_dict, name in _wrapped_fns_to_patch:        if name not in frame_dict:            if hasattr(builtins, name):                # 对于built-in函数，不存在于frame_dict中，单独进行处理来根据名称获取函数本身                orig_fn = getattr(builtins, name)            else:                # 如果是oneflow中指定需要包装的函数，那么就进行获取，否则抛出名称无法识别的异常                is_oneflow_wrapped_function, func = is_oneflow_wrapped_function_and_try_get(name)                if is_oneflow_wrapped_function:                    orig_fn = func                else:                    raise NameError("Cannot deal with the function %s."%name)        else:            # 如果函数名称已经存在于frame_dict中，直接通过字典查询来获得函数            orig_fn = frame_dict[name]        # 创建包装后的函数并进行`patch`，即定义当trace过程结束的时候，如何还原现场        patcher.patch(frame_dict, name, _create_wrapped_func(orig_fn))        # 对于类中的方法，直接包装并patch。    for cls, name in _wrapped_methods_to_patch:        patcher.patch_method(cls, name, _create_wrapped_method(cls, name))

import oneflow
oneflow.fx.wrap(len)class MyModule(oneflow.nn.Module):    def __init__(self):        super().__init__()        self.linear = oneflow.nn.Linear(512, 512)
    def forward(self, x):        x = self.linear(x) + len(x.shape)        return x
traced = oneflow.fx.symbolic_trace(MyModule())print(traced.code)"""def forward(self, x):    linear = self.linear(x)    getattr_1 = x.shape;  x = None    len_1 = len(getattr_1);  getattr_1 = None    add = linear + len_1;  linear = len_1 = None    return add"""

# MyModule来自其它包with oneflow.fx.global_wrap(len):    m = MyModule()
    traced = oneflow.fx.symbolic_trace(m)    print(traced.code)    """    def forward(self, x):        linear = self.linear(x);  x = None        getattr_1 = linear.shape        len_1 = len(getattr_1);  getattr_1 = None        relu = oneflow.relu(linear);  linear = None        add = relu + len_1;  relu = len_1 = None        return add    """

import oneflow
class MyModule(oneflow.nn.Module):    def __init__(self, do_activation : bool = False):        super().__init__()        self.do_activation = do_activation        self.linear = oneflow.nn.Linear(512, 512)
    def forward(self, x):        x = self.linear(x)        y = oneflow.ones([2, 3])
        x = oneflow.topk(x, 10)        return x.relu() + y
traced = oneflow.fx.symbolic_trace(MyModule())traced.graph.print_tabular()
"""opcode         name               target                    args                       kwargs-------------  -----------------  ------------------------  -------------------------  --------placeholder    x                  x                         ()                         {}call_module    linear             linear                    (x,)                       {}call_function  topk               <built-in function topk>  (linear, 10)               {}call_method    relu               relu                      (topk,)                    {}get_attr       _tensor_constant0  _tensor_constant0         ()                         {}call_function  add                <built-in function add>   (relu, _tensor_constant0)  {}output         output             output                    (add,)                     {}"""

import oneflowfrom oneflow.fx import symbolic_traceimport operator
class M(oneflow.nn.Module):    def forward(self, x, y):        return x + y, oneflow.add(x, y), x.add(y)
if __name__ == '__main__':    traced = symbolic_trace(M())
    patterns = set([operator.add, oneflow.add, "add"])
    for n in traced.graph.nodes:        if any(n.target == pattern for pattern in patterns):            with traced.graph.inserting_after(n):                new_node = traced.graph.call_function(oneflow.mul, n.args, n.kwargs)                n.replace_all_uses_with(new_node)            traced.graph.erase_node(n)
    traced.recompile()
    traced.graph.print_tabular()
    print(traced.code)

import oneflowimport flowvision.models as modelsimport statistics, tabulate, timefrom typing import Any, Dict, List
class ProfilingInterpreter(oneflow.fx.Interpreter):    def __init__(self, mod : oneflow.nn.Module):        gm = oneflow.fx.symbolic_trace(mod)        super().__init__(gm)
        # 记录总运行时间        self.total_runtime_sec : List[float] = []        # 记录各个节点运行时间        self.runtimes_sec : Dict[oneflow.fx.Node, List[float]] = {}
    # 重写`run`方法，本质上是对基类`run`方法的简单封装，在运行前后记录时间点。    # 这一方法是Graph整体运行的入口。    def run(self, *args) -> Any:        t_start = time.time()        return_val = super().run(*args)        t_end = time.time()        self.total_runtime_sec.append(t_end - t_start)        return return_val
    # 同上，重写`run_node`方法，不需要自己写细节实现，只需要在对基类的`run_node`调用前后记录时间点即可    # 这一方法是Graph中运行每个Node的入口。    def run_node(self, n : oneflow.fx.Node) -> Any:        t_start = time.time()        return_val = super().run_node(n)        t_end = time.time()        self.runtimes_sec.setdefault(n, [])        self.runtimes_sec[n].append(t_end - t_start)        return return_val
    # 定义如何打印性能测试结果    def summary(self, should_sort : bool = False) -> str:        # 存储每个节点的打印信息        node_summaries : List[List[Any]] = []        # 由于模块会被调用多次，所以这里计算一下平均的运行总时长        mean_total_runtime = statistics.mean(self.total_runtime_sec)
        for node, runtimes in self.runtimes_sec.items():            mean_runtime = statistics.mean(runtimes)            # 计算节点运行时间占总时间的比例            pct_total = mean_runtime / mean_total_runtime * 100            # 记录节点信息、节点平均运行时长和节点运行时间占总时间的比例            node_summaries.append(                [node.op, str(node), mean_runtime, pct_total])
        # 如果需要，安按照运行时间进行排序        if should_sort:            node_summaries.sort(key=lambda s: s[2], reverse=True)
        # 以下是借助tabulate库进行格式化来美化显示效果        headers : List[str] = [            'Op type', 'Op', 'Average runtime (s)', 'Pct total runtime'        ]        return tabulate.tabulate(node_summaries, headers=headers)

if __name__ == '__main__':    rn18 = models.resnet18()    rn18.eval()    input = oneflow.randn(5, 3, 224, 224)    output = rn18(input)    interp = ProfilingInterpreter(rn18)    interp.run(input)    print(interp.summary(True))

import sysimport oneflowimport oneflow.nn as nnimport numpy as npimport copyfrom typing import Dict, Any, Tuple
# 通过直接对权重进行运算的方式进行Conv和BN的融合def fuse_conv_bn_eval(conv, bn):    assert(not (conv.training or bn.training)), "Fusion only for eval!"    fused_conv = copy.deepcopy(conv)
    fused_conv.weight, fused_conv.bias = \        fuse_conv_bn_weights(fused_conv.weight, fused_conv.bias,                             bn.running_mean, bn.running_var, bn.eps, bn.weight, bn.bias)
    return fused_conv
# 权重融合方式def fuse_conv_bn_weights(conv_w, conv_b, bn_rm, bn_rv, bn_eps, bn_w, bn_b):    if conv_b is None:        conv_b = oneflow.zeros_like(bn_rm)    if bn_w is None:        bn_w = oneflow.ones_like(bn_rm)    if bn_b is None:        bn_b = oneflow.zeros_like(bn_rm)    bn_var_rsqrt = oneflow.rsqrt(bn_rv + bn_eps)
    conv_w = conv_w * (bn_w * bn_var_rsqrt).reshape([-1] + [1] * (len(conv_w.shape) - 1))    conv_b = (conv_b - bn_rm) * bn_var_rsqrt * bn_w + bn_b
    return oneflow.nn.Parameter(conv_w), oneflow.nn.Parameter(conv_b)
# 根据字符串对名称进行分割，比如`foo.bar.baz` -> (`foo.bar`, `baz`)def _parent_name(target : str) -> Tuple[str, str]:    *parent, name = target.rsplit('.', 1)    return parent[0] if parent else '', name
def replace_node_module(node: oneflow.fx.Node, modules: Dict[str, Any], new_module: oneflow.nn.Module):    assert(isinstance(node.target, str))    parent_name, name = _parent_name(node.target)    setattr(modules[parent_name], name, new_module)
# 定义对模型进行融合操作的过程def fuse(model: oneflow.nn.Module) -> oneflow.nn.Module:    model = copy.deepcopy(model)    # 先通过fx.symbolic_trace获取一个GraphModule    fx_model: oneflow.fx.GraphModule = oneflow.fx.symbolic_trace(model)    modules = dict(fx_model.named_modules())
    # 遍历GraphModule中的所有节点，分别进行操作    for node in fx_model.graph.nodes:        # 跳过所有不是module的节点        if node.op != 'call_module':            continue        # 检测到conv+bn的结构后进行融合操作        if type(modules[node.target]) is nn.BatchNorm2d and type(modules[node.args[0].target]) is nn.Conv2d:            # conv的输出同时被其它节点使用，即conv后连接两个节点时无法融合            if len(node.args[0].users) > 1:                continue            conv = modules[node.args[0].target]            bn = modules[node.target]            fused_conv = fuse_conv_bn_eval(conv, bn)            replace_node_module(node.args[0], modules, fused_conv)            # 对图中的边进行置换，对于用到bn输出的节点，要更改它们的输入            node.replace_all_uses_with(node.args[0])            # 移除旧的节点            fx_model.graph.erase_node(node)    fx_model.graph.lint()    # 重新建图（构造模型）    fx_model.recompile()    return fx_model

if __name__ == '__main__':    # 以下引入flowvision中的resnet 18模型，并进行融合前后的benchmark比较    import flowvision.models as models    import time
    rn18 = models.resnet18().cuda()    rn18.eval()
    inp = oneflow.randn(10, 3, 224, 224).cuda()    output = rn18(inp)
    def benchmark(model, iters=20):        for _ in range(10):            model(inp)        oneflow.cuda.synchronize()        begin = time.time()        for _ in range(iters):            model(inp)        return str(time.time()-begin)
    fused_rn18 = fuse(rn18)    unfused_time = benchmark(rn18)    fused_time = benchmark(fused_rn18)    print("Unfused time: ", benchmark(rn18))    print("Fused time: ", benchmark(fused_rn18))    assert unfused_time > fused_time