OneFlow: Python 端构图

前言

在写这篇之前，已经写过一篇博客 [1] 分析了如何获取前端的表达式，不过我认为写的不够清晰，代码没怎么贴，后面再看发现只能看当时的思考，而不是源代码，所以就重新写一遍了。这篇文章会更加专注于 Python 端是如何构图的，更加关注算子的内部的执行过程。

带着如下几个问题去看代码：

• 什么时候获取前端表达式？
• 获取的前端表达式是什么样的格式的？
• 对于一个 InferenceSession，读取本地模型，逐个逐个将 Op 添加到 OneFlow 当中。对于一个 Job 函数，它是如何将 Op 添加到 OneFlow 当中的呢？

流程分析

首先先走一遍流程，这里截取了 [1] 的流程部分，顺便补上代码，加上自己的一些点评。

• Python 逐行执行，执行到了 @flow.global_function()。

Python 解释器运行的时候，会逐行读取代码并执行。文末有附代码。

# 文档中的例子 lenet 
@flow.global_function(type="train")
def train_job(
    images: tp.Numpy.Placeholder((BATCH_SIZE, 1, 28, 28), dtype=flow.float),
    labels: tp.Numpy.Placeholder((BATCH_SIZE,), dtype=flow.int32),
) -> tp.Numpy:
    with flow.scope.placement("gpu", "0:0"):
        logits = lenet(images, train=True)
        loss = flow.nn.sparse_softmax_cross_entropy_with_logits(
            labels, logits, name="softmax_loss"
        )
    lr_scheduler = flow.optimizer.PiecewiseConstantScheduler([], [0.1])
    flow.optimizer.SGD(lr_scheduler, momentum=0).minimize(loss)
    return loss

• 跳进 api_oneflow_function，获取一个函数，这个函数是装饰器。根据静态图或者动态图，返回一个装饰器。

# python/oneflow/compatible/single_client/framework/function_util.py: 85
def api_oneflow_function(
    type: str = "predict", function_config: FunctionConfig = None
) -> Callable[[Callable], Callable]:
    """Creates a callable OneFlow global function from a Python function.

    For instance::

        @oneflow.compatible.single_client.global_function(flow.FunctionConfig())
        def train():
            # your model

    Args:
        function_config (FunctionConfig, optional): a `FunctionConfig` object. Defaults to FunctionConfig().

    Returns:
        Callable[[Callable], Callable]: a callable which is called to execute the compiled function
    """
    if isinstance(type, FunctionConfig):
        function_config = type
        print(
            "WARNING: flow.global_function(func_config) is deprecated. Please replace it with flow.global_function(type, func_config).
            "
        )
        print(traceback.format_stack()[-2])
    else:
        assert type in ["train", "predict"]
        if function_config is None:
            function_config = FunctionConfig()
        if type == "train":
            function_config.function_desc.job_config_proto.mutable_train_conf()
        else:
            function_config.function_desc.job_config_proto.mutable_predict_conf()
    api = enable_if.unique([eager_oneflow_function, lazy_oneflow_function])
    return api(function_config)

• 上面代码获取到的是 lazy_oneflow_function。这里会创建一个闭包，闭包里面包含有运行这个 Job 的 Session 等信息，以及用户定义的 Job 函数。

# python/oneflow/compatible/single_client/framework/function_util.py: 143
@enable_if.condition(
    hob.in_normal_mode & ~hob.eager_execution_enabled & ~hob.session_initialized
)
def lazy_oneflow_function(function_config=FunctionConfig()):
    assert isinstance(function_config, FunctionConfig)

    def Decorator(job_func):
        if not hasattr(job_func, "__oneflow_function_signature__"):
            job_func.__oneflow_function_signature__ = inspect.signature(job_func)
        oft_util.CheckGlobalFunctionAnnotation(job_func.__oneflow_function_signature__)
        sess = session_ctx.GetDefaultSession()

        @functools.wraps(job_func)
        def Func(*args, **kwargs):
            return _RunLazyJob(sess, job_func, *args, **kwargs)

        sess.AddJob(_CloneFunctionDesc(function_config.function_desc, job_func))
        for x in dir(job_func):
            if x.startswith("__oneflow_"):
                setattr(Func, x, getattr(job_func, x))
        return Func

    return Decorator

