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()