diff --git a/python/paddle/fluid/layers/nn.py b/python/paddle/fluid/layers/nn.py index cc1fdbd285611379cc4fa44d2373748aa6e24faf..00523c07987cb32fe3cf922a3cd76240b6b67139 100644 --- a/python/paddle/fluid/layers/nn.py +++ b/python/paddle/fluid/layers/nn.py @@ -5010,10 +5010,12 @@ def nce(input, alias_probs_[little[0]] = 1.0 alias_[little[0]] = -1 - probs = assign(input=np.array(custom_dist).astype('float32')) - custom_alias = assign(input=np.array(alias_).astype('int32')) + probs = assign( + input=np.array(custom_dist).astype('float32'), init_once=True) + custom_alias = assign( + input=np.array(alias_).astype('int32'), init_once=True) custom_alias_probs = assign( - input=np.array(alias_probs_).astype('float32')) + input=np.array(alias_probs_).astype('float32'), init_once=True) inputs['CustomDistProbs'] = probs inputs['CustomDistAlias'] = custom_alias diff --git a/python/paddle/fluid/layers/tensor.py b/python/paddle/fluid/layers/tensor.py index 49a486cf0c3d11b18417e8838aead07d748f3e02..d66d92b1df7fd4aa8312cca186c29a571d7c4af0 100644 --- a/python/paddle/fluid/layers/tensor.py +++ b/python/paddle/fluid/layers/tensor.py @@ -285,7 +285,7 @@ def sums(input, out=None): return out -def assign(input, output=None): +def assign(input, output=None, init_once=False): """ **Assign** @@ -294,6 +294,7 @@ def assign(input, output=None): Args: input(Variable|numpy.ndarray): The source variable output(Variable|None): The destination variable + init_once(bool|false): assign value into global var only in startup program. Returns: Variable: The destination variable that was supplied as the *output*. @@ -307,10 +308,18 @@ def assign(input, output=None): """ helper = LayerHelper('assign', **locals()) if output is None: - output = helper.create_variable_for_type_inference(dtype=input.dtype) + if init_once: + output = helper.create_parameter( + attr=ParamAttr(), shape=input.shape, dtype=input.dtype) + else: + output = helper.create_variable_for_type_inference( + dtype=input.dtype) if isinstance(input, Variable): + if init_once: + raise ValueError("init once only support numpy assign!") helper.append_op( type='assign', inputs={'X': [input]}, outputs={'Out': [output]}) + elif isinstance(input, numpy.ndarray): dtype = convert_np_dtype_to_dtype_(input.dtype) if dtype == VarDesc.VarType.FP32: @@ -325,14 +334,24 @@ def assign(input, output=None): raise ValueError("The size of input is too big. Please consider " "saving it to file and 'load_op' to load it") - helper.append_op( - type='assign_value', - outputs={'Out': [output]}, - attrs={ - 'dtype': dtype, - 'shape': list(input.shape), - value_name: values - }) + if init_once: + helper.startup_program.global_block().append_op( + type='assign_value', + outputs={'Out': [output]}, + attrs={ + 'dtype': dtype, + 'shape': list(input.shape), + value_name: values + }) + else: + helper.append_op( + type='assign_value', + outputs={'Out': [output]}, + attrs={ + 'dtype': dtype, + 'shape': list(input.shape), + value_name: values + }) else: raise ValueError("Wrong type for assign input: %s" % type(input)) diff --git a/python/paddle/fluid/tests/unittests/test_layers.py b/python/paddle/fluid/tests/unittests/test_layers.py index e180822c2b4b7cceaf9f66e7819477b48bf4941b..92065abb9bf0a8c7c60b24814869f52fec49264b 100644 --- a/python/paddle/fluid/tests/unittests/test_layers.py +++ b/python/paddle/fluid/tests/unittests/test_layers.py @@ -1015,6 +1015,18 @@ class TestBook(unittest.TestCase): print(str(program)) + def test_assign(self): + import numpy as np + startup = Program() + main = Program() + with program_guard(main, startup): + probs = layers.assign( + input=np.random.random([1, 2]).astype('float32'), + init_once=True) + + print(str(main)) + print(str(startup)) + if __name__ == '__main__': unittest.main()