def global_function_or_identity(*args, **kwargs):
if rt_mode.CurrentMode() == rt_mode.NORMAL_MODE:
return api_oneflow_function(*args, **kwargs)
else:
assert rt_mode.CurrentMode() == rt_mode.GLOBAL_MODE
identity_decorator = lambda func: func
return identity_decorator
def GenerateVariableOpConf(
name,
shape,
dtype=None,
initializer=None,
regularizer=None,
trainable=None,
model_name=None,
random_seed=None,
nd_sbp=None,
):
op_conf = op_conf_util.OperatorConf()
op_conf.name = name
op_conf.variable_conf.shape.dim.extend(shape)
assert dtype is not None
op_conf.variable_conf.data_type = oneflow._oneflow_internal.deprecated.GetProtoDtype4OfDtype(
dtype)
if rt_mode.CurrentMode() == rt_mode.NORMAL_MODE:
root_path = None
else:
root_path = (compile_context.GetCurJobConfigProto().
default_initialize_with_snapshot_path())
dir_path = os.path.join(root_path, name)
file_path = os.path.join(dir_path, "out")
if root_path and os.path.isfile(file_path):
op_conf.variable_conf.initialize_with_snapshot.path = dir_path
op_conf.variable_conf.initialize_with_snapshot.key = "out"
else:
if root_path:
print("{} not found, will be initialized".format(file_path))
if initializer is not None:
op_conf.variable_conf.initializer.CopyFrom(initializer)
if regularizer is not None:
op_conf.variable_conf.regularizer.CopyFrom(regularizer)
if trainable is not None:
op_conf.variable_conf.trainable = trainable
if model_name is not None:
op_conf.variable_conf.model_name = model_name
if nd_sbp is None:
nd_sbp = []
op_conf.variable_conf.nd_sbp.extend(nd_sbp)
if random_seed is not None:
op_conf.variable_conf.random_seed = random_seed
op_conf.variable_conf.out = "out"
return op_conf
def sync_default_session_if_normal():
if rt_mode.CurrentMode() == rt_mode.NORMAL_MODE:
flow.sync_default_session()
else:
pass
def in_device_mode(ctx):
return rt_mode.CurrentMode() == rt_mode.DEVICE_MODE
def in_global_mode(ctx):
return rt_mode.CurrentMode() == rt_mode.GLOBAL_MODE
def in_normal_mode(ctx):
return rt_mode.CurrentMode() == rt_mode.NORMAL_MODE