• 经过装饰器之后，被装饰的函数最终会执行的方法是 _RunLazyJob。如果用户第一次调用 Job 函数，会进行 Session 的初始化。初始化结束之后，会 LazyRun。被 @flow.global_function() 装饰的函数，第一次调用的时候，执行初始化 TryInit，后续调用不再重复初始化。

# python/oneflow/compatible/single_client/framework/function_util.py: 225
def _RunLazyJob(session, job_func, *args, **kwargs):
    return session.TryInit().LazyRun(job_func, *args, **kwargs)

• Session 初始化的时候，进行编译。调用 compiler 的 Compile 方法，对 Job 进行编译。

# python/oneflow/compatible/single_client/framework/session_util.py： 183
def Init(self):
    assert self.status_ is SessionStatus.OPEN
    self.status_ = SessionStatus.RUNNING
    if not oneflow._oneflow_internal.IsEnvInited():
        flow.env.init()
    _TryCompleteConfigProto(self.config_proto)
    self.resource_ = self.config_proto.resource
    if not oneflow._oneflow_internal.EagerExecutionEnabled():
        c_api_util.InitLazyGlobalSession(self.config_proto)
        for (job_name, func_desc) in self.job_name2function_desc_.items():
            compiler.Compile(self, func_desc, self.config_proto)
            self.existed_module_names_ = set()
        self.job_name2var_name2var_blob_ = dict()
        assert len(self.job_name2function_desc_.items()) > 0
        oneflow._oneflow_internal.StartLazyGlobalSession()
        self.inter_user_job_info_ = c_api_util.GetInterUserJobInfo()
        self.UpdateInfo4InterfaceOp()
        if not config_util.api_legacy_model_io_enabled():
            check_point_v2.Init()
    else:
        self.eager_config_proto_ctx_ = oneflow._oneflow_internal.LogicalConfigProtoContext(
            str(self.config_proto)
        )
    return self

• 在编译之前需要设置好环境，可以看到 Compile 中的 with，指向了一个方法 InterpretScope。在 InterpretScope 的最后，主要的工作是：打开 JobBuildAndInferCtx，设置 JobBuildAndInferCtx，设置 Python 的运行模式，设置运行时的 scope。这些设置放在了 with 里面，在这一段结束的时候，还会进行一些清理的工作，比如关闭 JobBuildAndInferCtx 等。

# python/oneflow/compatible/single_client/framework/compiler.py: 44
def Compile(session, function_desc, config_proto):
    with InterpretScope(session, function_desc, config_proto):
        _CompileJob(session, function_desc)
        session.StashJob(function_desc.job_func.__name__)
        oneflow._oneflow_internal.CurJobBuildAndInferCtx_Complete()
        session.StashJob(
            function_desc.job_func.__name__,
            function_desc.job_func.__name__ + "_after_complete",
        )

# python/oneflow/compatible/single_client/framework/compiler.py: 63
@contextmanager
def InterpretScope(session, function_desc, config_proto):
    job_conf = function_desc.job_config_proto
    job_conf.set_job_name(function_desc.job_func.__name__)
    placement_scope = function_desc.function_attribute.default_placement_scope
    if placement_scope is None:
        tag_and_dev_ids = placement_util.GetDefaultMachineDeviceIds(session.resource)
        hierarchy = None
    else:
        assert isinstance(placement_scope, placement_ctx.EmptyPlacementScope)
        tag_and_dev_ids = (
            placement_scope.device_tag,
            placement_scope.machine_device_ids,
        )
        hierarchy = placement_scope.hierarchy
    distribute_strategy = function_desc.function_attribute.default_distribute_strategy
    if distribute_strategy is None:
        distribute_strategy = distribute_util.DistributeConsistentStrategy()
    is_mirrored = isinstance(
        distribute_strategy, distribute_util.DistributeMirroredStrategy
    )
    assert isinstance(hierarchy, (list, tuple)) or hierarchy is None
    if hierarchy is not None:
        hierarchy = oneflow._oneflow_internal.Size(tuple(hierarchy))
    scope = scope_util.MakeInitialScope(
        job_conf, *tag_and_dev_ids, hierarchy, is_mirrored
    )
    with _JobBuildAndInferCtx(job_conf.job_name()), distribute_strategy:
        c_api_util.CurJobBuildAndInferCtx_SetJobConf(job_conf)
        with runtime_mode.ModeScope(runtime_mode.GLOBAL_MODE):
            with scope_util.ScopeContext(scope):
                yield

