未验证 提交 dbbe6e9c 编写于 作者: W wawltor 提交者: GitHub

fix the device supported of the op unique and unique_with_counts. (#21395)

* fix the device supported of the op unique and unique_with_counts.
test=develop
test=document_fix

* Fix the precision of test in the op of unique and unique_with_counts.
test=develop
test=document_fix
上级 32959e03
......@@ -35,6 +35,14 @@ class UniqueOp : public framework::OperatorWithKernel {
ctx->SetOutputDim("Out", {-1});
ctx->SetOutputDim("Index", in_dims);
}
protected:
framework::OpKernelType GetExpectedKernelType(
const framework::ExecutionContext& ctx) const override {
return framework::OpKernelType(
OperatorWithKernel::IndicateVarDataType(ctx, "X"),
platform::CPUPlace());
}
};
class UniqueOpMaker : public framework::OpProtoAndCheckerMaker {
......
......@@ -40,6 +40,14 @@ class UniqueWithCountsOp : public framework::OperatorWithKernel {
ctx->SetOutputDim("Index", in_dims);
ctx->SetOutputDim("Count", {-1});
}
protected:
framework::OpKernelType GetExpectedKernelType(
const framework::ExecutionContext& ctx) const override {
return framework::OpKernelType(
OperatorWithKernel::IndicateVarDataType(ctx, "X"),
platform::CPUPlace());
}
};
class UniqueWithCountsOpMaker : public framework::OpProtoAndCheckerMaker {
......
......@@ -12746,7 +12746,7 @@ def unique_with_counts(x, dtype='int32'):
This OP return a unique tensor for `x` , and count tensor that the count of unqiue result in raw input, \
and an index tensor pointing to this unique tensor.
**NOTICE**: This op just be supported in device of CPU, and support the variable type of Tensor only.
**NOTICE**: This op support the variable type of Tensor only.
Args:
x(Variable): A 1-D input tensor with input shape of :math:`[N]` , the input data type is float32, float64, int32, int64.
......
......@@ -68,5 +68,47 @@ class TestRandom(TestUniqueOp):
self.outputs = {'Out': target_out, 'Index': target_index}
@unittest.skipIf(not core.is_compiled_with_cuda(),
"core is not compiled with CUDA")
class TestOneGPU(TestUniqueOp):
def init_config(self):
self.inputs = {'X': np.array([2], dtype='int64'), }
self.attrs = {'dtype': int(core.VarDesc.VarType.INT32)}
self.outputs = {
'Out': np.array(
[2], dtype='int64'),
'Index': np.array(
[0], dtype='int32')
}
def test_check_output(self):
if core.is_compiled_with_cuda():
place = core.CUDAPlace(0)
self.check_output_with_place(place, atol=1e-5)
@unittest.skipIf(not core.is_compiled_with_cuda(),
"core is not compiled with CUDA")
class TestRandomGPU(TestUniqueOp):
def init_config(self):
self.inputs = {'X': np.random.randint(0, 100, (150, ), dtype='int64')}
self.attrs = {'dtype': int(core.VarDesc.VarType.INT64)}
np_unique, np_index, reverse_index = np.unique(self.inputs['X'], True,
True)
np_tuple = [(np_unique[i], np_index[i]) for i in range(len(np_unique))]
np_tuple.sort(key=lambda x: x[1])
target_out = np.array([i[0] for i in np_tuple], dtype='int64')
target_index = np.array(
[list(target_out).index(i) for i in self.inputs['X']],
dtype='int64')
self.outputs = {'Out': target_out, 'Index': target_index}
def test_check_output(self):
if core.is_compiled_with_cuda():
place = core.CUDAPlace(0)
self.check_output_with_place(place, atol=1e-5)
if __name__ == "__main__":
unittest.main()
......@@ -80,5 +80,57 @@ class TestRandom(TestUniqueWithCountsOp):
}
@unittest.skipIf(not core.is_compiled_with_cuda(),
"core is not compiled with CUDA")
class TestOneGPU(TestUniqueWithCountsOp):
def init_config(self):
self.inputs = {'X': np.array([2], dtype='int64'), }
self.attrs = {'dtype': int(core.VarDesc.VarType.INT32)}
self.outputs = {
'Out': np.array(
[2], dtype='int64'),
'Index': np.array(
[0], dtype='int32'),
'Count': np.array(
[1], dtype='int32')
}
def test_check_output(self):
if core.is_compiled_with_cuda():
place = core.CUDAPlace(0)
self.check_output_with_place(place, atol=1e-5)
@unittest.skipIf(not core.is_compiled_with_cuda(),
"core is not compiled with CUDA")
class TestRandomGPU(TestUniqueWithCountsOp):
def init_config(self):
input_data = np.random.randint(0, 100, (2000, ), dtype='int64')
self.inputs = {'X': input_data}
self.attrs = {'dtype': int(core.VarDesc.VarType.INT64)}
np_unique, np_index, reverse_index = np.unique(self.inputs['X'], True,
True)
np_tuple = [(np_unique[i], np_index[i]) for i in range(len(np_unique))]
np_tuple.sort(key=lambda x: x[1])
target_out = np.array([i[0] for i in np_tuple], dtype='int64')
target_index = np.array(
[list(target_out).index(i) for i in self.inputs['X']],
dtype='int64')
count = [0 for i in range(len(np_unique))]
for i in range(target_index.shape[0]):
count[target_index[i]] += 1
target_count = np.array(count, dtype='int64')
self.outputs = {
'Out': target_out,
'Index': target_index,
'Count': target_count
}
def test_check_output(self):
if core.is_compiled_with_cuda():
place = core.CUDAPlace(0)
self.check_output_with_place(place, atol=1e-5)
if __name__ == "__main__":
unittest.main()
Markdown is supported
0% .
You are about to add 0 people to the discussion. Proceed with caution.
先完成此消息的编辑!
想要评论请 注册