def unique_with_counts( x: input_blob_util.ArgBlobDef, out_idx: dtype_util.dtype = dtype_util.int32, 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("UniqueWithCounts_") else: op_conf.name = name op_conf.unique_with_counts_conf.x = x.unique_name op_conf.unique_with_counts_conf.y = "y" op_conf.unique_with_counts_conf.idx = "idx" op_conf.unique_with_counts_conf.count = "count" op_conf.unique_with_counts_conf.num_unique = "num_unique" op_conf.unique_with_counts_conf.out_idx = out_idx.oneflow_proto_dtype interpret_util.Forward(op_conf) y_lbi = logical_blob_id_util.LogicalBlobId() y_lbi.op_name = op_conf.name y_lbi.blob_name = "y" idx_lbi = logical_blob_id_util.LogicalBlobId() idx_lbi.op_name = op_conf.name idx_lbi.blob_name = "idx" count_lbi = logical_blob_id_util.LogicalBlobId() count_lbi.op_name = op_conf.name count_lbi.blob_name = "count" 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_lbi), remote_blob_util.RemoteBlob(idx_lbi), remote_blob_util.RemoteBlob(count_lbi), remote_blob_util.RemoteBlob(num_unique_lbi), )
def distribute_add(xs, name=None): assert oneflow.placement.current_scope().parallel_size == len(xs) if name is None: name = id_util.UniqueStr("DistributeAdd_") op_conf = op_conf_util.OperatorConf() op_conf.name = name getattr(op_conf.distribute_add_conf, "in").extend([_SoleConsistentLbn(x) for x in xs]) op_conf.distribute_add_conf.out = "out" interpret_util.ConsistentForward(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 __init__(self, shape, dtype, batch_axis, name=None): lbi = lbi_util.LogicalBlobId() if name is None: name = id_util.UniqueStr("Input_") lbi.op_name = name lbi.blob_name = "out" blob_desc.BlobDesc.__init__(self, lbi) assert type(shape) is tuple for dim in shape: assert type(dim) is int assert dim > 0 self.shape_ = shape self.dtype_ = dtype self.batch_axis_ = batch_axis
def image_decoder_random_crop_resize( input_blob: oneflow_api.BlobDesc, target_width: int, target_height: int, num_attempts: Optional[int] = None, seed: Optional[int] = None, random_area: Optional[Sequence[float]] = None, random_aspect_ratio: Optional[Sequence[float]] = None, num_workers: Optional[int] = None, warmup_size: Optional[int] = None, max_num_pixels: Optional[int] = None, name: Optional[str] = None, ) -> Tuple[oneflow_api.BlobDesc]: if name is None: name = id_util.UniqueStr("ImageDecoderRandomCropResize_") op_conf = op_conf_util.OperatorConf() op_conf.name = name setattr(op_conf.image_decoder_random_crop_resize_conf, "in", input_blob.unique_name) op_conf.image_decoder_random_crop_resize_conf.out = "out" op_conf.image_decoder_random_crop_resize_conf.target_width = target_width op_conf.image_decoder_random_crop_resize_conf.target_height = target_height if num_attempts is not None: op_conf.image_decoder_random_crop_resize_conf.num_attempts = num_attempts if seed is not None: op_conf.image_decoder_random_crop_resize_conf.seed = seed if random_area is not None: assert len(random_area) == 2 op_conf.image_decoder_random_crop_resize_conf.random_area_min = random_area[ 0] op_conf.image_decoder_random_crop_resize_conf.random_area_max = random_area[ 1] if random_aspect_ratio is not None: assert len(random_aspect_ratio) == 2 op_conf.image_decoder_random_crop_resize_conf.random_aspect_ratio_min = random_aspect_ratio[ 0] op_conf.image_decoder_random_crop_resize_conf.random_aspect_ratio_max = random_aspect_ratio[ 1] if num_workers is not None: op_conf.image_decoder_random_crop_resize_conf.num_workers = num_workers if warmup_size is not None: op_conf.image_decoder_random_crop_resize_conf.warmup_size = warmup_size if max_num_pixels is not None: op_conf.image_decoder_random_crop_resize_conf.max_num_pixels = max_num_pixels 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 unpack(input, unpack_num, name=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("Unpack_"), ) setattr(op_conf.unpack_conf, "in", input.unique_name) op_conf.unpack_conf.out = "out" op_conf.unpack_conf.unpack_num = unpack_num 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 _GenModelIOPathInputOpConfAndRetLbi(): op_conf = op_conf_util.OperatorConf() op_conf.name = "model_io_path_input" op_conf.device_tag = "cpu" op_conf.input_conf.out = "out" blob_conf = inter_face_blob_conf_util.InterfaceBlobConf() blob_conf.shape.dim.append(65536) blob_conf.data_type = oneflow._oneflow_internal.deprecated.GetProtoDtype4OfDtype( flow.int8) blob_conf.is_dynamic = True op_conf.input_conf.blob_conf.CopyFrom(blob_conf) 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 _GenModelLoadOpConfAndRetLbi(var_op_conf, path_lbi): variable_op_conf = op_conf_util.VariableOpConf() variable_op_conf.CopyFrom(var_op_conf.variable_conf) op_conf = op_conf_util.OperatorConf() op_conf.name = "model_load" op_conf.device_tag = "cpu" op_conf.model_load_conf.path = "{}/{}".format(path_lbi.op_name, path_lbi.blob_name) op_conf.model_load_conf.out.append("out_0") op_conf.model_load_conf.variable_op_name.append(var_op_conf.name) op_conf.model_load_conf.original_variable_conf.append(variable_op_conf) lbi = logical_blob_id_util.LogicalBlobId() lbi.op_name = op_conf.name lbi.blob_name = op_conf.model_load_conf.out[0] return (op_conf, lbi)
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): bn_in_op2blob_object = {} def BuildInstruction(builder): parallel_conf = oneflow.placement.current_scope().default_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.EagerLogicalBlob( lbi, blob_object=bn_in_op2blob_object["out"])
def distribute_concat(xs, axis=0, name=None): assert oneflow.current_scope( ).device_parallel_desc_symbol.parallel_num == len(xs) if name is None: name = id_util.UniqueStr("DistributeConcat_") op_conf = op_conf_util.OperatorConf() op_conf.name = name getattr(op_conf.distribute_concat_conf, "in").extend([_SoleConsistentLbn(x) for x in xs]) op_conf.distribute_concat_conf.axis = axis op_conf.distribute_concat_conf.out = "out" interpret_util.ConsistentForward(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 _GenModelLoadPathInputOpConfAndRetLbi(): op_conf = op_conf_util.OperatorConf() op_conf.name = "model_load_path_input" op_conf.input_conf.out = "out" blob_conf = op_conf_util.InterfaceBlobConf() blob_conf.shape.dim.append(65536) blob_conf.data_type = dtype_util.int8.oneflow_proto_dtype blob_conf.batch_axis.value = 0 blob_conf.is_dynamic = True op_conf.input_conf.blob_conf.CopyFrom(blob_conf) 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 acc(one, max_acc_num, name=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("Acc_"), ) op_conf.acc_conf.one = one.unique_name op_conf.acc_conf.acc = "acc" op_conf.acc_conf.max_acc_num = max_acc_num compile_context.CurJobAddOp(op_conf) lbi = logical_blob_id_util.LogicalBlobId() lbi.op_name = op_conf.name lbi.blob_name = "acc" return remote_blob_util.RemoteBlob(lbi)
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 tensor_list_to_tensor_buffer( input: remote_blob_util.BlobDef, name: Optional[str] = None) -> remote_blob_util.BlobDef: 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)
def square_sum( x: input_blob_util.ArgBlobDef, name: Optional[str] = None ) -> remote_blob_util.BlobDef: op_conf = op_conf_util.OperatorConf() if name is None: op_conf.name = id_util.UniqueStr("SquareSum_") else: op_conf.name = name op_conf.square_sum_conf.x = x.unique_name op_conf.square_sum_conf.y = "y" interpret_util.Forward(op_conf) lbi = logical_blob_id_util.LogicalBlobId() lbi.op_name = op_conf.name lbi.blob_name = "y" return remote_blob_util.RemoteBlob(lbi)
def decode_ofrecord( ofrecord_dir: str, blobs: Sequence[BlobConf], batch_size: int = 1, data_part_num: int = 1, part_name_prefix: str = "part-", part_name_suffix_length: int = -1, shuffle: bool = False, buffer_size: int = 1024, name: str = None, ) -> Tuple[remote_blob_util.BlobDef]: print( "WARNING:", "oneflow.data.decode_ofrecord is deprecated, and NOT work in eager mode, please use: \n", " 1) ofrecord = oneflow.data.ofrecord_reader(...) to read ofrecord; \n", " 2) image = oneflow.data.ofrecord_image_decoder(...) to decode image; \n", " 3) raw = oneflow.data.ofrecord_raw_decoder(...) to decode raw data like label; \n", traceback.format_stack()[-2], ) assert not flow.eager_execution_enabled() if name is None: name = id_util.UniqueStr("Decode_") lbis = [] op_conf = op_conf_util.OperatorConf() op_conf.name = name op_conf.decode_ofrecord_conf.data_dir = ofrecord_dir op_conf.decode_ofrecord_conf.data_part_num = data_part_num op_conf.decode_ofrecord_conf.batch_size = batch_size op_conf.decode_ofrecord_conf.part_name_prefix = part_name_prefix op_conf.decode_ofrecord_conf.part_name_suffix_length = part_name_suffix_length if shuffle == True: op_conf.decode_ofrecord_conf.random_shuffle_conf.buffer_size = buffer_size for blob_conf in blobs: op_conf.decode_ofrecord_conf.blob.extend([blob_conf.to_proto()]) lbi = logical_blob_id_util.LogicalBlobId() lbi.op_name = name lbi.blob_name = blob_conf.name lbis.append(lbi) interpret_util.ConsistentForward(op_conf) return tuple(map(lambda x: remote_blob_util.RemoteBlob(x), lbis))
def dynamic_binary_concat( input_blob_list: Sequence[remote_blob_util.BlobDef], source_blob: input_blob_util.ArgBlobDef, source_sbp: str = "S:0", name: Optional[str] = None, ) -> remote_blob_util.BlobDef: op_conf = op_conf_util.OperatorConf() if name is None: op_conf.name = id_util.UniqueStr("DynamicBinaryConcat_") else: op_conf.name = name in_lbns = [] for in_blob in input_blob_list: in_lbns.append(in_blob.unique_name) getattr(op_conf.dynamic_binary_concat_conf, "in").extend(in_lbns) # op_conf.dynamic_binary_concat_conf.in[:] = in_lbns op_conf.dynamic_binary_concat_conf.out = "out" op_conf.dynamic_binary_concat_conf.out_data_type = ( source_blob.dtype.oneflow_proto_dtype) op_conf.dynamic_binary_concat_conf.out_shape.dim.extend( list(source_blob.shape)) if source_blob.batch_axis is not None: op_conf.dynamic_binary_concat_conf.out_batch_axis.value = source_blob.batch_axis else: op_conf.dynamic_binary_concat_conf.out_batch_axis.SetInParent() if "S" in source_sbp: axis = int(source_sbp.split(":")[-1]) op_conf.dynamic_binary_concat_conf.out_sbp.split_parallel.axis = axis elif "B" in source_sbp: op_conf.dynamic_binary_concat_conf.out_sbp.broadcast_parallel.SetInParent( ) elif "P" in source_sbp: op_conf.dynamic_binary_concat_conf.out_sbp.partial_sum_parallel.SetInParent( ) else: print("Error! invalid sbp str:", source_sbp) op_conf.dynamic_binary_concat_conf.out_sbp.SetInParent() interpret_util.Forward(op_conf) out_lbi = logical_blob_id_util.LogicalBlobId() out_lbi.op_name = op_conf.name out_lbi.blob_name = "out" return remote_blob_util.RemoteBlob(out_lbi)
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 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 RemoteBlobList(self): remote_blob_list = [] for k in self.op_conf_.user_conf.output: if k not in self.output_arg_key_list_: raise ValueError( "output_arg_name {} of {} op is not set in python op builder" .format(k, self.op_conf_.name)) for output_arg_name in self.output_arg_key_list_: assert output_arg_name in self.op_conf_.user_conf.output for i in range( len(self.op_conf_.user_conf.output[output_arg_name].s)): lbi = logical_blob_id_util.LogicalBlobId() lbi.op_name = self.op_conf_.name lbi.blob_name = "{}_{}".format(output_arg_name, i) remote_blob_list.append(self.MakeRemoteBlob(lbi)) return tuple(remote_blob_list)
def distribute_clone(x, name=None): if name is None: name = id_util.UniqueStr("DistributeClone_") op_conf = op_conf_util.OperatorConf() op_conf.name = name setattr(op_conf.distribute_clone_conf, "in", x.unique_name) parallel_size = oneflow.placement.current_scope().parallel_size op_conf.distribute_clone_conf.out.extend( ["out_%d" % i for i in range(parallel_size)]) interpret_util.ConsistentForward(op_conf) ret = [] for i in range(parallel_size): out = "out_%d" % i lbi = logical_blob_id_util.LogicalBlobId() lbi.op_name = op_conf.name lbi.blob_name = out ret.append(remote_blob_util.RemoteBlob(lbi)) return tuple(ret)
def dynamic_reshape( x: remote_blob_util.BlobDef, shape: Sequence[int], name: Optional[str] = None ) -> remote_blob_util.BlobDef: assert isinstance(shape, tuple) or isinstance(shape, list) shape = list(shape) 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 elem_cnt( inputs: remote_blob_util.BlobDef, 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("ElemCnt_")) op_conf.shape_elem_cnt_conf.x = inputs.unique_name op_conf.shape_elem_cnt_conf.exclude_axis_conf.SetInParent() 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() setattr(out_lbi, "op_name", op_conf.name) setattr(out_lbi, "blob_name", "y") return remote_blob_util.RemoteBlob(out_lbi)
def _GetReturnOpConfAndOutLbiAndScope(remote_blob, allow_cpu_return_op=True): op_conf = op_conf_util.OperatorConf() op_conf.name = id_util.UniqueStr("Return_") setattr(op_conf.return_conf, "in", remote_blob.unique_name) op_conf.return_conf.out = "out" if allow_cpu_return_op: op_conf.device_tag = "cpu" lbi = logical_blob_id_util.LogicalBlobId() lbi.op_name = op_conf.name lbi.blob_name = "out" parallel_conf = placement_cfg.ParallelConf() parallel_conf.CopyFrom(remote_blob.parallel_conf) def BuildScope(old_scope, builder): return builder.BuildScopeWithNewParallelConf(old_scope, parallel_conf) sess = session_ctx.GetDefaultSession() scope = scope_util.MakeScope(BuildScope) return (op_conf, lbi, scope)
def decode_random( shape: Sequence[int], dtype: flow.dtype, batch_size: int = 1, initializer: Optional[initializer_conf_util.InitializerConf] = None, tick: Optional[oneflow_api.BlobDesc] = None, name: Optional[str] = None, ) -> oneflow_api.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_api.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 distribute_split(x, axis=0, name=None): if name is None: name = id_util.UniqueStr("DistributeSplit_") op_conf = op_conf_util.OperatorConf() op_conf.name = name setattr(op_conf.distribute_split_conf, "in", x.unique_name) op_conf.distribute_split_conf.axis = axis parallel_size = oneflow.current_scope( ).device_parallel_desc_symbol.parallel_num op_conf.distribute_split_conf.out.extend( ["out_%d" % i for i in range(parallel_size)]) interpret_util.ConsistentForward(op_conf) ret = [] for i in range(parallel_size): out = "out_%d" % i lbi = logical_blob_id_util.LogicalBlobId() lbi.op_name = op_conf.name lbi.blob_name = out ret.append(remote_blob_util.RemoteBlob(lbi)) return tuple(ret)
def RemoteBlobList(self): remote_blob_list = [] for k in self.op_conf_.user_conf.output: if k not in self.output_arg_key_list_: raise ValueError( "output_arg_name {} of {} op is not set in python op builder" .format(k, self.op_conf_.name)) for output_arg_name in self.output_arg_key_list_: assert output_arg_name in self.op_conf_.user_conf.output for i in range( len(self.op_conf_.user_conf.output[output_arg_name].s)): lbi = logical_blob_id_util.LogicalBlobId() lbi.op_name = self.op_conf_.name lbi.blob_name = "{}_{}".format(output_arg_name, i) remote_blob_obj = self.MakeRemoteBlob(lbi) remote_blob_list.append(remote_blob_obj) if flow.eager_execution_enabled(): gradient_util.GetDefaultBackwardBlobRegister( ).TrySetObject4BlobName(remote_blob_obj.logical_blob_name, remote_blob_obj.blob_object) return tuple(remote_blob_list)
def stack( inputs: Sequence[remote_blob_util.BlobDef], axis: int, name: Optional[str] = None ) -> remote_blob_util.BlobDef: if not isinstance(inputs, (list, tuple)): inputs = [inputs] if axis < 0: axis = axis + len(inputs[0].shape) assert axis == 0, "Only support dim0 stack now." op_conf = op_conf_util.OperatorConf() setattr(op_conf, "name", name or id_util.UniqueStr("Stack_")) getattr(op_conf.stack_conf, "in").extend([input.unique_name for input in inputs]) setattr(op_conf.stack_conf, "axis", axis) setattr(op_conf.stack_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 sync_dynamic_resize( inputs: remote_blob_util.BlobDef, size: remote_blob_util.BlobDef, 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("SyncDynamicResize_"), ) setattr(op_conf.sync_dynamic_resize_conf, "in", inputs.unique_name) setattr(op_conf.sync_dynamic_resize_conf, "size", size.unique_name) setattr(op_conf.sync_dynamic_resize_conf, "axis", 0) setattr(op_conf.sync_dynamic_resize_conf, "out", "out") setattr(op_conf.sync_dynamic_resize_conf, "eager", flow.eager_execution_enabled()) 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 tensor_list_split( input_tensor_list: remote_blob_util.BlobDef, name: Optional[str] = None) -> Tuple[remote_blob_util.BlobDef]: 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)