From b53eb7dcda6cb2c8e5a3f49bee15189fc2232401 Mon Sep 17 00:00:00 2001
From: Qiao Longfei <qiaolongfei@baidu.com>
Date: Wed, 26 Dec 2018 19:43:38 +0800
Subject: [PATCH] add init once for assign layer

---
 python/paddle/fluid/layers/nn.py              |  8 ++--
 python/paddle/fluid/layers/tensor.py          | 39 ++++++++++++++-----
 .../fluid/tests/unittests/test_layers.py      | 12 ++++++
 3 files changed, 46 insertions(+), 13 deletions(-)

diff --git a/python/paddle/fluid/layers/nn.py b/python/paddle/fluid/layers/nn.py
index cc1fdbd2856..00523c07987 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 49a486cf0c3..d66d92b1df7 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 e180822c2b4..92065abb9bf 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()
-- 
GitLab