def _EagerRunModelLoad(var_op_conf, snapshot_path):
assert isinstance(snapshot_path, str)
assert os.path.basename(snapshot_path) == "out"
snapshot_path = os.path.dirname(snapshot_path)
assert os.path.basename(snapshot_path) == var_op_conf.name
snapshot_path = os.path.dirname(snapshot_path)
path_input_op_conf, path_lbi = _GenModelIOPathInputOpConfAndRetLbi()
path_input_blob_objects = {}
(
BuildModelIOPathInputInstruction,
BuildFeedPathInstruction,
) = _MakeModelIOPathInputBuilds(path_input_op_conf, snapshot_path,
path_input_blob_objects)
model_load_op_conf, _ = _GenModelLoadOpConfAndRetLbi(var_op_conf, path_lbi)
model_load_blob_objects = {}
def BuildModelLoadInstruction(builder):
path_blob_object = path_input_blob_objects["out"]
model_load_blob_objects["path"] = path_blob_object
op_attribute = op_infer_util.Infer(
model_load_op_conf, ibn2blob_object=model_load_blob_objects)
parallel_conf = path_blob_object.parallel_desc_symbol.parallel_conf
builder.StatelessCall(op_attribute,
parallel_conf,
bn_in_op2blob_object=model_load_blob_objects)
sess = session_ctx.GetDefaultSession()
with scope_util.ScopeContext(scope_util.MakeScope(_BuildNotMirroredScope)):
vm_util.LogicalRun(BuildModelIOPathInputInstruction)
vm_util.LogicalRun(BuildFeedPathInstruction)
vm_util.LogicalRun(BuildModelLoadInstruction)
return model_load_blob_objects["out_0"]
def _EagerRunModelSave(var_blobs, snapshot_path):
path_input_op_conf, path_lbi = _GenModelIOPathInputOpConfAndRetLbi()
path_input_blob_objects = {}
(
BuildModelIOPathInputInstruction,
BuildFeedPathInstruction,
) = _MakeModelIOPathInputBuilds(path_input_op_conf, snapshot_path,
path_input_blob_objects)
model_save_op_conf = _GenModelSaveOpConf(var_blobs, path_lbi)
model_save_blob_objects = {}
def BuildModelSaveInstruction(builder):
path_blob_object = path_input_blob_objects["out"]
model_save_blob_objects["path"] = path_blob_object
for i, blob in enumerate(var_blobs):
model_save_blob_objects["in_{}".format(i)] = blob.blob_object
op_attribute = op_infer_util.Infer(
model_save_op_conf, ibn2blob_object=model_save_blob_objects)
parallel_conf = path_blob_object.parallel_desc_symbol.parallel_conf
builder.StatelessCall(op_attribute,
parallel_conf,
bn_in_op2blob_object=model_save_blob_objects)
sess = session_ctx.GetDefaultSession()
with scope_util.ScopeContext(scope_util.MakeScope(_BuildNotMirroredScope)):
vm_util.LogicalRun(BuildModelIOPathInputInstruction)
vm_util.LogicalRun(BuildFeedPathInstruction)
vm_util.LogicalRun(BuildModelSaveInstruction)
def _EagerRunModelLoad(var_op_conf, snapshot_path):
assert isinstance(snapshot_path, str)
assert os.path.basename(snapshot_path) == "out"
snapshot_path = os.path.dirname(snapshot_path)
assert os.path.basename(snapshot_path) == var_op_conf.name
snapshot_path = os.path.dirname(snapshot_path)
path_input_op_conf, path_lbi = _GenModelLoadPathInputOpConfAndRetLbi()
model_load_op_conf, _ = _GenModelLoadOpConfAndRetLbi(var_op_conf, path_lbi)
path_input_blob_objects = {}
model_load_blob_objects = {}
def BuildNotMirroredScope(old_scope, builder):
return old_scope.BuildWithNewIsMirrored(builder, False)
def BuildModelLoadPathInputInstruction(builder):
op_attribute = op_infer_util.Infer(path_input_op_conf,
ibn2blob_object={})
parallel_conf = oneflow.placement.current_scope().default_parallel_conf
builder.StatelessCall(op_attribute,
parallel_conf,
bn_in_op2blob_object=path_input_blob_objects)
def FeedPath(ofblob):
ofblob.CopyFromNdarray(
np.frombuffer(snapshot_path.encode("ascii"), dtype=np.int8))
def BuildFeedPathInstruction(builder):
blob_object = path_input_blob_objects["out"]
builder.FeedBlob(blob_object, FeedPath)
builder.InsertRemoveForeignCallbackInstruction(blob_object.object_id,
FeedPath)
def BuildModeLoadInstruction(builder):
path_blob_object = path_input_blob_objects["out"]
model_load_blob_objects["path"] = path_blob_object
op_attribute = op_infer_util.Infer(
model_load_op_conf, ibn2blob_object=model_load_blob_objects)
parallel_conf = path_blob_object.parallel_desc_symbol.parallel_conf
builder.StatelessCall(op_attribute,
parallel_conf,
bn_in_op2blob_object=model_load_blob_objects)
sess = session_ctx.GetDefaultSession()
with sess.NewCurrentScope(sess.MakeScope(BuildNotMirroredScope)):
vm_util.LogicalRun(BuildModelLoadPathInputInstruction)
vm_util.LogicalRun(BuildFeedPathInstruction)
vm_util.LogicalRun(BuildModeLoadInstruction)
return model_load_blob_objects["out_0"]
def AsyncSlice(Yield):
def build(builder):
op_conf = op_conf_pb.OperatorConf()
# device_tag doesn't matter for logical_slice op
device_tag = oneflow.current_scope(
).device_parallel_desc_symbol.device_tag
op_conf.device_tag = device_tag
op_conf.name = op_name
op_conf.user_conf.op_type_name = "logical_slice"
op_conf.user_conf.input["x"].s.append("{}/x_0".format(op_name))
op_conf.user_conf.output["y"].s.append("{}/y_0".format(op_name))
input_blob_object = input_blob.blob_object
parallel_conf = input_blob_object.parallel_desc_symbol.parallel_conf
op_conf.user_conf.attr["parallel_conf"].at_string = str(
parallel_conf)
op_conf.user_conf.attr["start"].at_list_int64.val[:] = start
op_conf.user_conf.attr["stop"].at_list_int64.val[:] = stop
op_conf.user_conf.attr["step"].at_list_int64.val[:] = [
1
] * len(start)
bn_in_op2blob_object = {"x_0": input_blob_object}
scope_symbol_id = _GetScopeSymbolIdFromEagerBlob(input_blob)
op_attribute = op_infer_util.Infer(op_conf, bn_in_op2blob_object,
scope_symbol_id)
builder.StatelessCall(
op_attribute,
parallel_conf=parallel_conf,
bn_in_op2blob_object=bn_in_op2blob_object,
)
Yield(bn_in_op2blob_object["y_0"])
vm_util.LogicalRun(build)
def CreateEagerVariableBlob(op_attribute, job_name=""):
bn_in_op2blob_object = oneflow_api.deprecated.BnInOp2BlobObject()
def BuildInstruction(builder):
parallel_conf = (
oneflow.current_scope().device_parallel_desc_symbol.parallel_conf
)
cfg_op_attribute = oneflow_api.deprecated.MakeOpAttributeByString(
str(op_attribute)
)
builder.StatelessCall(
cfg_op_attribute,
parallel_conf,
bn_in_op2blob_object,
boxing_util.BoxingTo,
vm_util._FindOrCreateDelegateBlobObject,
)
vm_util.LogicalRun(BuildInstruction)
lbi = lbi_util.LogicalBlobId()
lbi.set_op_name(op_attribute.op_conf.name)
lbi.set_blob_name(op_attribute.op_conf.variable_conf.out)
if not isinstance(lbi, lbi_util.LogicalBlobId):
cfg_lbi = lbi_util.LogicalBlobId()
cfg_lbi.set_op_name(lbi.op_name)
cfg_lbi.set_blob_name(lbi.blob_name)
lbi = cfg_lbi
return oneflow_api.EagerConsistentBlob(
lbi,
blob_object=bn_in_op2blob_object["out"],
blob_register=blob_register,
job_name=job_name,
)
def FetchBlobNumpy(blob_object):
consistent_blob_name = None
def BoxingToSingleDevice(builder):
parallel_conf = placement_pb.ParallelConf()
parallel_conf.device_tag = blob_object.parallel_desc_symbol.device_tag
parallel_conf.device_name.append("{}:{}".format(0, 0))
tmp_parallel_desc_symbol = builder.GetParallelDescSymbol(parallel_conf)
tmp_op_arg_parallel_attr = op_arg_util.OpArgParallelAttribute(
tmp_parallel_desc_symbol,
blob_object.op_arg_parallel_attr.sbp_parallel,
blob_object.op_arg_parallel_attr.opt_mirrored_parallel,
)
with oneflow.scope.placement(
self.parallel_conf.device_tag, list(self.parallel_conf.device_name)
):
tmp_blob_object = boxing_util.BoxingTo(
builder, blob_object, tmp_op_arg_parallel_attr
)
nonlocal consistent_blob_name
consistent_blob_name = "{}-consistent".format(self.logical_blob_name)
if not blob_register.HasObject4BlobName(consistent_blob_name):
blob_register.SetObject4BlobName(
consistent_blob_name, tmp_blob_object
)
vm_util.LogicalRun(BoxingToSingleDevice)
return eager_blob_util.EagerPhysicalBlob(consistent_blob_name).numpy()
def DistributeConcatOrAdd(op_attribute, parallel_conf, blob_register):
op_parallel_desc_sym = symbol_storage.GetSymbol4Id(
op_attribute.parallel_signature.op_parallel_desc_symbol_id)
parallel_size = len(op_attribute.input_bns)
op_arg_parallel_attr = op_arg_util.GetOpArgParallelAttribute(
op_parallel_desc_sym, op_attribute, "out")
op_arg_blob_attr = op_arg_util.GetOpArgBlobAttribute(op_attribute, "out")
parallel_sig = op_attribute.parallel_signature.bn_in_op2parallel_desc_symbol_id
def GetInBlobObject(builder, i, bn_in_op2blob_object):
ibn = "in_%s" % i
origin_blob_object = bn_in_op2blob_object[ibn]
in_op_parallel_desc_sym = symbol_storage.GetSymbol4Id(
parallel_sig[ibn])
in_op_arg_parallel_attr = op_arg_util.GetOpArgParallelAttribute(
in_op_parallel_desc_sym, op_attribute, ibn)
return boxing_util.BoxingTo(builder, origin_blob_object,
in_op_arg_parallel_attr)
def BuildInstruction(builder):
with blob_register.BnInOp2BlobObjectScope(
op_attribute) as bn_in_op2blob_object:
def GetPhysicalInBlob(i):
return GetInBlobObject(builder, i, bn_in_op2blob_object)
in_blob_objects = [
GetPhysicalInBlob(i) for i in range(parallel_size)
]
bn_in_op2blob_object[
"out"] = builder.PackPhysicalBlobsToLogicalBlob(
in_blob_objects, op_arg_parallel_attr, op_arg_blob_attr)
vm_util.LogicalRun(BuildInstruction)
def CreateEagerVariableBlob(op_attribute, job_name=""):
bn_in_op2blob_object = {}
def BuildInstruction(builder):
parallel_conf = (
oneflow.current_scope().device_parallel_desc_symbol.parallel_conf)
builder.StatelessCall(op_attribute,
parallel_conf,
bn_in_op2blob_object=bn_in_op2blob_object)
vm_util.LogicalRun(BuildInstruction)
lbi = lbi_util.LogicalBlobId()
lbi.set_op_name(op_attribute.op_conf.name)
lbi.set_blob_name(op_attribute.op_conf.variable_conf.out)
if not isinstance(lbi, lbi_util.LogicalBlobId):
cfg_lbi = lbi_util.LogicalBlobId()
cfg_lbi.set_op_name(lbi.op_name)
cfg_lbi.set_blob_name(lbi.blob_name)
lbi = cfg_lbi
return oneflow_api.EagerConsistentBlob(
lbi,
blob_object=bn_in_op2blob_object["out"],
blob_register=blob_register,
job_name=job_name,
)
def EagerReturnRemoteBlob(remote_blob, allow_cpu_return_op=True):
if not hob.is_trainable(None):
return remote_blob
op_conf, lbi, scope = _GetReturnOpConfAndOutLbiAndScope(
remote_blob, allow_cpu_return_op)
if remote_blob.blob_object.op_arg_parallel_attr.is_mirrored():
add_and_infer = compile_context.CurJobAddMirroredOp
else:
add_and_infer = compile_context.CurJobAddConsistentOp
op_attribute = add_and_infer(op_conf, scope)
def BuildInstruction(builder):
get_blob_scope = blob_register_util.BnInOp2BlobObjectScope
with get_blob_scope(blob_register,
op_attribute) as bn_in_op2blob_object:
cfg_op_attribute = oneflow_api.deprecated.MakeOpAttributeByString(
str(op_attribute))
builder.StatelessCall(
cfg_op_attribute,
remote_blob.blob_object.parallel_desc_symbol.parallel_conf,
bn_in_op2blob_object,
boxing_util.BoxingTo,
vm_util._FindOrCreateDelegateBlobObject,
)
vm_util.LogicalRun(BuildInstruction)
return remote_blob_util.RemoteBlob(lbi)
def _EagerRunModelInit(var_op_conf):
op_conf, _ = _GenModelInitOpConfAndRetLbi(var_op_conf)
bn_in_op2blob_object = oneflow_api.deprecated.BnInOp2BlobObject()
def BuildModelInitInstruction(builder):
upstream_signature = op_node_signature_pb.OpNodeSignature()
op_conf.scope_symbol_id = oneflow.current_scope().symbol_id
op_attribute = c_api_util.InferOpConf(op_conf, upstream_signature)
parallel_conf = (
oneflow.current_scope().device_parallel_desc_symbol.parallel_conf)
cfg_op_attribute = oneflow_api.deprecated.MakeOpAttributeByString(
str(op_attribute))
builder.StatelessCall(
cfg_op_attribute,
parallel_conf,
bn_in_op2blob_object,
boxing_util.BoxingTo,
vm_util._FindOrCreateDelegateBlobObject,
)
sess = session_ctx.GetDefaultSession()
with scope_util.ScopeContext(scope_util.MakeScope(_BuildNotMirroredScope)):
vm_util.LogicalRun(BuildModelInitInstruction)
return bn_in_op2blob_object["out_0"]
def AsyncFeedValueToInterfaceBlob(Yield):
def build(builder):
blob_object = GetEagerInterfaceBlob(op_name).blob_object
FeedValueToInterfaceBlobObject(blob_object, ndarray)
Yield()
vm_util.LogicalRun(build)
def AsyncFeedValueToInterfaceBlob(Yield):
def build(builder):
blob_object = builder.MakeLazyRefBlobObject(op_name)
if blob_object.op_arg_blob_attr.is_tensor_list:
input_blob_def = input_blob_def_util.MirroredTensorListDef(
[x.shape for x in ndarray],
dtype=dtype_util.convert_numpy_dtype_to_oneflow_dtype(
ndarray.dtype),
)
elif blob_object.op_arg_parallel_attr.is_mirrored():
input_blob_def = input_blob_def_util.MirroredTensorDef(
ndarray.shape,
dtype=dtype_util.convert_numpy_dtype_to_oneflow_dtype(
ndarray.dtype),
)
else:
input_blob_def = input_blob_def_util.FixedTensorDef(
ndarray.shape,
dtype=dtype_util.convert_numpy_dtype_to_oneflow_dtype(
ndarray.dtype),
)
push_util.FeedValueToEagerBlob(blob_object, input_blob_def,
ndarray)
Yield()
vm_util.LogicalRun(build)
def FeedValueToInterfaceBlobObject(blob_object, ndarray):
flow.sync_default_session()
def build(builder):
if blob_object.op_arg_blob_attr.is_tensor_list:
input_blob_def = input_blob_def_util.MirroredTensorListDef(
[x.shape for x in ndarray],
dtype=dtype_util.convert_numpy_dtype_to_oneflow_dtype(
ndarray.dtype),
)
elif blob_object.op_arg_parallel_attr.is_mirrored():
input_blob_def = input_blob_def_util.MirroredTensorDef(
ndarray.shape,
dtype=dtype_util.convert_numpy_dtype_to_oneflow_dtype(
ndarray.dtype),
)
else:
input_blob_def = input_blob_def_util.FixedTensorDef(
ndarray.shape,
dtype=dtype_util.convert_numpy_dtype_to_oneflow_dtype(
ndarray.dtype),
)
push_util.FeedValueToEagerBlob(blob_object, input_blob_def, ndarray)
vm_util.LogicalRun(build)
def api_one_to_one_assign(ref, value):
assert hob.eager_execution_enabled(None)
vm_util.LogicalRun(
lambda builder: builder.Build121AssignInstruction(
ref.blob_object, value.blob_object
)
)
return ref
def eager_system_assign(ref, value, validate_shape=None, use_locking=None, name=None):
op_conf = _SystemAssignOpConf(ref, value, name=name)
# no backward for assign
vm_util.LogicalRun(
lambda builder: boxing_util.BuildAssignInstruction(
builder, ref.blob_object, value.blob_object, op_conf
)
)
return ref
def MakeParallelDescSymbol(parallel_conf):
symbol_id = None
def BuildInstruction(builder):
nonlocal symbol_id
symbol_id = builder.GetParallelDescSymbol(parallel_conf).symbol_id
vm_util.LogicalRun(BuildInstruction)
return symbol_id
def MakeParallelDescSymbol(parallel_conf_str):
parallel_conf = text_format.Parse(parallel_conf_str, placement_pb.ParallelConf())
symbol_id = None
def BuildInstruction(builder):
nonlocal symbol_id
symbol_id = builder.GetParallelDescSymbol(parallel_conf).symbol_id
vm_util.LogicalRun(BuildInstruction)
return symbol_id
def MakeInitialScope(job_conf, device_tag, machine_device_ids, is_mirrored):
scope = None
def BuildInitialScope(builder):
nonlocal scope
scope = scope_util.BuildInitialScope(builder, job_conf, device_tag,
machine_device_ids, is_mirrored)
vm_util.LogicalRun(BuildInitialScope)
return scope
def InitOpKernel(self):
def BuildInstruction(builder):
if not isinstance(
self.op_conf,
oneflow_api.oneflow.core.operator.op_conf.OperatorConf):
cfg_op_conf = oneflow_api.deprecated.MakeOpConfByString(
str(self.op_conf))
self.set_opkernel_object(builder.NewOpKernelObject(cfg_op_conf))
vm_util.LogicalRun(BuildInstruction)
def _GetPhysicalBlobObjects(logical_blob_object, lbi):
blob_register = blob_register_util.GetDefaultBlobRegister()
physical_blob_objects = None
def BuildLogical2PhysicalInstruction(builder):
nonlocal physical_blob_objects
physical_blob_objects = builder.UnpackLogicalBlobToPhysicalBlobs(
logical_blob_object)
vm_util.LogicalRun(BuildLogical2PhysicalInstruction)
return physical_blob_objects
def OpKernelCall(opkernel_object, op_attribute, blob_register):
def BuildInstruction(builder):
with blob_register.BnInOp2BlobObjectScope(
op_attribute) as bn_in_op2blob_object:
builder.StatefulCall(
op_attribute,
opkernel_object=opkernel_object,
bn_in_op2blob_object=bn_in_op2blob_object,
)
vm_util.LogicalRun(BuildInstruction)
def _NaiveInterpret(op_attribute, parallel_conf, blob_register):
def BuildInstruction(builder):
with blob_register.BnInOp2BlobObjectScope(
op_attribute) as bn_in_op2blob_object:
builder.StatelessCall(
op_attribute,
parallel_conf,
bn_in_op2blob_object=bn_in_op2blob_object,
)
vm_util.LogicalRun(BuildInstruction)
def MakeInitialScope(job_conf, device_tag, machine_device_ids, is_mirrored):
scope = None
def BuildInitialScope(builder):
nonlocal scope
session_id = session_ctx.GetDefaultSession().id
scope = builder.BuildInitialScope(session_id, job_conf, device_tag,
machine_device_ids, is_mirrored)
vm_util.LogicalRun(BuildInitialScope)
return scope
def _EagerRunModelLoad(var_op_conf, snapshot_path):
assert isinstance(snapshot_path, str)
assert os.path.basename(snapshot_path) == "out"
snapshot_path = os.path.dirname(snapshot_path)
assert os.path.basename(snapshot_path) == var_op_conf.name
snapshot_path = os.path.dirname(snapshot_path)
path_input_op_conf, path_lbi = _GenModelIOPathInputOpConfAndRetLbi()
path_input_blob_objects = {}
(
BuildModelIOPathInputInstruction,
BuildFeedPathInstruction,
) = _MakeModelIOPathInputBuilds(path_input_op_conf, snapshot_path,
path_input_blob_objects)
model_load_op_conf, _ = _GenModelLoadOpConfAndRetLbi(var_op_conf, path_lbi)
model_load_blob_objects = oneflow_api.deprecated.BnInOp2BlobObject()
def BuildModelLoadInstruction(builder):
path_blob_object = path_input_blob_objects["out"]
model_load_blob_objects["path"] = path_blob_object
op_attribute = op_infer_util.Infer(
model_load_op_conf, ibn2blob_object=model_load_blob_objects)
parallel_conf = path_blob_object.parallel_desc_symbol.parallel_conf
cfg_op_attribute = oneflow_api.deprecated.MakeOpAttributeByString(
str(op_attribute))
builder.StatelessCall(
cfg_op_attribute,
parallel_conf,
model_load_blob_objects,
boxing_util.BoxingTo,
vm_util._FindOrCreateDelegateBlobObject,
)
sess = session_ctx.GetDefaultSession()
with scope_util.ScopeContext(scope_util.MakeScope(_BuildNotMirroredScope)):
vm_util.LogicalRun(BuildModelIOPathInputInstruction)
vm_util.LogicalRun(BuildFeedPathInstruction)
vm_util.LogicalRun(BuildModelLoadInstruction)
return model_load_blob_objects["out_0"]
def MakeScope(build_func):
scope = None
old_scope = oneflow_api.GetCurrentScope()
assert old_scope is not None
def BuildScope(builder):
nonlocal scope
scope = build_func(old_scope, builder)
assert scope is not None
vm_util.LogicalRun(BuildScope)
return scope
def MirroredCast(op_attribute, blob_register):
def BuildInstruction(builder):
with blob_register.BnInOp2BlobObjectScope(
op_attribute) as bn_in_op2blob_object:
in_blob_object = bn_in_op2blob_object["in"]
parallel_desc_symbol = in_blob_object.parallel_desc_symbol
op_arg_parallel_attr = op_arg_util.GetOpArgParallelAttribute(
parallel_desc_symbol, op_attribute, "out")
out_blob_object = builder.MakeReferenceBlobObject(
in_blob_object, op_arg_parallel_attr)
bn_in_op2blob_object["out"] = out_blob_object
vm_util.LogicalRun(BuildInstruction)
def _MakeInputBlobObject(arg_blob_def):
input_op_conf, lbi = _MakeInputOpConfAndRetLbi(arg_blob_def)
bn_in_op2blob_object = {}
def BuildInputInstruction(builder):
op_attribute = arg_blob_def.EagerAddAndInferOp(input_op_conf)
scope = oneflow.current_scope()
parallel_conf = scope.device_parallel_desc_symbol.parallel_conf
builder.StatelessCall(
op_attribute, parallel_conf, bn_in_op2blob_object=bn_in_op2blob_object
)
vm_util.LogicalRun(BuildInputInstruction)
return bn_in_op2blob_object["out"], lbi
def _Assign(var_blob_object, value_blob_object):
def BuildAssignInstruction(builder):
new_parallel_desc_symbol = boxing_util.TryReplaceDeviceTag(
builder, var_blob_object.parallel_desc_symbol, "cpu")
consumer_op_arg_parallel_attr = op_arg_util.OpArgParallelAttribute(
new_parallel_desc_symbol,
var_blob_object.op_arg_parallel_attr.sbp_parallel,
var_blob_object.op_arg_parallel_attr.opt_mirrored_parallel,
)
tmp_blob_object = boxing_util.BoxingTo(builder, value_blob_object,
consumer_op_arg_parallel_attr)
boxing_util.Assign(builder, var_blob_object, tmp_blob_object)
vm_util.LogicalRun(BuildAssignInstruction)
def _CreateEagerVariableBlob(op_attribute):
bn_in_op2blob_object = {}
def BuildInstruction(builder):
parallel_conf = oneflow.placement.current_scope().default_parallel_conf
builder.StatelessCall(op_attribute,
parallel_conf,
bn_in_op2blob_object=bn_in_op2blob_object)
vm_util.LogicalRun(BuildInstruction)
lbi = logical_blob_id_util.LogicalBlobId()
lbi.op_name = op_attribute.op_conf.name
lbi.blob_name = op_attribute.op_conf.variable_conf.out
return remote_blob_util.EagerLogicalBlob(
lbi, blob_object=bn_in_op2blob_object["out"])
def _EagerRunModelSave(var_blobs, snapshot_path):
path_input_op_conf, path_lbi = _GenModelIOPathInputOpConfAndRetLbi()
path_input_blob_objects = {}
(
BuildModelIOPathInputInstruction,
BuildFeedPathInstruction,
) = _MakeModelIOPathInputBuilds(path_input_op_conf, snapshot_path,
path_input_blob_objects)
model_save_op_conf = _GenModelSaveOpConf(var_blobs, path_lbi)
model_save_blob_objects = oneflow_api.deprecated.BnInOp2BlobObject()
def BuildModelSaveInstruction(builder):
path_blob_object = path_input_blob_objects["out"]
model_save_blob_objects["path"] = path_blob_object
for i, blob in enumerate(var_blobs):
model_save_blob_objects["in_{}".format(i)] = blob.blob_object
op_attribute = op_infer_util.Infer(
model_save_op_conf, ibn2blob_object=model_save_blob_objects)
parallel_conf = path_blob_object.parallel_desc_symbol.parallel_conf
cfg_op_attribute = oneflow_api.deprecated.MakeOpAttributeByString(
str(op_attribute))
builder.StatelessCall(
cfg_op_attribute,
parallel_conf,
model_save_blob_objects,
boxing_util.BoxingTo,
vm_util._FindOrCreateDelegateBlobObject,
)
sess = session_ctx.GetDefaultSession()
with scope_util.ScopeContext(scope_util.MakeScope(_BuildNotMirroredScope)):
vm_util.LogicalRun(BuildModelIOPathInputInstruction)
vm_util.LogicalRun(BuildFeedPathInstruction)
vm_util.LogicalRun(BuildModelSaveInstruction)
