def _MakeEagerLogicalBlob(op_attribute, obn, blob_register):
lbi = op_attribute.arg_signature.bn_in_op2lbi[obn]
blob_object = blob_register.GetObject4BlobName(
"%s/%s" % (lbi.op_name, lbi.blob_name))
mirrored_sig_map = op_attribute.mirrored_signature.bn_in_op2opt_mirrored_parallel
if mirrored_sig_map[obn].HasField("mirrored_parallel"):
return remote_blob_util.EagerMirroredBlob(lbi, blob_object)
else:
return remote_blob_util.EagerConsistentBlob(lbi, blob_object)
def _Watch(op_attribute, parallel_conf, blob_register):
lbi = op_attribute.arg_signature.bn_in_op2lbi["in"]
uuid = op_attribute.op_conf.foreign_watch_conf.handler_uuid
lbn = "%s/%s" % (lbi.op_name, lbi.blob_name)
in_blob_object = blob_register.GetObject4BlobName(lbn)
if in_blob_object.op_arg_parallel_attr.is_mirrored():
blob = remote_blob_util.EagerMirroredBlob(lbi, in_blob_object)
else:
blob = remote_blob_util.EagerConsistentBlob(lbi, in_blob_object)
uuid2watch_handler = session_ctx.GetDefaultSession().uuid2watch_handler
assert uuid in uuid2watch_handler
uuid2watch_handler[uuid](blob)
del uuid2watch_handler[uuid]
def _LogicalSlice(input_blob: EagerBlobTrait, start: Sequence[int],
stop: Sequence[int]) -> np.ndarray:
"""
Construct a logical_slice op and run it by oneflow eager,
return the sliced result as a numpy ndarray
"""
op_name = id_util.UniqueStr(OP_PREFIX)
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)
lbi = logical_blob_id_util.LogicalBlobId()
lbi.op_name = op_name
lbi.blob_name = op_name
blob_object = async_util.Await(1, AsyncSlice)[0]
blob = remote_blob_util.EagerConsistentBlob(lbi,
blob_object=blob_object,
job_name=FAKE_JOB_NAME)
return blob.numpy()
def Build(builder, Yield):
blob_object = _GetInterfaceBlobObject(builder, op_name)
lbi = logical_blob_id_util.LogicalBlobId()
lbi.op_name = op_name
op_attribute = sess.OpAttribute4InterfaceOpName(op_name)
assert len(op_attribute.output_bns) == 1
lbi.blob_name = op_attribute.output_bns[0]
if blob_object.op_arg_parallel_attr.is_mirrored():
remote_blob = remote_blob_util.EagerMirroredBlob(
lbi, blob_object, job_name)
else:
remote_blob = remote_blob_util.EagerConsistentBlob(
lbi, blob_object, job_name)
Yield(remote_blob)
def CreateEagerVariableBlob(op_attribute, job_name=None):
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 = 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.EagerConsistentBlob(
lbi, blob_object=bn_in_op2blob_object["out"], job_name=job_name)
def build(builder):
blob_object = GetEagerInterfaceBlob(op_name).blob_object
lbi = logical_blob_id_util.LogicalBlobId()
lbi.op_name = op_name
op_attribute = sess.OpAttribute4InterfaceOpName(op_name)
assert len(op_attribute.output_bns) == 1
lbi.blob_name = op_attribute.output_bns[0]
if blob_object.op_arg_parallel_attr.is_mirrored():
remote_blob = remote_blob_util.EagerMirroredBlob(
lbi, blob_object, job_name)
else:
remote_blob = remote_blob_util.EagerConsistentBlob(
lbi, blob_object, job_name)
if blob_object.op_arg_blob_attr.is_tensor_list:
value = remote_blob.numpy_list()
else:
value = remote_blob.numpy()
Yield(value)
def get_blob(lbi, blob_object):
blob = remote_blob_util.EagerConsistentBlob(lbi, blob_object)
with oneflow.scope.consistent_view():
return oneflow.identity(blob)