# python/oneflow/compatible/single_client/framework/compiler.py: 145
@contextmanager
def _JobBuildAndInferCtx(job_name):
    c_api_util.JobBuildAndInferCtx_Open(job_name)
    try:
        yield
    finally:
        oneflow._oneflow_internal.JobBuildAndInferCtx_Close()

# python/oneflow/compatible/single_client/framework/scope_util.py: 105
@contextmanager
def ScopeContext(scope):
    old_scope = oneflow._oneflow_internal.GetCurrentScope()
    oneflow._oneflow_internal.GlobalScopeStackPush(scope)
    try:
        yield
    finally:
        assert oneflow._oneflow_internal.GetCurrentScope() is scope
        oneflow._oneflow_internal.GlobalScopeStackPop()
        assert oneflow._oneflow_internal.GetCurrentScope() is old_scope

• 接下来执行 _CompileJob，其中会执行用户定义的 Job 函数。我可以看到调用了 func 方法，这个实际上就是 Job 函数。

# python/oneflow/compatible/single_client/framework/compiler.py: 97
def _CompileJob(session, function_desc):
    func = function_desc.job_func
    parameters = func.__oneflow_function_signature__.parameters
    if len(parameters) == 0:
        func.__oneflow_input_blob_defs__ = ()
    elif all((p.annotation is inspect._empty for (_, p) in parameters.items())):
        func.__oneflow_input_blob_defs__ = _GetArgDefault(func)
    elif all((p.annotation is not inspect._empty for (_, p) in parameters.items())):
        func.__oneflow_input_blob_defs__ = _MakeInputBlobDefFromParameterSignature(
            parameters
        )
    else:
        raise NotImplementedError(
            "All parameters of global function should be annotated"
        )
    inputs = _RecursiveMakeInputBlobs(func.__oneflow_input_blob_defs__)
    ret = func(*inputs)
    return_annotation = func.__oneflow_function_signature__.return_annotation
    oft_util.CheckReturnByAnnotation(func.__name__, ret, return_annotation)
    func.__oneflow_output_remote_blobs__ = _RecursiveMakeRetRemoteBlobs(
        ret, allow_cpu_return_op=function_desc.function_attribute.allow_cpu_return_op
    )

• Job 函数内部执行的时候，就是到了具体的算子。每个算子的结尾都会去调用 user_op_builder.py 的方法。比较重要的一个方法是 InferAndTryRun，这个方法将会调用 CurJobAddOp 添加算子到图里。下面以 pad 算子为例子。Python 端的算子会检查用户输入是否合法，最后调用 user_op_builder 将算子加入到计算图中。

# python/oneflow/compatible/single_client/ops/pad.py: 229
def reflection_pad2d(
    x: oneflow._oneflow_internal.BlobDesc,
    padding: Union[int, tuple, list],
    name: Optional[str] = None,
) -> oneflow._oneflow_internal.BlobDesc:
    (H, W) = (x.shape[2], x.shape[3])
    if isinstance(padding, (tuple, list)):
        assert len(padding) == len(x.shape), ValueError(
            "padding boundry must be the same size of input dims"
        )
        assert (
            padding[2] < H and padding[3] < H and (padding[0] < W) and (padding[1] < W)
        ), ValueError(
            "Padding size should be less than the corresponding input dimension!"
        )
        boundry = [padding[0], padding[1], padding[2], padding[3]]
    elif isinstance(padding, int):
        assert padding < H and padding < W, ValueError(
            "Padding size should be less than the corresponding input dimension!"
        )
        boundry = [padding, padding, padding, padding]
    else:
        raise ValueError("padding must be in or list or tuple!")
    return (
        flow.user_op_builder(
            name if name is not None else id_util.UniqueStr("Reflection_Pad2d_")
        )
        .Op("reflection_pad2d")
        .Input("x", [x])
        .Output("y")
        .Attr("padding", list(boundry))
        .Build()
        .InferAndTryRun()
        .RemoteBlobList()[0]
    )

