def print_job_set(self): self._check_status(self.SessionStatus.OPEN, self.SessionStatus.RUNNING) job_set = c_api_util.GetJobSet() for job in job_set.job: print("job_name:", job.job_conf.job_name) for op_conf in job.net.op: print("\top_name:", op_conf.name)
def Export( job_func: Callable, model_save_dir: Text, onnx_filename: Text, continue_on_error: bool = False, opset: Optional[int] = None, extra_opset: Optional[int] = None, shape_override: Optional[Dict[Text, List[int]]] = None, external_data: bool = False, ): r"""Export a oneflow model into ONNX format. Args: job_func: The job function model_save_dir: The directory containing oneflow model weights. Users are expected to call check_point.save(dir), wait for the model saving finishing, and pass the argument 'dir' as model_save_dir. onnx_filename: a string for the output filename continue_on_error: if an op can't be processed (aka there is no mapping), continue opset: the opset to be used (int, default is oneflow.python.onnx.constants.PREFERRED_OPSET) extra_opset: list of extra opset's, for example the opset's used by custom ops shape_override: dict with inputs that override the shapes given by oneflow external_data: Save weights as ONNX external data, usually to bypass the 2GB file size limit of protobuf. """ assert os.getenv("ENABLE_USER_OP") != "False" assert os.path.isdir(model_save_dir) job_set = c_api_util.GetJobSet() job_name = job_func.__name__ for job in job_set.job: # TODO(OYY) Modify the interface before modifying it if job.job_conf.job_name == job_name: onnx_graph = ProcessFlowGraph( job, model_save_dir, continue_on_error=continue_on_error, opset=opset, extra_opset=extra_opset, shape_override=shape_override, ) onnx_graph = optimizer.OptimizeGraph(onnx_graph) model_proto = onnx_graph.MakeModel( job_name, onnx_filename, external_data=external_data ) with open(onnx_filename, "wb") as f: try: f.write(model_proto.SerializeToString()) except ValueError as e: raise ValueError( "Error occured when running model_proto.SerializeToString(). If the model is larger than 2GB, please specify external_data=True when calling flow.onnx.export. Original error message:\n{}".format( e ) ) return raise ValueError('Cannot find job "{}" in jobset'.format(job_name))
def UpdateInfo4InterfaceOp(self): for op_attr in c_api_util.GetInterfaceOpAttributes().op_attribute: self.interface_op_name2op_attr_[op_attr.op_conf.name] = op_attr for job in c_api_util.GetJobSet().job: op_name2parallel_conf = {} for placement_group in job.placement.placement_group: for op_name in placement_group.op_set.op_name: op_name2parallel_conf[ op_name] = placement_group.parallel_conf for op_conf in job.net.op: if c_api_util.IsInterfaceOpConf(op_conf): self.interface_op_name2job_name_[ op_conf.name] = job.job_conf.job_name self.lazy_interface_op_name2parallel_conf_[ op_conf.name] = op_name2parallel_conf[op_conf.name]