def _sort_at_last_dim(
input: remote_blob_util.BlobDef,
direction: str = "ASCENDING",
name: Optional[str] = None,
) -> remote_blob_util.BlobDef:
assert direction in ["ASCENDING", "DESCENDING"]
return (
flow.user_op_builder(name if name is not None else id_util.UniqueStr("Sort_"))
.Op("sort")
.Input("in", [input])
.Output("out")
.Attr("direction", direction)
.Build()
.InferAndTryRun()
.RemoteBlobList()[0]
)
def image_normalize(
image: BlobDef,
std: Sequence[float],
mean: Sequence[float],
name: Optional[str] = None,
) -> BlobDef:
if name is None:
name = id_util.UniqueStr("ImageNormalize_")
assert isinstance(std, (list, tuple))
assert isinstance(mean, (list, tuple))
op = (flow.user_op_builder(name).Op("image_normalize").Input(
"in", [image]).Output("out").Attr("std", std).Attr("mean",
mean).Build())
return op.InferAndTryRun().SoleOutputBlob()
def gather_nd(
params: remote_blob_util.BlobDef,
indices: remote_blob_util.BlobDef,
name: Optional[str] = None,
) -> remote_blob_util.BlobDef:
if name is None:
name = id_util.UniqueStr("GatherNd_")
op = (
flow.user_op_builder(name)
.Op("gather_nd")
.Input("params", [params])
.Input("indices", [indices])
.Output("out")
.Build()
)
return op.InferAndTryRun().RemoteBlobList()[0]
def tanh(x: remote_blob_util.BlobDef,
name: Optional[str] = None) -> remote_blob_util.BlobDef:
r"""Computes hyperbolic tangent of `x` element-wise.
Args:
x: Input `Blob`.
Returns:
A `Blob`
"""
return (
flow.user_op_builder(name if name is not None else id_util.UniqueStr(
"TanH_")).Op("tanh").Input(
"in",
[x
]).Output("out").Build().InferAndTryRun().RemoteBlobList()[0])
def CropMirrorNormalize(
input_blob: BlobDef,
mirror_blob: Optional[BlobDef] = None,
color_space: str = "BGR",
output_layout: str = "NCHW",
crop_h: int = 0,
crop_w: int = 0,
crop_pos_y: float = 0.5,
crop_pos_x: float = 0.5,
mean: Sequence[float] = [0.0],
std: Sequence[float] = [1.0],
output_dtype: dtype_util.dtype = dtype_util.float,
name: Optional[str] = None,
):
if name is None:
name = id_util.UniqueStr("CropMirrorNormalize_")
op_type_name = ""
if input_blob.dtype is dtype_util.tensor_buffer:
op_type_name = "crop_mirror_normalize_from_tensorbuffer"
elif input_blob.dtype is dtype_util.uint8:
op_type_name = "crop_mirror_normalize_from_uint8"
else:
print(
"ERROR! oneflow.data.crop_mirror_normalize op",
" NOT support input data type : ",
input_blob.dtype,
)
raise NotImplementedError
op = flow.user_op_builder(name).Op(op_type_name).Input("in", [input_blob])
if mirror_blob is not None:
op = op.Input("mirror", [mirror_blob])
return (
op.Output("out")
.Attr("color_space", color_space)
.Attr("output_layout", output_layout)
.Attr("mean", mean)
.Attr("std", std)
.Attr("crop_h", crop_h)
.Attr("crop_w", crop_w)
.Attr("crop_pos_y", crop_pos_y)
.Attr("crop_pos_x", crop_pos_x)
.Attr("output_dtype", output_dtype)
.Build()
.InferAndTryRun()
.RemoteBlobList()[0]
)
def reshape_like(
x: remote_blob_util.BlobDef,
like: remote_blob_util.BlobDef,
name: Optional[str] = None,
) -> remote_blob_util.BlobDef:
if name is None:
name = id_util.UniqueStr("ReshapeLike_")
return (
flow.user_op_builder(name)
.Op("reshape_like")
.Input("in", [x])
.Input("like", [like])
.Output("out")
.Build()
.InferAndTryRun()
.RemoteBlobList()[0]
)
def distributed_partial_fc_sample(
weight: oneflow_api.BlobDesc,
label: oneflow_api.BlobDesc,
num_sample: int,
name: Optional[str] = None,
) -> oneflow_api.BlobDesc:
parallel_num = flow.current_scope(
).device_parallel_desc_symbol.parallel_num
assert num_sample % parallel_num == 0
assert weight.shape[0] % parallel_num == 0
return (flow.user_op_builder(name if name is not None else id_util.
UniqueStr("DistributedPartialFcSample_")).
Op("distributed_partial_fc_sample").Input(
"weight", [weight]).Input("label", [label]).Attr(
"num_sample",
num_sample).Output("mapped_label").Output("sampled_label").
Output("sampled_weight").Build().InferAndTryRun().RemoteBlobList())
def tanh_grad(
y: remote_blob_util.BlobDef,
dy: remote_blob_util.BlobDef,
name: Optional[str] = None,
) -> remote_blob_util.BlobDef:
return (
flow.user_op_builder(
name if name is not None else id_util.UniqueStr("TanhGrad_")
)
.Op("tanh_grad")
.Input("y", [y])
.Input("dy", [dy])
.Output("dx")
.Build()
.InferAndTryRun()
.RemoteBlobList()[0]
)
def top_k(
input: remote_blob_util.BlobDef,
k: int = 1,
sorted: bool = True,
name: Optional[str] = None,
) -> remote_blob_util.BlobDef:
return (
flow.user_op_builder(name if name is not None else id_util.UniqueStr("TopK_"))
.Op("top_k")
.Input("in", [input])
.Output("out")
.Attr("k", k)
.Attr("sorted", sorted)
.Build()
.InferAndTryRun()
.RemoteBlobList()[0]
)
def squeeze(
input: remote_blob_util.BlobDef,
axis: Optional[Sequence[int]] = None,
name: Optional[str] = None,
) -> remote_blob_util.BlobDef:
if axis is None:
axis = [idx for idx, dim in enumerate(input.shape) if dim == 1]
else:
assert isinstance(axis, list) or isinstance(axis, tuple)
in_num_axes = len(input.shape)
for x in axis:
assert x >= -in_num_axes and x < in_num_axes
return (
flow.user_op_builder(name if name is not None else id_util.UniqueStr(
"Squeeze_")).Op("squeeze").Input("in", [input]).Output("out").Attr(
"axes",
list(axis)).Build().InferAndTryRun().RemoteBlobList()[0])
def generate_random_batch_permutation_indices(
value: remote_blob_util.BlobDef,
seed: Optional[int] = None,
name: Optional[str] = None,
) -> remote_blob_util.BlobDef:
import random
op = (flow.user_op_builder(name if name is not None else id_util.UniqueStr(
value.op_name + "_random_batch_permutation_indices")).Op(
"generate_random_batch_permutation_indices").Input(
"x", [value]).Output("y"))
if seed is not None:
op.Attr("seed", seed)
assert name is not None
else:
op.Attr("seed", random.randint(-(2**63) + 1, 2**63 - 1))
return op.Build().InferAndTryRun().RemoteBlobList()[0]
def expand_dims(
input: remote_blob_util.BlobDef, axis: int, name: Optional[str] = None
) -> remote_blob_util.BlobDef:
in_num_axes = len(input.shape)
assert axis >= -(in_num_axes + 1) and axis <= in_num_axes
return (
flow.user_op_builder(
name if name is not None else id_util.UniqueStr("ExpandDims_")
)
.Op("expand_dims")
.Input("in", [input])
.Output("out")
.Attr("axis", axis)
.Build()
.InferAndTryRun()
.RemoteBlobList()[0]
)
def write_image(value, step=None, tag=None, name=None):
r"""Write image to log file
Args:
value: A 'Blob' with dtype in 'flow.uint8'
step: A 'Blob' with 1 value and dtype is 'flow.int64'
tag: A 'Blob' with 1 value and dtype is 'flow.int8'
name: This operator's name
"""
if name is None:
name = id_util.UniqueStr("WriteImage_")
if tag is None:
tag = "image"
(flow.user_op_builder(name).Op("summary_write_image").Input(
"in", [value]).Input("step",
[step]).Input("tag",
[tag]).Build().InferAndTryRun())
def layer_norm_param_grad(
dy: remote_blob_util.BlobDef,
norm: remote_blob_util.BlobDef,
gamma: remote_blob_util.BlobDef,
begin_params_axis: int = -1,
name: str = "LayerNormParamGrad",
) -> Tuple[
remote_blob_util.BlobDef, remote_blob_util.BlobDef, remote_blob_util.BlobDef
]:
r"""Backward pass for layer normalization
Args:
dy (remote_blob_util.BlobDef): Upstream derivstives.
norm (remote_blob_util.BlobDef): Normalized output.
gamma (remote_blob_util.BlobDef): Scale parameter.
begin_params_axis (int, optional): From which parameters to begin with. Defaults to -1.
name (Optional[str], optional): This layer's name. Defaults to 'LayerNormParamGrad'.
Returns:
Tuple[remote_blob_util.BlobDef]:
normalized_diff: Gradient with respect to input `Blob`.
beta_diff: Gradient with respect to shift parameter beta.
gamma_diff: Gradient with respect to scale parameter gamma.
"""
op = (
flow.user_op_builder(name)
.Op("layer_norm_param_grad")
.Input("dy", [dy])
.Input("normalized", [norm])
.Input("gamma", [gamma])
.Output("normalized_diff")
.Output("beta_diff")
.Output("gamma_diff")
.Output("reduce_buf")
.Attr("begin_params_axis", begin_params_axis)
.Build()
)
(
normalized_diff,
beta_diff,
gamma_diff,
reduce_buf,
) = op.InferAndTryRun().RemoteBlobList()
return normalized_diff, beta_diff, gamma_diff
def argmax(input: remote_blob_util.BlobDef,
name: Optional[str] = None) -> remote_blob_util.BlobDef:
r"""Returns a 'Blob' whose value is the index with largest value accross the last axes of 'input'
Args:
input (remote_blob_util.BlobDef): a 'Blob'
name (Optional[str], optional): This operator's name. Defaults to None.
Returns:
remote_blob_util.BlobDef: A 'Blob' with type 'flow.int32'
"""
return (
flow.user_op_builder(name if name is not None else id_util.UniqueStr(
"ArgMax_")).Op("argmax").Input(
"in",
[input
]).Output("out").Build().InferAndTryRun().RemoteBlobList()[0])
def tensor_to_tensor_buffer(
x: BlobDef,
instance_dims: int,
name: Optional[str] = None,
) -> BlobDef:
r"""This operator converts the Blob's type from Tensor to TensorBuffer.
Args:
x (BlobDef): Input `Blob`.
instance_dims (int): The dimensions of dynamic tensor instance.
name (Optional[str], optional): The name for the operation. Defaults to None.
Returns:
BlobDef: The result Blob.
For example:
.. code-block:: python
import oneflow as flow
import numpy as np
import oneflow.typing as tp
@flow.global_function()
def tensor_buffer_to_tensor_Job(x: tp.Numpy.Placeholder(shape=(4, 16, 64, 64), dtype=flow.float32),
) -> tp.Numpy:
x = flow.tensor_to_tensor_buffer(x,
instance_dims=2)
return flow.tensor_buffer_to_tensor(x,
instance_shape=(64, 64),
dtype=flow.float)
x = np.random.randn(4, 16, 64, 64).astype(np.float32)
out = tensor_buffer_to_tensor_Job(x)
# out.shape (4, 16, 64, 64)
"""
if name is None:
name = id_util.UniqueStr("TensorToTensorBuffer_")
return (flow.user_op_builder(name).Op("tensor_to_tensor_buffer").Input(
"in", [x]).Output("out").Attr(
"instance_dims",
instance_dims).Build().InferAndTryRun().RemoteBlobList()[0])
def reshape(
x: remote_blob_util.BlobDef, shape: Sequence[int], name: Optional[str] = None
) -> remote_blob_util.BlobDef:
r"""Reshapes a blob.
Args:
x: A `Blob`.
shape: Shape of the output blob.
name: A name for the operation (optional).
Returns:
A `Blob`, has the same type as `x`.
"""
x = flow.cast_to_current_logical_view(x)
assert isinstance(shape, tuple) or isinstance(shape, list)
shape = list(shape)
assert all(dim == -1 or dim > 0 for dim in shape)
assert shape.count(-1) <= 1
if not x.is_dynamic:
if name is None:
name = id_util.UniqueStr("Reshape_")
return (
flow.user_op_builder(name)
.Op("reshape")
.Input("in", [x])
.Output("out")
.Attr("shape", infer_shape(x, shape))
.Build()
.InferAndTryRun()
.RemoteBlobList()[0]
)
else:
op_conf = op_conf_util.OperatorConf()
setattr(
op_conf,
"name",
name if name is not None else id_util.UniqueStr("DynamicReshape_"),
)
setattr(op_conf.dynamic_reshape_conf, "in", x.unique_name)
op_conf.dynamic_reshape_conf.shape.dim.extend(list(shape))
setattr(op_conf.dynamic_reshape_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)
def ssp_variable_proxy(
var: oneflow._oneflow_internal.BlobDesc, buffer_size: int = 1, name=None
) -> Tuple[oneflow._oneflow_internal.BlobDesc, oneflow._oneflow_internal.BlobDesc]:
r""" return ref_blob, value_blob """
if name is None:
name = id_util.UniqueStr("SspVariableProxy_")
blob_dict = (
flow.user_op_builder(name)
.Op("ssp_variable_proxy")
.Input("var", [var])
.Output("ref")
.Output("value")
.Attr("buffer_size", buffer_size)
.Build()
.InferAndTryRun()
.RemoteBlobDict()
)
return blob_dict["ref"][0], blob_dict["value"][0]
def Resize(
input_blob: BlobDef,
color_space: str = "BGR",
interp_type: str = "Linear",
resize_shorter: int = 0,
resize_x: int = 0,
resize_y: int = 0,
name: Optional[str] = None,
) -> BlobDef:
if name is None:
name = id_util.UniqueStr("ImageResize_")
return (flow.user_op_builder(name).Op("image_resize").Input(
"in", [input_blob]).Output("out").Attr(
"color_space", color_space).Attr("interp_type", interp_type).Attr(
"resize_shorter",
resize_shorter).Attr("resize_x", resize_x).Attr(
"resize_y",
resize_y).Build().InferAndTryRun().RemoteBlobList()[0])
def scatter_nd(
indices: remote_blob_util.BlobDef,
updates: remote_blob_util.BlobDef,
shape: Sequence[int],
name: Optional[str] = None,
):
if name is None:
name = id_util.UniqueStr("ScatterNd_")
op = (
flow.user_op_builder(name)
.Op("scatter_nd")
.Input("indices", [indices])
.Input("updates", [updates])
.Attr("shape", shape)
.Output("out")
.Build()
)
return op.InferAndTryRun().RemoteBlobList()[0]
def yolo_prob_loss(bbox_objness,
bbox_clsprob,
pos_inds,
pos_cls_label,
neg_inds,
valid_num,
num_classes,
name=None):
op = (flow.user_op_builder(name).Op("yolo_prob_loss").Input(
"bbox_objness",
[bbox_objness]).Input("bbox_clsprob", [bbox_clsprob]).Input(
"pos_inds",
[pos_inds]).Input("pos_cls_label", [pos_cls_label]).Input(
"neg_inds", [neg_inds]).Input("valid_num", [
valid_num
]).Output("bbox_objness_out").Output("bbox_clsprob_out").Attr(
"num_classes", num_classes).Build().InferAndTryRun())
return op.RemoteBlobList()
def tensor_scatter_nd_add(
params: remote_blob_util.BlobDef,
indices: remote_blob_util.BlobDef,
updates: remote_blob_util.BlobDef,
name: Optional[str] = None,
) -> remote_blob_util.BlobDef:
if name is None:
name = id_util.UniqueStr("TensorScatterNdAdd_")
op = (
flow.user_op_builder(name)
.Op("tensor_scatter_nd_add")
.Input("params", [params])
.Input("updates", [updates])
.Input("indices", [indices])
.Output("out")
.Build()
)
return op.InferAndTryRun().RemoteBlobList()[0]
def ofrecord_image_classification_reader(
ofrecord_dir: str,
image_feature_name: str,
label_feature_name: str,
batch_size: int = 1,
data_part_num: int = 1,
part_name_prefix: str = "part-",
part_name_suffix_length: int = -1,
random_shuffle: bool = False,
shuffle_buffer_size: int = 1024,
shuffle_after_epoch: bool = False,
color_space: str = "BGR",
decode_buffer_size_per_thread: int = 32,
num_decode_threads_per_machine: Optional[int] = None,
name: Optional[str] = None,
) -> BlobDef:
if name is None:
name = id_util.UniqueStr("OFRecordImageClassificationReader_")
(image, label) = (
flow.user_op_builder(name).Op("ofrecord_image_classification_reader").
Output("image").Output("label").Attr("data_dir", ofrecord_dir).Attr(
"data_part_num",
data_part_num).Attr("batch_size", batch_size).Attr(
"part_name_prefix",
part_name_prefix).Attr("random_shuffle", random_shuffle).Attr(
"shuffle_buffer_size", shuffle_buffer_size).Attr(
"shuffle_after_epoch", shuffle_after_epoch).Attr(
"part_name_suffix_length",
part_name_suffix_length).Attr(
"color_space",
color_space).Attr("image_feature_name",
image_feature_name).Attr(
"label_feature_name",
label_feature_name).
Attr("decode_buffer_size_per_thread",
decode_buffer_size_per_thread).Attr(
"num_decode_threads_per_machine",
num_decode_threads_per_machine
or 0).Build().InferAndTryRun().RemoteBlobList())
label = flow.tensor_buffer_to_tensor(label,
dtype=flow.int32,
instance_shape=[1])
label = flow.squeeze(label, axis=[-1])
return image, label
def gen_tensor_buffer(
shape: Sequence[int],
shape_list: Sequence[Sequence[int]],
value_list: Sequence[float],
data_type: Optional[flow.dtype] = flow.float32,
dynamic_out: Optional[bool] = False,
name: Optional[str] = None,
) -> oneflow._oneflow_internal.BlobDesc:
r"""This operator generates a tensor buffer blob.
Args:
shape (Sequence[int]): shape of output blob
shape_list ( Sequence[Sequence[int]]): shapes for tensor buffer in output blob
value_list (Sequence[float]): values for tensor buffer in output blob
data_type (Optional[flow.dtype]): data type for tensor buffer in output blob
dynamic_out (Optional[bool]): if output is a dynamic blob
name (Optional[str]): The name for the operation. Defaults to None.
Returns:
BlobDesc: The result Blob.
For example:
.. code-block:: python
import oneflow as flow
@flow.global_function(function_config=func_config)
def GenTensorBufferJob():
with flow.scope.placement("cpu", "0:0"):
x = flow.gen_tensor_buffer([(2,)], [(2, 1), (1, 2)], [0.0, 1.0])
y = flow.tensor_buffer_to_list_of_tensors(x, (100, 100), flow.float, True)
return y
# y_0.shape (2, 1), y_1.shape (1. 2)
"""
return (flow.user_op_builder(
name if name is not None else id_util.UniqueStr("GenTensorBuffer_")
).Op("gen_tensor_buffer").Output("out").Attr("shape", shape).Attr(
"shape_list", shape_list).Attr(
"value_list", value_list).Attr("data_type", data_type).Attr(
"dynamic_out",
dynamic_out).Build().InferAndTryRun().RemoteBlobList()[0])
def object_segm_poly_to_mask(poly: BlobDef,
poly_index: BlobDef,
image_size: BlobDef,
name: Optional[str] = None) -> BlobDef:
assert isinstance(poly, BlobDef)
assert isinstance(poly_index, BlobDef)
assert isinstance(image_size, BlobDef)
assert poly.shape[0] == poly_index.shape[0]
assert poly.shape[0] == image_size.shape[0]
if name is None:
name = id_util.UniqueStr("ObjectSegmPolyToMask_")
op = (flow.user_op_builder(name).Op(
"object_segmentation_polygon_to_mask").Input("poly", [poly]).Input(
"poly_index",
[poly_index]).Input("image_size",
[image_size]).Output("out").Build())
return op.InferAndTryRun().SoleOutputBlob()
def object_segm_poly_scale(
poly: BlobDef, scale: BlobDef, name: Optional[str] = None
) -> BlobDef:
assert isinstance(poly, BlobDef)
assert isinstance(scale, BlobDef)
assert poly.shape[0] == scale.shape[0]
if name is None:
name = id_util.UniqueStr("ObjectSegmPolyFilp_")
op = (
flow.user_op_builder(name)
.Op("object_segmentation_polygon_scale")
.Input("poly", [poly])
.Input("scale", [scale])
.Output("out")
.Build()
)
return op.InferAndTryRun().SoleOutputBlob()
def object_bbox_scale(
bbox: BlobDef, scale: BlobDef, name: Optional[str] = None
) -> BlobDef:
assert isinstance(bbox, BlobDef)
assert isinstance(scale, BlobDef)
assert bbox.shape[0] == scale.shape[0]
if name is None:
name = id_util.UniqueStr("ObjectBboxScale_")
op = (
flow.user_op_builder(name)
.Op("object_bbox_scale")
.Input("bbox", [bbox])
.Input("scale", [scale])
.Output("out")
.Build()
)
return op.InferAndTryRun().SoleOutputBlob()
def image_target_resize(
images: BlobDef, target_size: int, max_size: int, name: Optional[str] = None
) -> Sequence[BlobDef]:
# TODO: check target_size and max_size valid
if name is None:
name = id_util.UniqueStr("ImageTargetResize_")
op = (
flow.user_op_builder(name)
.Op("image_target_resize")
.Input("in", [images])
.Output("out")
.Output("size")
.Output("scale")
.Attr("target_size", target_size)
.Attr("max_size", max_size)
.Build()
)
return op.InferAndTryRun().RemoteBlobList()
def OFRecordImageDecoder(
input_blob: BlobDef,
blob_name: str,
color_space: str = "BGR",
name: Optional[str] = None,
) -> BlobDef:
if name is None:
name = id_util.UniqueStr("OFRecordImageDecoder_")
return (
flow.user_op_builder(name)
.Op("ofrecord_image_decoder")
.Input("in", [input_blob])
.Output("out")
.Attr("name", blob_name)
.Attr("color_space", color_space)
.Build()
.InferAndTryRun()
.RemoteBlobList()[0]
)
def unique_with_counts(
x: input_blob_util.ArgBlobDef,
out_idx: flow.dtype = flow.int32,
name: Optional[str] = None,
) -> Tuple[oneflow_api.BlobDesc]:
op = (
flow.user_op_builder(
name if name is not None else id_util.UniqueStr("UniqueWithCounts_")
)
.Op("unique_with_counts")
.Input("x", [x])
.Attr("out_idx", out_idx)
.Output("y")
.Output("idx")
.Output("count")
.Output("num_unique")
.Build()
)
return op.InferAndTryRun().RemoteBlobList()
