def GetInterfaceBlobConf(job_name, lbn, blob_conf=None): assert isinstance(job_name, str) assert isinstance(lbn, str) if blob_conf is None: blob_conf = interface_blob_conf_pb.InterfaceBlobConf() else: assert isinstance(blob_conf, interface_blob_conf_pb.InterfaceBlobConf) shape = c_api_util.JobBuildAndInferCtx_GetStaticShape(job_name, lbn) dtype = c_api_util.JobBuildAndInferCtx_GetDataType(job_name, lbn) split_axis = c_api_util.JobBuildAndInferCtx_GetSplitAxisFromProducerView( job_name, lbn) batch_axis = c_api_util.JobBuildAndInferCtx_GetBatchAxis(job_name, lbn) is_dynamic = c_api_util.JobBuildAndInferCtx_IsDynamic(job_name, lbn) is_tensor_list = c_api_util.JobBuildAndInferCtx_IsTensorList(job_name, lbn) blob_conf.shape.dim.extend(shape) blob_conf.data_type = dtype if split_axis is not None: blob_conf.split_axis.value = split_axis if batch_axis is not None: blob_conf.batch_axis.value = batch_axis blob_conf.is_dynamic = is_dynamic blob_conf.is_tensor_list = is_tensor_list return blob_conf
def _make_pull_job_cb(self, output_name, user_job_name, future): output_lbn = oneflow_api.JobBuildAndInferCtx_GetOpBlobLbn( user_job_name, output_name, "out") split_axis = c_api_util.JobBuildAndInferCtx_GetSplitAxisFromProducerView( user_job_name, output_lbn) def pull_fn(ofblob): ndarray = ofblob.CopyToNdarray() self.event_loop_.call_soon_threadsafe(future.set_result, ndarray) return pull_fn
def _make_pull_job_cb(self, output_name, future): output_lbn = self.inferface_name2lbn_[output_name] job_name = self.inferface_name2job_name_[output_name] split_axis = c_api_util.JobBuildAndInferCtx_GetSplitAxisFromProducerView( job_name, output_lbn ) loop = asyncio.get_event_loop() def pull_fn(ofblob): ndarray_lists = ofblob.CopyToNdarrayLists() assert len(ndarray_lists) == 1 ndarray_list = ndarray_lists[0] if len(ndarray_list) == 1: loop.call_soon_threadsafe(future.set_result, ndarray_list[0]) else: assert split_axis is not None pull_result = np.concatenate(ndarray_list, axis=split_axis) loop.call_soon_threadsafe(future.set_result, pull_result) return pull_fn
def _make_pull_job_cb(self, output_name, user_job_name, future): output_lbn = oneflow_api.JobBuildAndInferCtx_GetOpBlobLbn( user_job_name, output_name, "out") split_axis = c_api_util.JobBuildAndInferCtx_GetSplitAxisFromProducerView( user_job_name, output_lbn) def pull_fn(ofblob): ndarray_lists = ofblob.CopyToNdarrayLists() assert len(ndarray_lists) == 1 ndarray_list = ndarray_lists[0] if len(ndarray_list) == 1: self.event_loop_.call_soon_threadsafe(future.set_result, ndarray_list[0]) else: assert split_axis is not None pull_result = np.concatenate(ndarray_list, axis=split_axis) self.event_loop_.call_soon_threadsafe(future.set_result, pull_result) return pull_fn
def GetInterfaceBlobConf(job_name, lbn, blob_conf=None): assert isinstance(job_name, str) assert isinstance(lbn, str) if blob_conf is None: blob_conf = interface_blob_conf_pb.InterfaceBlobConf() else: assert isinstance(blob_conf, interface_blob_conf_pb.InterfaceBlobConf) shape = c_api_util.JobBuildAndInferCtx_GetStaticShape(job_name, lbn) dtype = c_api_util.JobBuildAndInferCtx_GetDataType(job_name, lbn) split_axis = c_api_util.JobBuildAndInferCtx_GetSplitAxisFromProducerView( job_name, lbn) is_dynamic = c_api_util.JobBuildAndInferCtx_IsDynamic(job_name, lbn) blob_conf.shape.dim.extend(shape) blob_conf.data_type = dtype if split_axis is not None: sbp_parallel = sbp_parallel_pb.SbpParallel() sbp_parallel.split_parallel.axis = split_axis blob_conf.parallel_distribution.sbp_parallel.extend([sbp_parallel]) blob_conf.is_dynamic = is_dynamic return blob_conf
def split_axis(self): return c_api_util.JobBuildAndInferCtx_GetSplitAxisFromProducerView( self.job_name_, self.lbn_ )