# python/oneflow/compatible/single_client/ops/user_op_builder.py: 163
class LazyUserOp(UserOp):
    def __init__(self, op_name, op_type_name):
        UserOp.__init__(self, op_name, op_type_name)

    def InferAndTryRun(self):
        compile_context.CurJobAddOp(self.op_conf_)
        return self

    def MakeRemoteBlob(self, lbi):
        return remote_blob_util.RemoteBlob(lbi)

• 之后再次执行这些函数的时候，会跳过 TryInit，执行 LazyRun。调用 LaunchUserJob 启动 Job 函数，进行前向传播。

# python/oneflow/compatible/single_client/framework/session_util.py: 275
def LazyRun(self, job_func, *arg):
    assert self.status_ is SessionStatus.RUNNING
    remote_blobs = self.LaunchUserJob(job_func, *arg)
    if remote_blobs is None:
        return
    future_blob = LazyFutureRemoteBlobs(self).SetResult(remote_blobs).Inited()
    annotation = inspect.signature(job_func).return_annotation
    return oft_util.TransformGlobalFunctionResult(future_blob, annotation)

计算图构建

流程分析完了，来分析一下用户的计算图是如何构建的？CurJobAddOp 的执行意味着什么？

在 InferenceSession 中，不需要构图，直接从文件系统读取一个计算图，然后对每个 Op 调用 CurJobAddOp。CurJobAddOp 输入是一个 OpConf，下面以一个卷积的激活函数为例子，这个 OpConf 是我从 lenet 保存下来的计算图中截取出来的。我们可以看到有名字，设备，输入，输出，op 类型等信息。

    op_list {
      name: "conv1-activation"
      device_tag: "gpu"
      scope_symbol_id: 4611686018427531262
      user_conf {
        op_type_name: "relu"
        input {
          key: "in"
          value {
            s: "conv1-bias_add/out_0"
          }
        }
        output {
          key: "out"
          value {
            s: "conv1-activation/out_0"
          }
        }
      }
    }

• 在 Job 函数中，我们如何知道 OpConf 的信息是什么呢？以 ReLU 为例子。relu 函数调用了 user_op_builder 来添加算子到计算图中，添加的时候调用的是 InferAndTryRun，将算子加到计算图中。CurJobAddOp 需要的信息，就是上面的 OpConf。这个 OpConf 由 user_op_builder 使用 Op、Input、Output、Build 构建而来。调用了这一些操作之后，我们就有了一个 OpConf 了，接着调用 InferAndTryRun 将 OpConf 添加到计算图当中。

# python/oneflow/compatible/single_client/ops/math_ops.py: 564
def relu(
    x: oneflow._oneflow_internal.BlobDesc, name: Optional[str] = None
) -> oneflow._oneflow_internal.BlobDesc:
    return (
        flow.user_op_builder(name if name is not None else id_util.UniqueStr("Relu_"))
        .Op("relu")
        .Input("in", [x])
        .Output("out")
        .Build()
        .InferAndTryRun()
        .RemoteBlobList()[0]
    )

# python/oneflow/compatible/single_client/ops/user_op_builder.py: 167
def InferAndTryRun(self):
    compile_context.CurJobAddOp(self.op_conf_)
    return self

# python/oneflow/compatible/single_client/ops/user_op_builder.py: 257
def Input(self, input_name, input_blob_list):
    """Set input blob of op

    Args:
        input_name (str): input name of blob
        input_blob_list : list of blobs

    Returns:
        self
    """
    assert isinstance(input_blob_list, (tuple, list))
    input_conf = self.user_op_.op_conf_.user_conf.input
    input_conf[input_name].ClearField("s")
    for input_blob in input_blob_list:
        input_conf[input_name].s.append(input_blob.unique_name)
    return self

