OneFlow: 从 Op 到 Job
前言
前面在初始化 Session 的时候，通过 CurJobAddOp 将 Op 加入到计算图当中。在启动 Session 之前，需要将这些 Op 组成的计算图编译成 JobSet。这篇文章主要分析如何从一个个 op 变成一个 Job，Job 如何在 Complete 函数内进行优化的。
流程分析
上一篇文章中，_CompileJob 将会去调用用户定义的 Job 函数，将 Op 加入到计算图中。这个过程执行完成之后，会调用 CurJobBuildAndInferCtx_Complete，生成 Job 或 JobSet(如果用户定义了多个 Job)。Complete 的时候，会执行一些 Pass，将计算图进行改写优化。
# 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",
)
CurJobBuildAndInferCtx_Complete
Complete 的时候做了什么呢？它的输入是什么？输出又是什么？输入和输出都是 job，问题是这个输入的 job 是从哪里来的呢？我们可以看到 job 和 mut_job 等方法的调用，这些方法返回的都是类成员 job_。这个 job_ 是怎么构造来的呢？
// oneflow/core/job/job_build_and_infer_ctx.cpp: 960
Maybe<void> LazyJobBuildAndInferCtx::Complete() {
CHECK_GT_OR_RETURN(job().net().op_size(), 0)
<< " Sorry, nn.Graph need at least 1 op in net, but get 0 now.";
CHECK_NOTNULL(Global<JobDesc>::Get());
Global<JobDesc>::Delete();
auto scope = std::make_unique<GlobalJobDescScope>(mut_job()->job_conf(), job_id());
JobPassCtx job_pass_ctx(GlobalJobDesc());
auto DoPass = [&](const std::string& pass_name) -> Maybe<void> {
return JobPass4Name(pass_name)(mut_job(), &job_pass_ctx);
};
// ... DoPass
return Maybe<void>::Ok();
}
Job 是哪里来的
来看看构造函数，和 SetJobConf，这两个函数都会修改 JobBuildAndInferCtx 内部的 job_ 成员。
// oneflow/core/job/job_build_and_infer_ctx.cpp: 103
JobBuildAndInferCtx::JobBuildAndInferCtx(Job* job, int64_t job_id)
: job_(job), job_id_(job_id), unique_op_name_index_(0) {
is_job_conf_frozen_ = false;
has_job_conf_ = false;
}
Maybe<void> JobBuildAndInferCtx::SetJobConf(const JobConfigProto& job_conf) {
CHECK_OR_RETURN(!is_job_conf_frozen_) << Error::JobConfFrozenError();
CHECK_OR_RETURN(!has_job_conf_) << Error::JobConfRepeatedSetError();
has_job_conf_ = true;
CHECK_EQ_OR_RETURN(job_->job_conf().job_name(), job_conf.job_name())
<< Error::JobNameNotEqualError() << "job name you set: " << job_conf.job_name()
<< " not equal to origin job name: " << job_->job_conf().job_name();
job_->mutable_job_conf()->CopyFrom(job_conf);
CHECK_ISNULL_OR_RETURN(Global<JobDesc>::Get());
Global<JobDesc>::New(job_conf, job_id_);
return Maybe<void>::Ok();
}
- 如果你对上一篇还有点印象的话，会记得，在 CurJobAddOp 之前，需要调用 SetJobConf。下面是运行 lenet_model.py 的 JobConf。所以 SetJobConf 会对 JobBuildAndInferCtx 的 job_ 成员进行初始化。
job_name: "train_job"
train_conf {
}
- 设置了 JobConf 之后，就可以进行添加算子，调用 CurJobAddOp，然后调用 AddAndInferConsistentOp，最后到了 AddAndInferOp。AddAndInferOp 特别长，主要是进行一些 SBP 属性的推理，对于 SBP 我完全是不懂呢，所以先跳过。总之这个 AddAndInferOp 里面调用了一个重要的方法：AddOpAndUpdateJobParallelViewConf。正如名字暗示的那样，AddOp！
// oneflow/core/job/job_build_and_infer_ctx.cpp: 546
Maybe<OpAttribute> JobBuildAndInferCtx::AddAndInferConsistentOp(const OperatorConf& op_conf) {
CHECK_OR_RETURN(op_conf.has_scope_symbol_id());
const auto& scope = Global<symbol::Storage<Scope>>::Get()->Get(op_conf.scope_symbol_id());
const auto& parallel_desc = *JUST(scope.GetParallelDesc(op_conf));
const auto* job_desc = JUST(scope.job_desc());
return AddAndInferOp(op_conf, parallel_desc.parallel_conf(), job_desc, false);
}
// TODO(): add handle error of same interface op blob between jobs
Maybe<OpAttribute> JobBuildAndInferCtx::AddAndInferOp(const OperatorConf& op_conf,
const ParallelConf& origin_parallel_conf,
const JobDesc* job_desc,
bool is_mirrored_parallel_view) {
// ...
AddOpAndUpdateJobParallelViewConf(*new_op_conf, parallel_desc, nd_sbp_sig_conf,
is_mirrored_parallel_view);
// ...
return op->GetOpAttributeWithoutOpNameAndLbn();
}
- AddOpAndUpdateJobParallelViewConf 这个方法的最后一行，将 Op 添加到了 job_ 里面。于是呢，在我们调用 Complete 的时候，我们可以找到这个 Op，找到整个计算图。
// oneflow/core/job/job_build_and_infer_ctx.cpp: 176
void JobBuildAndInferCtx::AddOpAndUpdateJobParallelViewConf(
const OperatorConf& operator_conf, const ParallelDesc& parallel_desc,
const cfg::NdSbpSignature& nd_sbp_signature, bool is_mirrored_parallel_view) const {
auto* op_name2sbp_sig =
job_->mutable_job_parallel_view_conf()->mutable_op_name2sbp_signature_conf();
auto* op_name2nd_sbp_sig =
job_->mutable_job_parallel_view_conf()->mutable_op_name2nd_sbp_signature_conf();
if (nd_sbp_signature.bn_in_op2nd_sbp().size() > 0) {
nd_sbp_signature.ToProto(&(*op_name2nd_sbp_sig)[operator_conf.name()]);
if (parallel_desc.hierarchy()->NumAxes() == 1) {
cfg::SbpSignature sbp_signature;
NdSbpSignatureToSbpSignature(nd_sbp_signature, &sbp_signature);
sbp_signature.ToProto(&(*op_name2sbp_sig)[operator_conf.name()]);
}
}
auto* op_name2is_mirrored_parallel_view =
job_->mutable_job_parallel_view_conf()->mutable_op_name2is_mirrored_parallel_view();
if (is_mirrored_parallel_view) {
(*op_name2is_mirrored_parallel_view)[operator_conf.name()] = true;
}
job_->mutable_net()->add_op()->CopyFrom(operator_conf);
}
proto 定义
Job
前面的 job_ 其实就是一个 Proto 的 message，不妨将它的定义找来看一看。一个 Job 由 5 个部分组成，其中 DLNetConf 就是计算图，JobConfigProto 表示这个 job 的配置(train 或者 predict)。
// oneflow/core/job/job.proto: 29
message Job {
optional DLNetConf net = 1;
optional Placement placement = 2;
required JobConfigProto job_conf = 3;
optional JobParallelViewConf job_parallel_view_conf = 4;
optional JobHelperConf helper = 5;
}
DLNetConf
DLNetConf 是一个由 OperatorConf 构成的计算图。
message DLNetConf {
repeated OperatorConf op = 1;
}
OperatorConf
OperatorConf 定义了一个算子，算子的名字，算子的设备等信息。
message OperatorConf {
required string name = 1;
optional string device_tag = 4 [default = "invalid_device"];
repeated string ctrl_in_op_name = 7;
optional int64 scope_symbol_id = 8;
optional uint32 stream_index_hint = 9;
optional string pass_tag = 10;
oneof op_type {
// ...
}
}
Complete
既然我们知道了 Job 是怎么来的，它长什么样，那么我们就可以开始尝试着将它画出来！每个 Pass 之后就画一张图看看。OneFlow 中提供了 Graph 基类，里面不仅仅有一些常用的图算法，还有可视化，可以将图转为 .dot 文件。OneFlow 中还提供了一个 OpGraph 类，它接收一个 Job 对象，然后可以调用可视化的方法。
初始计算图
Complete 之前的计算图，就是用户定义的那样，按照 Job 函数里面的算子进行构图，调用 CurJobAddOp 加入到 Job 对象里面。下面就是在 Complete 开始之前可视化出来的计算图。文末附有训练脚本，这次为了方便查看，选择了只有几个节点的 mlp 识别手写字体模型。
最终计算图
在 Complete 中的 Pass 中，并不是所有的 Pass 都会改写计算图，有的计算图可能因为没有配置，所以不执行。目前肉眼可见对计算图改变最大的 Pass 是：GenerateBackwardAndOptimizerOpConfs。这个 Pass 会生成后向算子和优化器节点。
损失函数节点
损失函数节点，接受 Input_1 和 dense2-bias_add 两个节点的输出，Input_1 是标签，dense2-bias_add 是前向传播计算的结果。用这两个节点的值就可以计算出损失。在损失函数节点后面，有三个节点，一个是 Return_4，这个节点输出损失，还有另外两个节点用来输出损失函数的梯度，这两个节点的值我猜想应该是 1。
Optimizer 节点
这个训练脚本使用的是 Adam 优化器，Adam 在之前的博客中有记录，对于 dense1 这个权重矩阵，它需要记录动量和梯度平方和，也就是需要两个额外的和 dense1 形状一样的权重矩阵。
附：mlp 识别手写字体
from oneflow.compatible import single_client as flow
from oneflow.compatible.single_client import typing as tp
import numpy as np
BATCH_SIZE = 100
flow.enable_eager_execution(False)
def mlp_model(images, labels, train=True):
# [batch_size, image_sizes] -> [batch_size, pixels]
# reshape = flow.reshape(images, [images.shape[0], -1])
reshape = flow.flatten(images, start_dim=1)
# dense, [batch_size, pixels] -> [batch_size, 500]
initializer1 = flow.random_uniform_initializer(-1 / 28.0, 1 / 28.0)
hidden = flow.layers.dense(
reshape,
500,
activation=flow.nn.relu,
kernel_initializer=initializer1,
bias_initializer=initializer1,
name="dense1"
)
# dense, [batch_size, 500] -> [batch_size, logits]
initializer2 = flow.random_uniform_initializer(
-np.sqrt(1 / 500.0), np.sqrt(1 / 500.0)
)
logits = flow.layers.dense(
hidden,
10,
kernel_initializer=initializer2,
bias_initializer=initializer2,
name="dense2"
)
if train:
loss = flow.nn.sparse_softmax_cross_entropy_with_logits(labels, logits)
return loss
else:
return logits
@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"):
loss = mlp_model(images, labels)
flow.optimizer.Adam(
flow.optimizer.PiecewiseConstantScheduler([], [0.001])
).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("Epoch [{}/{}], Loss: {:.4f}".format(epoch + 1, 20, loss.mean()))
flow.checkpoint.save("./mlp_model")
