提交 493396d8 编写于 作者: Q Qiao Longfei 提交者: GitHub

add support_gpu (#3304)

* add support_gpu

* fix allclose

* fix name error and symplify code
上级 36ac89b9
......@@ -260,6 +260,12 @@ class OpRegistry {
return CreateOp(op_desc.type(), inputs, outputs, attrs);
}
static bool SupportGPU(const std::string& op_type) {
OperatorWithKernel::OpKernelKey key;
key.place_ = platform::GPUPlace();
return OperatorWithKernel::AllOpKernels().at(op_type).count(key) != 0;
}
static std::shared_ptr<OperatorBase> CreateGradOp(const OperatorBase& op) {
PADDLE_ENFORCE(!op.IsNetOp(),
"Use framework::Backward to get backward ops");
......
......@@ -200,6 +200,8 @@ All parameter, weight, gradient are variables in Paddle.
return OpRegistry::CreateOp(desc);
});
operator_base.def_static("support_gpu", &OpRegistry::SupportGPU);
operator_base.def("backward",
[](const OperatorBase &forwardOp,
const std::unordered_set<std::string> &no_grad_vars) {
......
......@@ -28,7 +28,7 @@ class OpTestMeta(type):
kwargs = dict()
places = []
places.append(core.CPUPlace())
if core.is_compile_gpu():
if core.is_compile_gpu() and core.Operator.support_gpu(self.type):
places.append(core.GPUPlace(0))
for place in places:
......@@ -66,7 +66,9 @@ class OpTestMeta(type):
for out_name in func.all_output_args:
actual = numpy.array(scope.find_var(out_name).get_tensor())
expect = self.outputs[out_name]
numpy.isclose(actual, expect)
self.assertTrue(
numpy.allclose(actual, expect),
"output name: " + out_name + "has diff")
obj.test_all = test_all
return obj
Markdown is supported
0% .
You are about to add 0 people to the discussion. Proceed with caution.
先完成此消息的编辑!
想要评论请 注册