From 3f371db82a7c720f8f87cb46b79d9c5dc550a663 Mon Sep 17 00:00:00 2001
From: Zhang Ting <709968123@qq.com>
Date: Thu, 19 Mar 2020 12:59:25 +0800
Subject: [PATCH] add op_device attr for backward op_desc, test=develop
 (#23062)

---
 python/paddle/fluid/backward.py                          | 3 +++
 python/paddle/fluid/tests/unittests/test_device_guard.py | 9 +++++++--
 2 files changed, 10 insertions(+), 2 deletions(-)

diff --git a/python/paddle/fluid/backward.py b/python/paddle/fluid/backward.py
index 25778e834a..095ca5772d 100644
--- a/python/paddle/fluid/backward.py
+++ b/python/paddle/fluid/backward.py
@@ -273,10 +273,13 @@ def _create_op_desc_(op_type, inputs, outputs, attrs):
                     args)))
 
     op_role_attr_name = core.op_proto_and_checker_maker.kOpRoleAttrName()
+    op_device_attr_name = core.op_proto_and_checker_maker.kOpDeviceAttrName()
 
     if op_role_attr_name not in attrs:
         attrs[
             op_role_attr_name] = core.op_proto_and_checker_maker.OpRole.Backward
+    if op_device_attr_name not in attrs:
+        attrs[op_device_attr_name] = ""
     for name, val in six.iteritems(attrs):
         if isinstance(val, framework.Block):
             op_desc.set_block_attr(name, val.desc)
diff --git a/python/paddle/fluid/tests/unittests/test_device_guard.py b/python/paddle/fluid/tests/unittests/test_device_guard.py
index 1384a22b25..7b05f6b0f2 100644
--- a/python/paddle/fluid/tests/unittests/test_device_guard.py
+++ b/python/paddle/fluid/tests/unittests/test_device_guard.py
@@ -144,16 +144,19 @@ class TestDeviceGuard(unittest.TestCase):
         for op in all_ops:
             self.assertEqual(op.desc.attr(device_attr_name), "gpu")
 
-    def test_loss_op_desc(self):
+    # check if op_descs have op_device attr
+    def test_op_descs_device_attr(self):
         main_program = fluid.Program()
         startup_program = fluid.Program()
         with fluid.program_guard(main_program, startup_program):
             data1 = fluid.layers.data(name="data_1", shape=[2], dtype="float32")
+            data2 = fluid.layers.data(name="data_2", shape=[2], dtype="float32")
             label = fluid.layers.data(name="label", shape=[1], dtype="int64")
             fc1 = fluid.layers.fc(input=data1, size=10)
+            fc2 = fluid.layers.fc(input=fc1, size=10)
             with fluid.device_guard("gpu"):
                 out = fluid.layers.softmax_with_cross_entropy(
-                    logits=fc1, label=label)
+                    logits=fc1 + fc2, label=label)
                 loss = fluid.layers.mean(out)
                 opt = fluid.optimizer.SGDOptimizer(0.1)
                 opt.minimize(loss)
@@ -162,6 +165,8 @@ class TestDeviceGuard(unittest.TestCase):
         device_attr_name = core.op_proto_and_checker_maker.kOpDeviceAttrName()
         for op in all_ops:
             self.assertEqual(True, op.desc.has_attr(device_attr_name))
+            # fill_constant(backward op) is append to mean op, which should have
+            # the same op_device value as mean op
             if op.desc == 'fill_constant':
                 self.assertEqual(op.desc.attr(device_attr_name), "gpu")
 
-- 
GitLab