def CheckAndCompleteUserOpConf(op_conf_proto):
serialized_op_conf = str(text_format.MessageToString(op_conf_proto))
new_op_conf, error = oneflow_api.CheckAndCompleteUserOpConf(
serialized_op_conf)
if error.has_error_type():
raise JobBuildAndInferCfgError(error)
return text_format.Parse(new_op_conf, op_conf_util.OperatorConf())
def dynamic_binary_split(
x: input_blob_util.ArgBlobDef,
base_shift: int = 2,
out_num: int = 2,
name: Optional[str] = None,
) -> List[remote_blob_util.BlobDef]:
op_conf = op_conf_util.OperatorConf()
if name is None:
op_conf.name = id_util.UniqueStr("DynamicBinarySplit_")
else:
op_conf.name = name
obns = []
out_remote_blobs = []
for i in range(out_num):
obns.append("out_" + str(i))
setattr(op_conf.dynamic_binary_split_conf, "in", x.unique_name)
# op_conf.dynamic_binary_split_conf.in = x.unique_name
op_conf.dynamic_binary_split_conf.out[:] = obns
op_conf.dynamic_binary_split_conf.base_shift = base_shift
interpret_util.Forward(op_conf)
for i in range(out_num):
out_lbi = logical_blob_id_util.LogicalBlobId()
out_lbi.op_name = op_conf.name
out_lbi.blob_name = obns[i]
out_remote_blobs.append(remote_blob_util.RemoteBlob(out_lbi))
return out_remote_blobs
def decode_random(
shape: Sequence[int],
dtype: flow.dtype,
batch_size: int = 1,
initializer: Optional[initializer_conf_util.InitializerConf] = None,
tick: Optional[oneflow._oneflow_internal.BlobDesc] = None,
name: Optional[str] = None,
) -> oneflow._oneflow_internal.BlobDesc:
op_conf = op_conf_util.OperatorConf()
if name is None:
name = id_util.UniqueStr("DecodeRandom_")
assert isinstance(name, str)
op_conf.name = name
assert isinstance(shape, (list, tuple))
op_conf.decode_random_conf.shape.dim.extend(shape)
assert dtype is not None
setattr(
op_conf.decode_random_conf,
"data_type",
oneflow._oneflow_internal.deprecated.GetProtoDtype4OfDtype(dtype),
)
op_conf.decode_random_conf.batch_size = batch_size
if initializer is not None:
op_conf.decode_random_conf.data_initializer.CopyFrom(initializer)
else:
op_conf.decode_random_conf.data_initializer.CopyFrom(
flow.random_uniform_initializer())
if tick:
op_conf.decode_random_conf.tick = tick.unique_name
op_conf.decode_random_conf.out = "out"
lbi = logical_blob_id_util.LogicalBlobId()
lbi.op_name = op_conf.name
lbi.blob_name = "out"
interpret_util.ConsistentForward(op_conf)
return remote_blob_util.RemoteBlob(lbi)
def elem_cnt(
input_blob: remote_blob_util.BlobDef,
axis: Optional[Sequence[int]] = None,
dtype: Optional[dtype_util.dtype] = None,
name: Optional[str] = None,
) -> remote_blob_util.BlobDef:
op_conf = op_conf_util.OperatorConf()
setattr(
op_conf,
"name",
name if name is not None else id_util.UniqueStr("ShapeElemCnt_"),
)
op_conf.shape_elem_cnt_conf.x = input_blob.unique_name
if axis is None:
op_conf.shape_elem_cnt_conf.exclude_axis_conf.SetInParent()
else:
assert isinstance(axis, (tuple, list))
op_conf.shape_elem_cnt_conf.include_axis_conf.axis.extend(axis)
if dtype is not None:
op_conf.shape_elem_cnt_conf.data_type = dtype.oneflow_proto_dtype
op_conf.shape_elem_cnt_conf.y = "y"
interpret_util.Forward(op_conf)
out_lbi = logical_blob_id_util.LogicalBlobId()
out_lbi.op_name = op_conf.name
out_lbi.blob_name = "y"
return remote_blob_util.RemoteBlob(out_lbi)
def indexed_slices_reduce_sum(
indices: input_blob_util.ArgBlobDef,
values: input_blob_util.ArgBlobDef,
name: Optional[str] = None,
) -> Tuple[remote_blob_util.BlobDef]:
op_conf = op_conf_util.OperatorConf()
if name is None:
op_conf.name = id_util.UniqueStr("IndexedSlicesReduceSum_")
else:
op_conf.name = name
op_conf.indexed_slices_reduce_sum_conf.x_indices = indices.unique_name
op_conf.indexed_slices_reduce_sum_conf.x_values = values.unique_name
op_conf.indexed_slices_reduce_sum_conf.y_indices = "y_indices"
op_conf.indexed_slices_reduce_sum_conf.y_values = "y_values"
op_conf.indexed_slices_reduce_sum_conf.num_unique = "num_unique"
interpret_util.Forward(op_conf)
y_indices_lbi = logical_blob_id_util.LogicalBlobId()
y_indices_lbi.op_name = op_conf.name
y_indices_lbi.blob_name = "y_indices"
y_values_lbi = logical_blob_id_util.LogicalBlobId()
y_values_lbi.op_name = op_conf.name
y_values_lbi.blob_name = "y_values"
num_unique_lbi = logical_blob_id_util.LogicalBlobId()
num_unique_lbi.op_name = op_conf.name
num_unique_lbi.blob_name = "num_unique"
return (
remote_blob_util.RemoteBlob(y_indices_lbi),
remote_blob_util.RemoteBlob(y_values_lbi),
remote_blob_util.RemoteBlob(num_unique_lbi),
)
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 = oneflow_api.deprecated.BnInOp2BlobObject()
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)
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,
)
Yield(bn_in_op2blob_object["y_0"])
def tensor_buffer_to_tensor_list(
input: remote_blob_util.BlobDef,
shape: Sequence[int],
dtype: dtype_util.dtype,
name: Optional[str] = None,
) -> remote_blob_util.BlobDef:
if name is None:
name = id_util.UniqueStr("TensorBufferToList_")
op_conf = op_conf_util.OperatorConf()
setattr(op_conf, "name", name)
setattr(op_conf.tensor_buffer_to_tensor_list_conf, "in", input.unique_name)
setattr(op_conf.tensor_buffer_to_tensor_list_conf, "out", "out")
op_conf.tensor_buffer_to_tensor_list_conf.shape.dim[:] = list(shape)
setattr(
op_conf.tensor_buffer_to_tensor_list_conf,
"data_type",
dtype.oneflow_proto_dtype,
)
interpret_util.Forward(op_conf)
lbi = logical_blob_id_util.LogicalBlobId()
lbi.op_name = op_conf.name
lbi.blob_name = "out"
return remote_blob_util.RemoteBlob(lbi)
def BuildAssignInstruction(builder):
op_conf = op_conf_pb.OperatorConf()
# device_tag doesn't matter for logical_slice_assign op
device_tag = oneflow.current_scope(
).device_parallel_desc_symbol.device_tag
op_conf.device_tag = device_tag
op_name = id_util.UniqueStr(OP_PREFIX)
op_conf.name = op_name
op_conf.user_conf.op_type_name = "logical_slice_assign"
op_conf.user_conf.input["value"].s.append("{}/value_0".format(op_name))
op_conf.user_conf.input["ref"].s.append("{}/ref_0".format(op_name))
parallel_conf = ref_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["ref_0"] = ref_blob_object
bn_in_op2blob_object["value_0"] = value_blob_object
scope_symbol_id = _GetScopeSymbolIdFromEagerBlob(ref_blob)
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,
vm_util._FindOrCreateDelegateBlobObject,
)
def parallel_cast(input, name=None, distribute=None, gradient_distribute=None):
assert not oneflow.eager_execution_enabled()
op_conf = op_conf_util.OperatorConf()
setattr(
op_conf,
"name",
name if name is not None else id_util.UniqueStr("ParallelCast_"),
)
op_conf.parallel_cast_conf.out = "out"
setattr(op_conf.parallel_cast_conf, "in", input.unique_name)
def to_split_axis(dist):
split_axis = data_type_util.OptInt64()
if type(dist) is distribute_util.SplitDistribute:
split_axis.value = dist.axis
elif type(dist) is distribute_util.BroadcastDistribute:
split_axis.ClearField("value")
else:
raise NotImplementedError
return split_axis
if distribute is not None:
op_conf.parallel_cast_conf.split_axis.CopyFrom(to_split_axis(distribute))
if gradient_distribute is not None:
op_conf.parallel_cast_conf.gradient_split_axis.CopyFrom(
to_split_axis(gradient_distribute)
)
compile_context.CurJobAddOp(op_conf)
lbi = logical_blob_id_util.LogicalBlobId()
lbi.op_name = op_conf.name
lbi.blob_name = "out"
return remote_blob_util.RemoteBlob(lbi)
def broadcast_to_compatible_with(
x: remote_blob_util.BlobDef,
compatible: Sequence[remote_blob_util.BlobDef],
name: Optional[str] = None,
) -> remote_blob_util.BlobDef:
r"""Returns a 'Blob' with the shape can be broadcasted by other shapes
Args:
x (remote_blob_util.BlobDef): a 'Blob'
compatible (Sequence[remote_blob_util.BlobDef]): Sequence of different shape
name (Optional[str], optional): This operator's name. Defaults to None.
Returns:
remote_blob_util.BlobDef: A 'Blob' with the biggest shape
"""
assert isinstance(compatible, (list, tuple))
if name is None:
name = id_util.UniqueStr("BroadcastToCompatibleWith_")
op_conf = op_conf_util.OperatorConf()
setattr(op_conf, "name", name)
setattr(op_conf.broadcast_to_compatible_with_conf, "x", x.unique_name)
setattr(op_conf.broadcast_to_compatible_with_conf, "y", "y")
op_conf.broadcast_to_compatible_with_conf.compatible.extend(
[cp.unique_name for cp in compatible])
interpret_util.Forward(op_conf)
ret_lbi = logical_blob_id_util.LogicalBlobId()
ret_lbi.op_name = op_conf.name
ret_lbi.blob_name = "y"
return remote_blob_util.RemoteBlob(ret_lbi)
def ofrecord_loader(
ofrecord_dir: str,
batch_size: int = 1,
data_part_num: int = 1,
part_name_prefix: str = "part-",
part_name_suffix_length: int = -1,
shuffle: bool = False,
shuffle_buffer_size: int = 1024,
name: Optional[str] = None,
) -> remote_blob_util.BlobDef:
if name is None:
name = id_util.UniqueStr("OFRecord_Loader_")
op_conf = op_conf_util.OperatorConf()
op_conf.name = name
op_conf.record_load_conf.out = "out"
op_conf.record_load_conf.data_dir = ofrecord_dir
op_conf.record_load_conf.data_part_num = data_part_num
op_conf.record_load_conf.batch_size = batch_size
op_conf.record_load_conf.part_name_prefix = part_name_prefix
if part_name_suffix_length != -1:
op_conf.record_load_conf.part_name_suffix_length = part_name_suffix_length
if shuffle:
op_conf.record_load_conf.random_shuffle_conf.buffer_size = shuffle_buffer_size
lbi = logical_blob_id_util.LogicalBlobId()
lbi.op_name = name
lbi.blob_name = "out"
interpret_util.ConsistentForward(op_conf)
return remote_blob_util.RemoteBlob(lbi)
def ConstructNaiveBoxingOpConf(
produced_blob_object,
consumer_op_arg_parallel_attr,
in_parallel_num,
out_parallel_num,
):
op_conf = op_conf_pb.OperatorConf()
op_conf.name = "undefined_boxing_op_name"
op_conf.device_tag = "cpu"
op_conf.boxing_conf.lbi.op_name = "undefined_boxing_op_name"
op_conf.boxing_conf.lbi.blob_name = "undefined_boxing_blob_name"
op_conf.boxing_conf.in_num = in_parallel_num
op_conf.boxing_conf.out_num = out_parallel_num
in_sbp_parallel = produced_blob_object.op_arg_parallel_attr.sbp_parallel
if in_sbp_parallel.HasField("split_parallel"):
op_conf.boxing_conf.concat_box.axis = in_sbp_parallel.split_parallel.axis
elif in_parallel_num == 1:
op_conf.boxing_conf.concat_box.axis = 0
else:
assert in_sbp_parallel.HasField("partial_sum_parallel")
op_conf.boxing_conf.add_box.SetInParent()
out_sbp_parallel = consumer_op_arg_parallel_attr.sbp_parallel
if out_sbp_parallel.HasField("split_parallel"):
out_axis = out_sbp_parallel.split_parallel.axis
else:
assert out_parallel_num == 1
out_axis = 0
op_conf.boxing_conf.split_box.axis = out_axis
shape = produced_blob_object.op_arg_blob_attr.shape
op_conf.boxing_conf.split_box.part_num.extend(
balanced_splitter.BalancedPartNums(shape[out_axis], out_parallel_num))
bn_in_op2blob_object = {("in_%s" % i): produced_blob_object
for i in range(in_parallel_num)}
return op_infer_util.Infer(op_conf, bn_in_op2blob_object)
def BuildAssignInstruction(builder):
op_conf = op_conf_pb.OperatorConf()
# device_tag doesn't matter for logical_slice_assign op
device_tag = oneflow.current_scope(
).device_parallel_desc_symbol.device_tag
op_conf.device_tag = device_tag
op_name = id_util.UniqueStr(OP_PREFIX)
op_conf.name = op_name
op_conf.user_conf.op_type_name = "logical_slice_assign"
op_conf.user_conf.input["value"].s.append("{}/value_0".format(op_name))
op_conf.user_conf.input["ref"].s.append("{}/ref_0".format(op_name))
parallel_conf = ref_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 = {
"ref_0": ref_blob_object,
"value_0": value_blob_object
}
scope_symbol_id = _GetScopeSymbolIdFromEagerBlob(ref_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,
)
def constant_like(
like: remote_blob_util.BlobDef,
value: Union[int, float],
dtype: Optional[dtype_util.dtype] = None,
name: Optional[str] = None,
) -> remote_blob_util.BlobDef:
op_conf = op_conf_util.OperatorConf()
setattr(
op_conf,
"name",
name if name is not None else id_util.UniqueStr("ConstantLike_"),
)
setattr(op_conf.constant_like_conf, "like", like.unique_name)
if isinstance(value, int):
op_conf.constant_like_conf.int_operand = value
elif isinstance(value, float):
op_conf.constant_like_conf.float_operand = value
else:
raise NotImplementedError
if dtype is not None:
setattr(op_conf.constant_like_conf, "data_type",
dtype.oneflow_proto_dtype)
setattr(op_conf.constant_like_conf, "out", "out")
interpret_util.Forward(op_conf)
out_lbi = logical_blob_id_util.LogicalBlobId()
setattr(out_lbi, "op_name", op_conf.name)
setattr(out_lbi, "blob_name", "out")
return remote_blob_util.RemoteBlob(out_lbi)
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"])
def build(builder):
op_conf = op_conf_pb.OperatorConf()
device_tag = flow.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._oneflow_internal.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._oneflow_internal.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"])
def argwhere(
condition: remote_blob_util.BlobDef,
dtype: Optional[dtype_util.dtype] = None,
name: Optional[str] = None,
) -> remote_blob_util.BlobDef:
if name is None:
name = id_util.UniqueStr("ArgWhere_")
op_conf = op_conf_util.OperatorConf()
setattr(op_conf, "name", name)
setattr(op_conf.arg_where_conf, "in", condition.unique_name)
setattr(op_conf.arg_where_conf, "out", "out")
setattr(op_conf.arg_where_conf, "out_size", "out_size")
if dtype is not None:
setattr(op_conf.arg_where_conf, "data_type", dtype.oneflow_proto_dtype)
interpret_util.Forward(op_conf)
arg_where_out_lbi = logical_blob_id_util.LogicalBlobId()
setattr(arg_where_out_lbi, "op_name", op_conf.name)
setattr(arg_where_out_lbi, "blob_name", "out")
arg_where_out_size_lbi = logical_blob_id_util.LogicalBlobId()
setattr(arg_where_out_size_lbi, "op_name", op_conf.name)
setattr(arg_where_out_size_lbi, "blob_name", "out_size")
arg_where_out = remote_blob_util.RemoteBlob(arg_where_out_lbi)
arg_where_out_size = remote_blob_util.RemoteBlob(arg_where_out_size_lbi)
return sync_dynamic_resize(arg_where_out, arg_where_out_size)
def InputOpByArgBlobDef(blob_def):
assert isinstance(blob_def, input_blob_util.ArgBlobDef)
op_conf = op_conf_util.OperatorConf()
op_conf.name = blob_def.op_name
op_conf.input_conf.out = blob_def.blob_name
op_conf.input_conf.blob_conf.CopyFrom(blob_def.ToInterfaceBlobConf())
blob_def.AddAndInferOp(op_conf)
return remote_blob_util.RemoteBlob(blob_def.lbi)
def CheckAndCompleteUserOpConf(op_conf_proto):
serialized_op_conf = str(text_format.MessageToString(op_conf_proto))
new_op_conf, error_str = oneflow_internal.CheckAndCompleteUserOpConf(
serialized_op_conf)
error = text_format.Parse(error_str, error_util.ErrorProto())
if error.HasField("error_type"):
raise JobBuildAndInferError(error)
return text_format.Parse(new_op_conf, op_conf_util.OperatorConf())
def _AssignOpConf():
op_conf = op_conf_pb.OperatorConf()
op_conf.name = "assign"
op_conf.assign_conf.ref = "assign/ref"
op_conf.assign_conf.value = "assign/value"
device_tag = oneflow.current_scope().device_parallel_desc_symbol.device_tag
op_conf.device_tag = device_tag
return op_conf
def __init__(self, op_name, op_type_name=None):
self.op_conf_ = op_conf_util.OperatorConf()
self.op_conf_.name = op_name
if op_type_name is not None:
self.op_conf_.user_conf.op_type_name = op_type_name
device_tag = oneflow.current_scope().device_parallel_desc_symbol.device_tag
self.op_conf_.device_tag = device_tag
self.output_arg_key_list_ = []
def _SystemAssignOpConf(ref, value, name=None):
if name is None:
name = id_util.UniqueStr("Assign_")
op_conf = op_conf_util.OperatorConf()
op_conf.name = name
op_conf.assign_conf.ref = ref.unique_name
op_conf.assign_conf.value = value.unique_name
return op_conf
def tensor_list_split(
input_tensor_list: oneflow_api.BlobDesc,
name: Optional[str] = None) -> Tuple[oneflow_api.BlobDesc]:
"""This operator splits the input `TensorList`.
Args:
input_tensor_list (oneflow_api.BlobDesc): The input `TensorList`.
name (Optional[str], optional): The name for the operation. Defaults to None.
Returns:
Tuple[oneflow_api.BlobDesc]: A Tuple of `ListNumpy`.
For example:
.. code-block:: python
import oneflow as flow
import numpy as np
import oneflow.typing as tp
from typing import Tuple
func_config = flow.FunctionConfig()
func_config.default_data_type(flow.float)
func_config.default_logical_view(flow.scope.mirrored_view())
@flow.global_function(function_config=func_config)
def tensorList_split_Job(x: tp.ListListNumpy.Placeholder(shape=(2, 5, 4), dtype=flow.float32),
) -> Tuple[tp.ListNumpy, tp.ListNumpy]:
return flow.tensor_list_split(x)
x = np.random.rand(1, 3, 2).astype(np.float32)
y = np.random.rand(1, 2, 2).astype(np.float32)
out = tensorList_split_Job([[x, y]])
# out[0][0].shape (3, 2)
# out[1][0].shape (2, 2)
"""
if name is None:
name = id_util.UniqueStr("TensorListSplit_")
output_size = input_tensor_list.shape[0]
op_conf = op_conf_util.OperatorConf()
setattr(op_conf, "name", name)
setattr(op_conf.tensor_list_split_conf, "in",
input_tensor_list.unique_name)
op_conf.tensor_list_split_conf.out.extend(
["out_{}".format(i) for i in range(output_size)])
interpret_util.Forward(op_conf)
ret = []
for i in range(output_size):
out_lbi = logical_blob_id_util.LogicalBlobId()
setattr(out_lbi, "op_name", op_conf.name)
setattr(out_lbi, "blob_name", "out_{}".format(i))
ret.append(remote_blob_util.RemoteBlob(out_lbi))
return tuple(ret)
def _GetEagerNcclAllReduce(parallel_conf, ibn2blob_object):
op_conf = op_conf_pb.OperatorConf()
op_conf.device_tag = "gpu"
op_conf.name = "eager_nccl_all_reduce"
op_conf.user_conf.op_type_name = "eager_nccl_all_reduce"
op_conf.user_conf.input["in"].s.append("eager_nccl_all_reduce/in_0")
op_conf.user_conf.output["out"].s.append("eager_nccl_all_reduce/out_0")
op_conf.user_conf.attr["parallel_conf"].at_string = str(parallel_conf)
return op_infer_util.Infer(op_conf, ibn2blob_object)
def _GenModelSaveOpConf(var_blobs, path_lbi):
op_conf = op_conf_util.OperatorConf()
op_conf.name = "model_save"
op_conf.device_tag = "cpu"
op_conf.model_save_conf.path = "{}/{}".format(path_lbi.op_name, path_lbi.blob_name)
for blob in var_blobs:
getattr(op_conf.model_save_conf, "in").append(blob.logical_blob_name)
getattr(op_conf.model_save_conf, "key").append(blob.logical_blob_name)
return op_conf
def _MakeCopyHdOpConfAndRetLbi():
op_conf = op_conf_pb.OperatorConf()
op_conf.name = "copy_hd"
op_conf.device_tag = "gpu"
setattr(op_conf.copy_conf, "in", "%s/in" % op_conf.name)
op_conf.copy_conf.out = "out"
lbi = logical_blob_id_util.LogicalBlobId()
lbi.op_name = op_conf.name
lbi.blob_name = "out"
return op_conf, lbi
def LazyConsistentWatch(blob_watched, handler):
handler_uuid = str(uuid.uuid1())
op_conf = op_conf_util.OperatorConf()
op_conf.name = id_util.UniqueStr("ForeignWatch_")
setattr(op_conf.foreign_watch_conf, "in", blob_watched.unique_name)
op_conf.foreign_watch_conf.handler_uuid = handler_uuid
device_name = blob_watched.parallel_conf.device_name(0)
with oneflow.scope.placement("cpu", "0:0"):
compile_context.CurJobAddOp(op_conf)
watcher_util.BindUuidAndHandler(handler_uuid, blob_watched, handler)
def _MakeInputOpConfAndRetLbi(arg_blob_def):
assert isinstance(arg_blob_def, input_blob_def.ArgBlobDef)
op_conf = op_conf_util.OperatorConf()
op_conf.name = id_util.UniqueStr("Input_")
op_conf.input_conf.out = "out"
op_conf.input_conf.blob_conf.CopyFrom(arg_blob_def.ToInterfaceBlobConf())
lbi = logical_blob_id_util.LogicalBlobId()
lbi.op_name = op_conf.name
lbi.blob_name = op_conf.input_conf.out
return op_conf, lbi
def GenerateVariableOpConf(
name,
shape,
dtype=None,
initializer=None,
regularizer=None,
trainable=None,
model_name=None,
random_seed=None,
distribute=oneflow_api.distribute.broadcast(),
):
op_conf = op_conf_util.OperatorConf()
op_conf.name = name
op_conf.variable_conf.shape.dim.extend(shape)
assert dtype is not None
op_conf.variable_conf.data_type = oneflow_api.deprecated.GetProtoDtype4OfDtype(
dtype)
if rt_mode.CurrentMode() == rt_mode.NORMAL_MODE:
root_path = None
else:
root_path = (compile_context.GetCurJobConfigProto().
default_initialize_with_snapshot_path())
dir_path = os.path.join(root_path, name)
file_path = os.path.join(dir_path, "out")
if root_path and os.path.isfile(file_path):
op_conf.variable_conf.initialize_with_snapshot.path = dir_path
op_conf.variable_conf.initialize_with_snapshot.key = "out"
else:
if root_path:
print("{} not found, will be initialized".format(file_path))
if initializer is not None:
op_conf.variable_conf.initializer.CopyFrom(initializer)
if regularizer is not None:
op_conf.variable_conf.regularizer.CopyFrom(regularizer)
if trainable is not None:
op_conf.variable_conf.trainable = trainable
if model_name is not None:
op_conf.variable_conf.model_name = model_name
if type(distribute) is oneflow_api.distribute.SplitDistribute:
op_conf.variable_conf.split_axis.value = distribute.axis
else:
op_conf.variable_conf.split_axis.ClearField("value")
if random_seed is not None:
op_conf.variable_conf.random_seed = random_seed
op_conf.variable_conf.out = "out"
return op_conf
def tensor_list_to_tensor_buffer(
input: oneflow_api.BlobDesc,
name: Optional[str] = None) -> oneflow_api.BlobDesc:
"""This operator converts `TensorList` to `TensorBuffer`.
Refer to `Concept Explanation <https://docs.oneflow.org/basics_topics/concept_explanation.html#3tensorbuffer-tensorlist>`_
for more about TensorList.
Args:
input (oneflow_api.BlobDesc): The input `TensorList`.
name (Optional[str], optional): The name for the operation. Defaults to None.
Returns:
oneflow_api.BlobDesc: The result Blob.
For example:
.. code-block:: python
import oneflow as flow
import numpy as np
import oneflow.typing as tp
func_config = flow.FunctionConfig()
func_config.default_data_type(flow.float)
func_config.default_logical_view(flow.scope.mirrored_view())
@flow.global_function(function_config=func_config)
def tensorList_to_tensorBuffer_Job(x: tp.ListListNumpy.Placeholder(shape=(2, 5, 4), dtype=flow.float32),
) -> tp.ListListNumpy:
x = flow.tensor_list_to_tensor_buffer(input=x)
return flow.tensor_buffer_to_tensor_list(x,
shape=(5, 4),
dtype=flow.float32)
x = np.random.rand(1, 3, 2).astype(np.float32)
y = np.random.rand(1, 2, 2).astype(np.float32)
out = tensorList_to_tensorBuffer_Job([[x, y]])
# out[0][0].shape (1, 3, 2)
"""
if name is None:
name = id_util.UniqueStr("TensorListToBuffer_")
op_conf = op_conf_util.OperatorConf()
setattr(op_conf, "name", name)
setattr(op_conf.tensor_list_to_tensor_buffer_conf, "in", input.unique_name)
setattr(op_conf.tensor_list_to_tensor_buffer_conf, "out", "out")
interpret_util.Forward(op_conf)
lbi = logical_blob_id_util.LogicalBlobId()
lbi.op_name = op_conf.name
lbi.blob_name = "out"
return remote_blob_util.RemoteBlob(lbi)
