def _LogicalSlice( input_blob_object: oneflow_api.BlobObject, start: Sequence[int], stop: Sequence[int], scope_symbol_id: 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)) 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 = oneflow_api.deprecated.BnInOp2BlobObject() bn_in_op2blob_object["x_0"] = input_blob_object op_attribute = op_infer_util.Infer(op_conf, bn_in_op2blob_object, scope_symbol_id) cfg_op_attribute = oneflow_api.deprecated.MakeOpAttributeByString( str(op_attribute)) builder.StatelessCall( cfg_op_attribute, parallel_conf, bn_in_op2blob_object, boxing_util.BoxingTo, ) Yield(bn_in_op2blob_object["y_0"]) oneflow_api.deprecated.LogicalRun(build) lbi = lbi_util.LogicalBlobId() lbi.set_op_name(op_name) lbi.set_blob_name(op_name) blob_object = async_util.Await(1, AsyncSlice)[0] blob = oneflow_api.EagerConsistentBlob( lbi, blob_object=blob_object, blob_register=blob_register, job_name=FAKE_JOB_NAME, ) return blob.numpy()
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, ) oneflow_api.deprecated.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 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 _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 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 if mirrored_sig_map[obn].HasField("mirrored_parallel"): return oneflow_api.EagerMirroredBlob(lbi, blob_object, default_blob_register) else: return oneflow_api.EagerConsistentBlob(lbi, blob_object, default_blob_register)
def Build(builder, Yield): blob_object = _GetInterfaceBlobObject(builder, op_name) lbi = lbi_util.LogicalBlobId() lbi.set_op_name(op_name) op_attribute = sess.OpAttribute4InterfaceOpName(op_name) assert len(op_attribute.output_bns) == 1 lbi.set_blob_name(op_attribute.output_bns[0]) if blob_object.op_arg_parallel_attr.is_mirrored(): remote_blob = oneflow_api.EagerMirroredBlob( lbi, blob_object, blob_register, job_name) else: remote_blob = oneflow_api.EagerConsistentBlob( lbi, blob_object, blob_register, job_name) Yield(remote_blob)
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 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 if in_blob_object.op_arg_parallel_attr.is_mirrored(): blob = oneflow_api.EagerMirroredBlob(lbi, in_blob_object, default_blob_register) else: blob = oneflow_api.EagerConsistentBlob(lbi, in_blob_object, default_blob_register) uuid2watch_handler = session_ctx.GetDefaultSession().uuid2watch_handler assert uuid in uuid2watch_handler uuid2watch_handler[uuid](blob) del uuid2watch_handler[uuid]
def build(builder): blob_object = GetEagerInterfaceBlob(op_name).blob_object lbi = lbi_util.LogicalBlobId() lbi.set_op_name(op_name) op_attribute = sess.OpAttribute4InterfaceOpName(op_name) assert len(op_attribute.output_bns) == 1 lbi.set_blob_name(op_attribute.output_bns[0]) 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 if blob_object.op_arg_parallel_attr.is_mirrored(): remote_blob = oneflow_api.EagerMirroredBlob( lbi, blob_object, blob_register, job_name) else: remote_blob = oneflow_api.EagerConsistentBlob( lbi, blob_object, blob_register, job_name) value = remote_blob.numpy() Yield(value)
def get_blob(lbi, blob_object, blob_register): blob = oneflow_api.EagerConsistentBlob(lbi, blob_object, blob_register) with oneflow.scope.consistent_view(): return oneflow.identity(blob)