def Output(self, output_name, num=1):
    """Set output blob of op

    Args:
        output_name (str): name of output blob
        num (int, optional):  Defaults to 1.

    Returns:
        self
    """
    assert isinstance(num, int) and num >= 1
    out_lbns = []
    for i in range(num):
        lbn = "{}/{}_{}".format(self.user_op_.op_conf_.name, output_name, i)
        out_lbns.append(lbn)
    self.user_op_.op_conf_.user_conf.output[output_name].s[:] = out_lbns
    self.user_op_.output_arg_key_list_.append(output_name)
    return self

def CheckAndComplete(self):
    assert self.user_op_.op_conf_.user_conf.op_type_name != ""
    self.user_op_.op_conf_ = c_api_util.CheckAndCompleteUserOpConf(
        self.user_op_.op_conf_
    )
    return self

def Build(self):
    """Build op when in/output and other attribute set up.

    Returns:
        self

    """
    return self.CheckAndComplete().user_op_

BlobDesc

接下来我们来看看 oneflow._oneflow_internal.BlobDesc，在构图的过程中，BlobDesc 保存了一些信息，有 shape，op 名字等。这些信息可以用在构图过程中，设置 OpConf 的属性，比如设置 OpConf 的输入就依赖了 BlobDesc 的 unique_name 属性。

BlobDesc 的实现在 C++ 里面，正如名字所暗示的那样，是 Blob 的描述，用来保存 Blob 的描述信息。在文末有 lenet 的完整训练脚本，lenet 函数里面的构造过程，像 conv1，pool1，conv2，pool2 等中间变量，其实就是一个 BlobDesc，而调用 flow.layers.conv2d 的方法，将会根据这些 BlobDesc 构建一个 OpConf，然后通过调用 CurJobAddOp 将 OpConf 放入计算图当中。

// oneflow/core/framework/py_blob_desc.h: 40
class Tensor {
 public:
  virtual ~Tensor() = default;

  virtual std::shared_ptr<cfg::LogicalBlobId> lbi() const = 0;
  virtual std::string logical_blob_name() const = 0;
  virtual std::string op_name() const = 0;
  virtual std::string blob_name() const = 0;
  virtual std::shared_ptr<Shape> shape() const = 0;
  virtual DataType dtype() const = 0;
  virtual std::shared_ptr<cfg::ParallelConf> parallel_conf() const = 0;
};

class BlobDesc : public Tensor {
 public:
  BlobDesc(const std::shared_ptr<cfg::LogicalBlobId>& lbi,
           const std::shared_ptr<Distribute>& distribute);

  BlobDesc(const BlobDesc& blob_desc) = default;
  virtual ~BlobDesc() override = default;

  virtual std::shared_ptr<cfg::LogicalBlobId> lbi() const override;
  virtual std::string logical_blob_name() const override;
  virtual std::string op_name() const override;
  virtual std::string blob_name() const override;
  virtual std::shared_ptr<Shape> shape() const override;
  virtual DataType dtype() const override;
  virtual std::shared_ptr<cfg::ParallelConf> parallel_conf() const override;

  virtual bool is_dynamic() const;
  virtual std::shared_ptr<Distribute> distribute() const;
  virtual std::string unique_name() const;

  void set_distribute(const std::shared_ptr<Distribute> distribute);

 protected:
  Maybe<std::string> Distribute2Str() const;

  std::shared_ptr<cfg::LogicalBlobId> lbi_;
  std::shared_ptr<Distribute> distribute_;
  std::string lbn_;
};

// oneflow/core/framework/py_remote_blob.h
class ConsistentBlob : public BlobDesc {};
class LazyConsistentBlob : public ConsistentBlob {};
class MirroredBlob : public BlobDesc {};
class LazyMirroredBlob : public MirroredBlob {};
class EagerBlobTrait {};
class EagerConsistentBlob : public EagerBlobTrait, public ConsistentBlob {};
class EagerMirroredBlob : public EagerBlobTrait, public MirroredBlob {};

