未验证 提交 3f371db8 编写于 作者: Z Zhang Ting 提交者: GitHub

add op_device attr for backward op_desc, test=develop (#23062)

上级 9ef7f6e8
......@@ -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)
......
......@@ -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")
......
Markdown is supported
0% .
You are about to add 0 people to the discussion. Proceed with caution.
先完成此消息的编辑!
想要评论请 注册