def get_eager_variable( name, shape=None, dtype=None, initializer=None, regularizer=None, trainable=None, model_name=None, random_seed=None, distribute=oneflow_api.distribute.broadcast(), reuse=True, ): assert isinstance(name, str) assert isinstance( shape, (list, tuple) ), "param shape should be a list or tuple of dimension" job_name = oneflow_api.JobBuildAndInferCtx_GetCurrentJobName() name = name_scope.GetJobNameScopePrefix(job_name) + name sess = session_ctx.GetDefaultSession() var_blob, job_var_blob = sess.TryGetVariableBlobOfJobFromStash(job_name, name) if reuse is False: assert job_var_blob is None, ( "varaible '{}' already exists, " "getting the same variable is not allowed " "when reuse is False".format(name) ) if job_var_blob is None: op_conf = GenerateVariableOpConf( name=name, shape=shape, dtype=dtype, initializer=initializer, regularizer=regularizer, trainable=trainable, model_name=model_name, random_seed=random_seed, distribute=distribute, ) op_attribute = compile_context.CurJobAddConsistentOp(op_conf) if var_blob is None: var_blob = CreateEagerVariableBlob(op_attribute) op_executor.EagerInitVariableBlob(sess, op_conf, var_blob) assert isinstance(var_blob, remote_blob_util.EagerConsistentBlob) sess.StashVariableBlob4Job(job_name, op_conf.name, var_blob) else: assert isinstance(job_var_blob, remote_blob_util.EagerConsistentBlob) assert isinstance(var_blob, remote_blob_util.EagerConsistentBlob) assert var_blob.IdenticalTo(job_var_blob) bw_blob_register = gradient_util.GetDefaultBackwardBlobRegister() bw_blob_register.TrySetObject4BlobName( var_blob.logical_blob_name, var_blob.blob_object ) return var_blob
def get_lazy_variable( name, shape=None, dtype=None, initializer=None, regularizer=None, trainable=None, model_name=None, random_seed=None, parallel_distribution=None, reuse=True, ): assert isinstance(name, str) assert isinstance( shape, (list, tuple)), "param shape should be a list or tuple of dimension" job_name = oneflow._oneflow_internal.JobBuildAndInferCtx_GetCurrentJobName( ) name = name_scope.GetJobNameScopePrefix(job_name) + name sess = session_ctx.GetDefaultSession() var_blob, job_var_blob = sess.TryGetVariableBlobOfJobFromStash( job_name, name) if reuse is False: assert job_var_blob is None, ( "variable '{}' already exists, " "getting the same variable is not allowed " "when param reuse is False".format(name)) if job_var_blob is None: op_conf = GenerateVariableOpConf( name=name, shape=shape, dtype=dtype, initializer=initializer, regularizer=regularizer, trainable=trainable, model_name=model_name, random_seed=random_seed, parallel_distribution=parallel_distribution, ) job_var_blob = _CreateVariableBlob(op_conf) assert isinstance(job_var_blob, oneflow._oneflow_internal.LazyConsistentBlob) sess.StashVariableBlob4Job(job_name, op_conf.name, job_var_blob) if var_blob is not None: assert isinstance(var_blob, oneflow._oneflow_internal.LazyConsistentBlob) assert var_blob.IdenticalTo(job_var_blob) else: assert isinstance(job_var_blob, oneflow._oneflow_internal.LazyConsistentBlob) assert isinstance(var_blob, oneflow._oneflow_internal.LazyConsistentBlob) assert var_blob.IdenticalTo(job_var_blob) return job_var_blob
def OpName(self, op_name): job_name = c_api_util.JobBuildAndInferCtx_GetCurrentJobName() op_name = name_scope.GetJobNameScopePrefix(job_name) + op_name self.user_op_.op_conf_.name = op_name user_conf = self.user_op_.op_conf_.user_conf def GetLbn(output_name, i): return "{}/{}_{}".format(op_name, output_name, i) for output_name, output in user_conf.output.items(): output.s[:] = [GetLbn(output_name, i) for i in range(len(output.s))] return self
def _FindOrCreateVarBlobObject(op_attribute, parallel_conf, blob_register): job_name = oneflow_api.JobBuildAndInferCtx_GetCurrentJobName() name = name_scope.GetJobNameScopePrefix(job_name) + op_attribute.op_conf.name sess = session_ctx.GetDefaultSession() var_blob, _ = sess.TryGetVariableBlobOfJobFromStash(job_name, name) if var_blob is not None: blob_register.SetObject4BlobName( var_blob.logical_blob_name, var_blob.blob_object ) return _NaiveInterpret(op_attribute, parallel_conf, blob_register) var_blob = _MakeEagerLogicalBlob(op_attribute, "out", blob_register=blob_register) EagerInitVariableBlob(sess, op_attribute.op_conf, var_blob) sess.StashVariableBlob4Job(job_name, op_attribute.op_conf.name, var_blob) return var_blob
def eager_consistent_user_op_module_builder(op_type_name): job_name = oneflow_api.JobBuildAndInferCtx_GetCurrentJobName() op_name = name_scope.GetJobNameScopePrefix(job_name) + op_type_name return UserOpModuleBuilder(EagerConsistentUserOpModule, op_name, op_type_name)
def api_consistent_user_op_builder(op_name): job_name = oneflow_api.JobBuildAndInferCtx_GetCurrentJobName() op_name = name_scope.GetJobNameScopePrefix(job_name) + op_name return UserOpConfBuilder(ConsistentUserOp, op_name, None)
def lazy_consistent_user_op_module_builder(op_type_name): job_name = c_api_util.JobBuildAndInferCtx_GetCurrentJobName() op_name = name_scope.GetJobNameScopePrefix(job_name) + op_type_name return UserOpModuleBuilder(LazyConsistentUserOpModule, op_name, op_type_name)
def eager_logical_user_op_module_builder(op_type_name): job_name = c_api_util.JobBuildAndInferCtx_GetCurrentJobName() op_name = name_scope.GetJobNameScopePrefix(job_name) + op_type_name return UserOpModuleBuilder(EagerLogicalUserOpModule, op_name, op_type_name)
def eager_user_op_builder(op_name): job_name = c_api_util.JobBuildAndInferCtx_GetCurrentJobName() op_name = name_scope.GetJobNameScopePrefix(job_name) + op_name return UserOpConfBuilder(EagerUserOp, op_name, None)
def __init__(self, job_name, op_name, user_op_class): name_scope_prefix = name_scope.GetJobNameScopePrefix(job_name) self.user_op_ = user_op_class(name_scope_prefix + op_name)
def lazy_user_op_module_builder(op_type_name): job_name = oneflow._oneflow_internal.JobBuildAndInferCtx_GetCurrentJobName( ) op_name = name_scope.GetJobNameScopePrefix(job_name) + op_type_name return UserOpModuleBuilder(LazyUserOpModule, op_name, op_type_name)