总结

Python 端是如何构图的呢？现在已经可以回答这个问题了，在 Job 函数内，使用 BlobDesc 来保存 Blob 描述信息，每个算子调用 user_op_builder 的方法利用 BlobDesc 的信息将 OpConf 构造出来，然后调用 CurJobAddOp 将 Op 添加到计算图当中。

这篇文章的开头，还讲了构图的触发时机。在用户第一次调用 Job 函数的时候，触发构图。原理是通过一个 global_function 装饰器，装饰用户定义的 Job 函数，用户第一次调用 Job 函数的时候，会尝试初始化 Session，如果已经初始化了，那么不再初始化。在初始化 Session 的时候，去走一遍用户定义的 Job 函数，构建计算图，也就是上面说的：每个算子调用 user_op_builder 的方法利用 BlobDesc 的信息将 OpConf 构造出来，然后调用 CurJobAddOp 将 Op 添加到计算图当中。

从上面的分析可以看出来，Python 端的算子，会检查用户的输入是否合理，比如卷积层的 groups 参数是否比 filters 参数大等。检查结束后，将参数输入给 user_op_builder 来构图。

参考链接

[1] https://www.cnblogs.com/zzk0/p/15009227.html

附: lenet 的训练脚本

# lenet_train.py
from oneflow.compatible import single_client as flow
from oneflow.compatible.single_client import typing as tp

BATCH_SIZE = 100
flow.config.enable_legacy_model_io(False)
# flow.config.gpu_device_num(2)
# flow.env.ctrl_port(9998)
# func_config = flow.function_config()
# func_config.default_logical_view(flow.scope.mirrored_view())


def lenet(data, train=False):
    initializer = flow.truncated_normal(0.1)
    conv1 = flow.layers.conv2d(
        data,
        32,
        5,
        padding="SAME",
        activation=flow.nn.relu,
        name="conv1",
        kernel_initializer=initializer,
    )
    pool1 = flow.nn.max_pool2d(
        conv1, ksize=2, strides=2, padding="SAME", name="pool1", data_format="NCHW"
    )
    conv2 = flow.layers.conv2d(
        pool1,
        64,
        5,
        padding="SAME",
        activation=flow.nn.relu,
        name="conv2",
        kernel_initializer=initializer,
    )
    pool2 = flow.nn.max_pool2d(
        conv2, ksize=2, strides=2, padding="SAME", name="pool2", data_format="NCHW"
    )
    # reshape = flow.reshape(pool2, [pool2.shape[0], -1])
    reshape = flow.flatten(pool2, start_dim=1)
    hidden = flow.layers.dense(
        reshape,
        512,
        activation=flow.nn.relu,
        kernel_initializer=initializer,
        name="dense1",
    )
    if train:
        hidden = flow.nn.dropout(hidden, rate=0.5, name="dropout")
    return flow.layers.dense(hidden, 10, kernel_initializer=initializer, name="dense2")


@flow.global_function(type="train")
def train_job(
    images: tp.Numpy.Placeholder((BATCH_SIZE, 1, 28, 28), dtype=flow.float),
    labels: tp.Numpy.Placeholder((BATCH_SIZE,), dtype=flow.int32),
) -> tp.Numpy:
    with flow.scope.placement("gpu", "0:0"):
        logits = lenet(images, train=True)
        loss = flow.nn.sparse_softmax_cross_entropy_with_logits(
            labels, logits, name="softmax_loss"
        )
    lr_scheduler = flow.optimizer.PiecewiseConstantScheduler([], [0.1])
    flow.optimizer.SGD(lr_scheduler, momentum=0).minimize(loss)
    return loss


if __name__ == "__main__":
    (train_images, train_labels), (test_images, test_labels) = flow.data.load_mnist(
        BATCH_SIZE, BATCH_SIZE
    )

    for epoch in range(20):
        for i, (images, labels) in enumerate(zip(train_images, train_labels)):
            loss = train_job(images, labels)
            if i % 20 == 0:
                print(loss.mean())
    flow.checkpoint.save("./lenet_models_1")  # need remove the existed folder
    print("model saved")