p85971406
/
mindspore
forked from mindspore-Ecosystem/mindspore
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

!11328 OneHot int64 input support 

From: @peilin-wang
Reviewed-by: @tom__chen,@robingrosman
Signed-off-by: @robingrosman
This commit is contained in:
mindspore-ci-bot 2021-01-21 06:01:21 +08:00 committed by Gitee
parent 868f0e9361 386b219ffc
commit ddf84551ab
4 changed files with 44 additions and 14 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

17
mindspore/ccsrc/backend/kernel_compiler/gpu/arrays/one_hot_gpu_kernel.cc
View File

@ -1,5 +1,5 @@
/**
* Copyright 2019 Huawei Technologies Co., Ltd
* Copyright 2021 Huawei Technologies Co., Ltd
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
@ -13,6 +13,7 @@
* See the License for the specific language governing permissions and
* limitations under the License.
*/
#include <cstdint>
#include "backend/kernel_compiler/gpu/arrays/one_hot_gpu_kernel.h"
@ -32,5 +33,19 @@ MS_REG_GPU_KERNEL_TWO(OneHot,
.AddInputAttr(kNumberTypeFloat16)
.AddOutputAttr(kNumberTypeFloat16),
OneHotGpuFwdKernel, half, int)
MS_REG_GPU_KERNEL_TWO(OneHot,
KernelAttr()
.AddInputAttr(kNumberTypeInt64)
.AddInputAttr(kNumberTypeFloat32)
.AddInputAttr(kNumberTypeFloat32)
.AddOutputAttr(kNumberTypeFloat32),
OneHotGpuFwdKernel, float, int64_t)
MS_REG_GPU_KERNEL_TWO(OneHot,
KernelAttr()
.AddInputAttr(kNumberTypeInt64)
.AddInputAttr(kNumberTypeFloat16)
.AddInputAttr(kNumberTypeFloat16)
.AddOutputAttr(kNumberTypeFloat16),
OneHotGpuFwdKernel, half, int64_t)
} // namespace kernel
} // namespace mindspore

6
mindspore/ccsrc/backend/kernel_compiler/gpu/cuda_impl/one_hot_impl.cu
View File

@ -49,3 +49,9 @@ template void OneHot<float, int>(const int *indices, size_t depth, const float *
size_t left_dim_size, size_t right_dim_size, float *output, cudaStream_t cuda_stream);
template void OneHot<half, int>(const int *indices, size_t depth, const half *on_value, const half *off_value,
size_t left_dim_size, size_t right_dim_size, half *output, cudaStream_t cuda_stream);
template void OneHot<float, int64_t>(const int64_t *indices, size_t depth, const float *on_value,
const float *off_value, size_t left_dim_size, size_t right_dim_size, float *output,
cudaStream_t cuda_stream);
template void OneHot<half, int64_t>(const int64_t *indices, size_t depth, const half *on_value, const half *off_value,
size_t left_dim_size, size_t right_dim_size, half *output,
cudaStream_t cuda_stream);

4
mindspore/ops/operations/nn_ops.py
View File

@ -2985,7 +2985,7 @@ class OneHot(PrimitiveWithInfer):
Inputs:
- **indices** (Tensor) - A tensor of indices. Tensor of shape :math:`(X_0, \ldots, X_n)`.
Data type must be int32.
Data type must be int32 or int64.
- **depth** (int) - A scalar defining the depth of the one hot dimension.
- **on_value** (Tensor) - A value to fill in output when `indices[j] = i`. With data type of float16 or float32.
- **off_value** (Tensor) - A value to fill in output when `indices[j] != i`.
@ -3015,7 +3015,7 @@ class OneHot(PrimitiveWithInfer):
def __infer__(self, indices, depth, on_value, off_value):
# check type
validator.check_tensor_dtype_valid("indices", indices['dtype'], (mstype.int32,), self.name)
validator.check_tensor_dtype_valid("indices", indices['dtype'], (mstype.int32, mstype.int64), self.name)
validator.check_type_name("depth", depth['dtype'], mstype.int_type, self.name)
args = {"on_value": on_value['dtype'], "off_value": off_value['dtype']}
validator.check_tensors_dtypes_same_and_valid(args, (mstype.float16, mstype.float32), self.name)

31
tests/st/ops/gpu/test_one_hot_op.py
View File

@ -1,4 +1,4 @@
# Copyright 2019 Huawei Technologies Co., Ltd
# Copyright 2021 Huawei Technologies Co., Ltd
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
@ -44,16 +44,13 @@ class NetOneHot(nn.Cell):
self.one_hot_3(indices3), self.one_hot_4(indices4))
@pytest.mark.level0
@pytest.mark.platform_x86_gpu_training
@pytest.mark.env_onecard
def test_one_hot():
one_hot = NetOneHot()
indices1 = Tensor(np.array([[0, 1], [4, 5], [2, 6]]).astype(np.int32))
indices2 = Tensor(np.array([1, 2, 3]).astype(np.int32))
indices3 = Tensor(np.array([[0, 1], [1, 0]]).astype(np.int32))
indices4 = Tensor(np.array([[0, 1], [4, 5], [2, 6]]).astype(np.int32))
output = one_hot(indices1, indices2, indices3, indices4)
def one_hot(nptype):
one_hot_net = NetOneHot()
indices1 = Tensor(np.array([[0, 1], [4, 5], [2, 6]]).astype(nptype))
indices2 = Tensor(np.array([1, 2, 3]).astype(nptype))
indices3 = Tensor(np.array([[0, 1], [1, 0]]).astype(nptype))
indices4 = Tensor(np.array([[0, 1], [4, 5], [2, 6]]).astype(nptype))
output = one_hot_net(indices1, indices2, indices3, indices4)
expect_0 = np.array([
[[2., 3., 3., 3., 3., 3.], [3., 2., 3., 3., 3., 3.]],
[[3., 3., 3., 3., 2., 3.], [3., 3., 3., 3., 3., 2.]],
@ -80,3 +77,15 @@ def test_one_hot():
assert (output[1].asnumpy() == expect_1).all()
assert (output[2].asnumpy() == expect_2).all()
assert (output[3].asnumpy() == expect_3).all()
@pytest.mark.level0
@pytest.mark.platform_x86_gpu_training
@pytest.mark.env_onecard
def test_one_hot_int32():
one_hot(np.int32)
@pytest.mark.level0
@pytest.mark.platform_x86_gpu_training
@pytest.mark.env_onecard
def test_one_hot_int64():
one_hot(np.int64)