m53297601
/
mindspore
1
0
Fork 0
Code Issues Pull Requests Releases Wiki Activity

GatherUpdate

Browse Source 
test file finishing pending

Update to GatherV2_Bug_Fix

lint fix

lint fix - 2

lint fix

Update to GatherV2 - fixed default inferImpl func + CudeStreamSync

lint fix

SyncDevice

added dynamic shape init_size input

lint
pull/8825/head
danishnxt 4 years ago
parent 85a020575a
commit 241c8f3d96
8 changed files with 223 additions and 215 deletions
Show Stats Download Patch File Download Diff File

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

@ -26,14 +26,6 @@ MS_REG_GPU_KERNEL_TWO(
GatherV2,
KernelAttr().AddInputAttr(kNumberTypeFloat16).AddInputAttr(kNumberTypeInt32).AddOutputAttr(kNumberTypeFloat16),
GatherV2GpuFwdKernel, half, int)
MS_REG_GPU_KERNEL_TWO(
SparseGatherV2,
KernelAttr().AddInputAttr(kNumberTypeFloat32).AddInputAttr(kNumberTypeInt32).AddOutputAttr(kNumberTypeFloat32),
GatherV2GpuFwdKernel, float, int)
MS_REG_GPU_KERNEL_TWO(
SparseGatherV2,
KernelAttr().AddInputAttr(kNumberTypeFloat16).AddInputAttr(kNumberTypeInt32).AddOutputAttr(kNumberTypeFloat16),
GatherV2GpuFwdKernel, half, int)
MS_REG_GPU_KERNEL_TWO(GatherV2,
KernelAttr()
.AddInputAttr(kNumberTypeFloat32)
@ -48,5 +40,14 @@ MS_REG_GPU_KERNEL_TWO(GatherV2,
.AddInputAttr(kNumberTypeInt64)
.AddOutputAttr(kNumberTypeFloat16),
GatherV2GpuFwdKernel, half, int)
MS_REG_GPU_KERNEL_TWO(
SparseGatherV2,
KernelAttr().AddInputAttr(kNumberTypeFloat32).AddInputAttr(kNumberTypeInt32).AddOutputAttr(kNumberTypeFloat32),
GatherV2GpuFwdKernel, float, int)
MS_REG_GPU_KERNEL_TWO(
SparseGatherV2,
KernelAttr().AddInputAttr(kNumberTypeFloat16).AddInputAttr(kNumberTypeInt32).AddOutputAttr(kNumberTypeFloat16),
GatherV2GpuFwdKernel, half, int)
} // namespace kernel
} // namespace mindspore

108
mindspore/ccsrc/backend/kernel_compiler/gpu/arrays/gatherv2_gpu_kernel.h
Unescape View File

@ -14,8 +14,8 @@
* limitations under the License.
*/
#ifndef MINDSPORE_GATHER_GPU_KERNEL_H
#define MINDSPORE_GATHER_GPU_KERNEL_H
#ifndef MINDSPORE_CCSRC_BACKEND_KERNEL_COMPILER_GPU_GATHER_V2_GPU_KERNEL_H_
#define MINDSPORE_CCSRC_BACKEND_KERNEL_COMPILER_GPU_GATHER_V2_GPU_KERNEL_H_
#include <vector>
#include <algorithm>
@ -41,45 +41,17 @@ class GatherV2GpuFwdKernel : public GpuKernel {
T *input_addr = GetDeviceAddress<T>(inputs, 0);
S *indices_addr = GetDeviceAddress<S>(inputs, 1);
T *output_addr = GetDeviceAddress<T>(outputs, 0);
if (is_dynamic_shape_) {
// if we are in dynamic shape mode, we don't know dims_, so we need to store the input_shape_ and indices_shape_,
// and axis_ in the workspace to calculate dims_
size_t *input_shape_device_address = GetDeviceAddress<size_t>(workspace, 0);
size_t *indices_shape_device_address = GetDeviceAddress<size_t>(workspace, 1);
int64_t *axis_device_address = GetDeviceAddress<int64_t>(workspace, 2);
CHECK_CUDA_RET_WITH_EXCEPT(
cudaMemcpyAsync(input_shape_device_address, input_shapes_.data(), workspace_size_list_[0],
cudaMemcpyHostToDevice, reinterpret_cast<cudaStream_t>(stream_ptr)),
"cudaMemcpyAsync input_shape failed");
CHECK_CUDA_RET_WITH_EXCEPT(
cudaMemcpyAsync(indices_shape_device_address, indices_shapes_.data(), workspace_size_list_[1],
cudaMemcpyHostToDevice, reinterpret_cast<cudaStream_t>(stream_ptr)),
"cudaMemcpyAsync indices_shape failed");
CHECK_CUDA_RET_WITH_EXCEPT(cudaMemcpyAsync(axis_device_address, &axis_, workspace_size_list_[2],
cudaMemcpyHostToDevice, reinterpret_cast<cudaStream_t>(stream_ptr)),
int64_t *axis_device_address = GetDeviceAddress<int64_t>(inputs, 2); // only get this if in dynamic mode
CHECK_CUDA_RET_WITH_EXCEPT(cudaMemcpyAsync(&axis_, axis_device_address, sizeof(int64_t), cudaMemcpyDeviceToHost,
reinterpret_cast<cudaStream_t>(stream_ptr)),
"cudaMemcpyAsync axis_ failed");
// output shape will be here for us to copy back to host
size_t *output_shape_device_address = GetDeviceAddress<size_t>(workspace, 3);
CalGatherV2DynamicShape(input_addr, indices_addr, output_addr, input_shape_device_address, input_shapes_.size(),
indices_shape_device_address, indices_shapes_.size(), axis_device_address,
output_shape_device_address, max_output_size_,
reinterpret_cast<cudaStream_t>(stream_ptr));
size_t output_rank = input_shapes_.size() - 1 + indices_shapes_.size();
real_output_shape_.resize(output_rank);
CHECK_CUDA_RET_WITH_ERROR(
cudaMemcpyAsync(&real_output_shape_[0], output_shape_device_address, output_rank * sizeof(int32_t),
cudaMemcpyDeviceToHost, reinterpret_cast<cudaStream_t>(stream_ptr)),
"Failed to copy gpu memory.");
} else {
auto input_dim1 = input_shapes_[IntToSize(axis_)];
CalGatherV2StaticShape(input_addr, indices_addr, output_addr, dims_[0], dims_[1], dims_[2], input_dim1,
reinterpret_cast<cudaStream_t>(stream_ptr));
CHECK_CUDA_RET_WITH_EXCEPT(cudaDeviceSynchronize(), "cudaDeviceSyncFailed - GatherV2 - in dynamic mode");
Reshape();
}
auto input_dim1 = input_shapes_[IntToSize(axis_)];
GatherV2(input_addr, indices_addr, output_addr, dims_[0], dims_[1], dims_[2], input_dim1,
reinterpret_cast<cudaStream_t>(stream_ptr));
return true;
}
bool Init(const CNodePtr &kernel_node) override {
@ -87,33 +59,24 @@ class GatherV2GpuFwdKernel : public GpuKernel {
size_t input_num = AnfAlgo::GetInputTensorNum(kernel_node);
if (input_num == 3) {
is_dynamic_shape_ = true;
} else if (input_num != 2) {
MS_LOG(EXCEPTION) << "Argument number is " << input_num << ", but GatherGpuV2FwdKernel needs 2.";
MS_LOG(INFO) << " GatherGpuV2FwdKernel running in Dynamic Mode.";
} else if (input_num == 2) {
MS_LOG(INFO) << " GatherGpuV2FwdKernel running in Normal Mode.";
} else {
MS_LOG(EXCEPTION) << "Argument number is " << input_num << ", but GatherGpuV2FwdKernel needs 2 or 3.";
}
input_shapes_ = AnfAlgo::GetInputRealDeviceShapeIfExist(kernel_node, 0);
indices_shapes_ = AnfAlgo::GetInputRealDeviceShapeIfExist(kernel_node, 1);
output_shapes_ = AnfAlgo::GetOutputRealDeviceShapeIfExist(kernel_node, 0);
if (is_dynamic_shape_) {
c_node_ptr_ = kernel_node;
size_t input_shape_min = *std::min_element(input_shapes_.begin(), input_shapes_.end());
max_output_size_ = (GetSize(input_shapes_) / input_shape_min) * GetSize(indices_shapes_);
} else {
if (!is_dynamic_shape_) {
axis_ = static_cast<int>(GetAttr<int64_t>(kernel_node, "axis"));
if (axis_ < 0) {
axis_ = axis_ + SizeToInt(input_shapes_.size());
}
Reshape();
}
InitSizeLists();
return true;
}
void ResetResource() noexcept override {
is_dynamic_shape_ = false;
max_output_size_ = -1;
input_shapes_.clear();
indices_shapes_.clear();
output_shapes_.clear();
@ -128,52 +91,32 @@ class GatherV2GpuFwdKernel : public GpuKernel {
void InitSizeLists() override {
size_t size = GetSize(input_shapes_);
input_size_list_.push_back(size);
size = GetSize(indices_shapes_);
input_size_list_.push_back(size);
if (is_dynamic_shape_) {
// add by chenweifeng
input_size_list_.push_back(sizeof(S));
// allocate maximum size needed
output_size_list_.push_back(max_output_size_);
// allocate workspace memory for input, indices, axis, and output shape respectively
size = GetSize(input_shapes_);
workspace_size_list_.push_back(size);
size = GetSize(indices_shapes_);
workspace_size_list_.push_back(size);
size = sizeof(int32_t);
workspace_size_list_.push_back(size);
size = GetSize(input_shapes_);
workspace_size_list_.push_back(size);
} else {
size = GetSize(output_shapes_);
output_size_list_.push_back(size);
input_size_list_.push_back(sizeof(int64_t));
}
size = GetSize(output_shapes_);
output_size_list_.push_back(size);
}
private:
void Reshape() {
if (axis_ < 0) {
axis_ = axis_ + SizeToInt(input_shapes_.size());
}
size_t dim_before_axis = 1;
for (size_t i = 0; i < IntToSize(axis_); i++) {
dim_before_axis *= output_shapes_[i];
}
size_t dim_of_indices = 1;
for (size_t i = 0; i < indices_shapes_.size(); i++) {
dim_of_indices *= indices_shapes_[i];
}
size_t dim_after_indices = 1;
for (size_t i = IntToSize(axis_) + indices_shapes_.size(); i < output_shapes_.size(); i++) {
dim_after_indices *= output_shapes_[i];
}
dims_[0] = dim_before_axis;
dims_[1] = dim_of_indices;
dims_[2] = dim_after_indices;
@ -193,14 +136,9 @@ class GatherV2GpuFwdKernel : public GpuKernel {
std::vector<size_t> input_shapes_;
std::vector<size_t> indices_shapes_;
std::vector<size_t> output_shapes_;
size_t dims_[3] = {};
int64_t axis_;
bool is_dynamic_shape_;
int max_output_size_;
std::vector<size_t> real_output_shape_;
CNodePtr c_node_ptr_;
std::vector<size_t> input_size_list_;
std::vector<size_t> output_size_list_;
std::vector<size_t> workspace_size_list_;
@ -208,4 +146,4 @@ class GatherV2GpuFwdKernel : public GpuKernel {
} // namespace kernel
} // namespace mindspore
#endif // MINDSPORE_GATHER_GPU_KERNEL_H
#endif // MINDSPORE_CCSRC_BACKEND_KERNEL_COMPILER_GPU_GATHER_V2_GPU_KERNEL_H_

91
mindspore/ccsrc/backend/kernel_compiler/gpu/cuda_impl/gatherv2.cu
Unescape View File

@ -18,7 +18,7 @@
#include "backend/kernel_compiler/gpu/cuda_impl/gatherv2.cuh"
#include "runtime/device/gpu/cuda_common.h"
template <typename T, typename S>
__device__ void GatherV2Kernel(T *input, S *indices, T *output, size_t output_dim0, size_t output_dim1,
__global__ void GatherV2Kernel(T *input, S *indices, T *output, size_t output_dim0, size_t output_dim1,
size_t output_dim2, size_t input_dim1) {
int num = output_dim0 * output_dim1 * output_dim2;
int i, j, k;
@ -38,90 +38,17 @@ __device__ void GatherV2Kernel(T *input, S *indices, T *output, size_t output_di
return;
}
template <typename T, typename S>
__global__ void GatherV2StaticShapeWrapper(T *input, S *indices, T *output, size_t output_dim0, size_t output_dim1,
size_t output_dim2, size_t input_dim1) {
GatherV2Kernel(input, indices, output, output_dim0, output_dim1, output_dim2, input_dim1);
}
template <typename T, typename S>
__global__ void GatherV2DynamicShape(T *input, S *indices, T *output, size_t *input_shape_wksp, size_t input_rank,
size_t *indices_shape_wksp, size_t indices_rank, int64_t *axis_wksp,
size_t *output_shape_wksp, const int max_output_size) {
int gt_id = blockIdx.x * blockDim.x + threadIdx.x;
size_t axis = (size_t)(*axis_wksp);
int output_shape_index = 0;
size_t output_dim0 = 1;
for (size_t i = 0; i < axis; i++) {
output_dim0 *= input_shape_wksp[i];
if (gt_id == 0) {
output_shape_wksp[output_shape_index] = input_shape_wksp[i];
output_shape_index++;
}
}
size_t output_dim1 = 1;
for (size_t i = 0; i < indices_rank; i++) {
output_dim1 *= indices_shape_wksp[i];
if (gt_id == 0) {
output_shape_wksp[output_shape_index] = indices_shape_wksp[i];
output_shape_index++;
}
}
size_t output_dim2 = 1;
for (size_t i = axis + 1; i < input_rank; i++) {
output_dim2 *= indices_shape_wksp[i];
if (gt_id == 0) {
output_shape_wksp[output_shape_index] = input_shape_wksp[i];
output_shape_index++;
}
}
size_t input_dim1 = (size_t)(input_shape_wksp[axis]);
GatherV2Kernel(input, indices, output, output_dim0, output_dim1, output_dim2, input_dim1);
}
// entry points from gpu kernel's .h file
template <typename T, typename S>
void CalGatherV2StaticShape(T *input, S *indices, T *output, size_t output_dim0, size_t output_dim1, size_t output_dim2,
size_t input_dim1, cudaStream_t stream) {
void GatherV2(T *input, S *indices, T *output, size_t output_dim0, size_t output_dim1, size_t output_dim2,
size_t input_dim1, cudaStream_t stream) {
int size = output_dim0 * output_dim1 * output_dim2;
GatherV2StaticShapeWrapper<<<GET_BLOCKS(size), GET_THREADS, 0, stream>>>(input, indices, output, output_dim0,
output_dim1, output_dim2, input_dim1);
GatherV2Kernel<<<GET_BLOCKS(size), GET_THREADS, 0, stream>>>(input, indices, output, output_dim0, output_dim1,
output_dim2, input_dim1);
return;
}
template <typename T, typename S>
void CalGatherV2DynamicShape(T *input, S *indices, T *output, size_t *input_shape_wksp, size_t input_rank,
size_t *indices_shape_wksp, size_t indices_rank, int64_t *axis_wksp,
size_t *output_shape_wksp, const int max_output_size, cudaStream_t stream) {
GatherV2DynamicShape<<<GET_BLOCKS(max_output_size), GET_THREADS, 0, stream>>>(
input, indices, output, input_shape_wksp, input_rank, indices_shape_wksp, indices_rank, axis_wksp,
output_shape_wksp, max_output_size);
}
// template instantiations
template void CalGatherV2StaticShape<float, int>(float *input, int *indices, float *output, size_t output_dim0,
size_t output_dim1, size_t output_dim2, size_t input_dim1,
cudaStream_t stream);
template void CalGatherV2StaticShape<half, int>(half *input, int *indices, half *output, size_t output_dim0,
size_t output_dim1, size_t output_dim2, size_t input_dim1,
cudaStream_t stream);
template void CalGatherV2DynamicShape<float, int>(float *input, int *indices, float *output, size_t *input_shape_wksp,
size_t input_rank, size_t *indices_shape_wksp, size_t indices_rank,
int64_t *axis_wksp, size_t *output_shape_wksp,
const int max_output_size, cudaStream_t stream);
template void GatherV2<float, int>(float *input, int *indices, float *output, size_t output_dim0, size_t output_dim1,
size_t output_dim2, size_t input_dim1, cudaStream_t stream);
template void CalGatherV2DynamicShape<half, int>(half *input, int *indices, half *output, size_t *input_shape_wksp,
size_t input_rank, size_t *indices_shape_wksp, size_t indices_rank,
int64_t *axis_wksp, size_t *output_shape_wksp,
const int max_output_size, cudaStream_t stream);
template void GatherV2<half, int>(half *input, int *indices, half *output, size_t output_dim0, size_t output_dim1,
size_t output_dim2, size_t input_dim1, cudaStream_t stream);

12
mindspore/ccsrc/backend/kernel_compiler/gpu/cuda_impl/gatherv2.cuh
Unescape View File

@ -14,14 +14,10 @@
* limitations under the License.
*/
#ifndef MINDSPORE_GATHER_GPU_CU_H
#define MINDSPORE_GATHER_GPU_CU_H
#ifndef MINDSPORE_CCSRC_KERNEL_GPU_CUDA_IMP_GATHER_V2_CU_H_
#define MINDSPORE_CCSRC_KERNEL_GPU_CUDA_IMP_GATHER_V2_CU_H_
template <typename T, typename S>
void CalGatherV2StaticShape(T *input, S *indices, T *output, size_t output_dim0, size_t output_dim1, size_t output_dim2,
size_t input_dim1, cudaStream_t stream);
void GatherV2(T *input, S *indices, T *output, size_t output_dim0, size_t output_dim1, size_t output_dim2,
size_t input_dim1, cudaStream_t stream);
template <typename T, typename S>
void CalGatherV2DynamicShape(T *input, S *indices, T *output, size_t *input_shape_wksp, size_t input_rank,
size_t *indices_shape_wksp, size_t indices_rank, int64_t *axis_wksp,
size_t *output_shape_wksp, const int max_output_size, cudaStream_t stream);
#endif

34
mindspore/core/abstract/prim_arrays.cc
Unescape View File

@ -408,7 +408,8 @@ AbstractBasePtr InferImplGatherV2(const AnalysisEnginePtr &, const PrimitivePtr
CheckArgsSize(op_name, args_spec_list, 3);
AbstractTensorPtr params = CheckArg<AbstractTensor>(op_name, args_spec_list, 0);
AbstractTensorPtr indices = CheckArg<AbstractTensor>(op_name, args_spec_list, 1);
bool ind_dyn = (!indices->shape()->min_shape().empty() && !indices->shape()->max_shape().empty());
bool param_dyn = (!params->shape()->min_shape().empty() && !params->shape()->max_shape().empty());
int64_t axis_val = 0;
// 3rd input is a Tensor when GatherV2 is a dynamic shape operator
if (args_spec_list[2]->isa<AbstractTensor>()) {
@ -425,31 +426,36 @@ AbstractBasePtr InferImplGatherV2(const AnalysisEnginePtr &, const PrimitivePtr
} else {
MS_LOG(EXCEPTION) << "Invalid abstract type:" << args_spec_list[2]->type_name();
}
auto params_shp = params->shape()->shape();
auto indices_shp = indices->shape()->shape();
auto params_rank = static_cast<int64_t>(params_shp.size());
// either inputs or both can be dynamic and computation requires min/max shapes for both
ShapeVector param_shp_min = (param_dyn) ? params->shape()->min_shape() : params->shape()->shape();
ShapeVector param_shp_max = (param_dyn) ? params->shape()->max_shape() : params->shape()->shape();
ShapeVector indices_shp_min = (ind_dyn) ? indices->shape()->min_shape() : indices->shape()->shape();
ShapeVector indices_shp_max = (ind_dyn) ? indices->shape()->max_shape() : indices->shape()->shape();
// check axis_val within interval: [-params_rank, params_rank)
if (!(-params_rank <= axis_val) || !(axis_val < params_rank)) {
MS_LOG(EXCEPTION) << "For GatherV2 - Axis value must be within [ " << -params_rank << ", " << params_rank << " ) "
<< "Got " << axis_val << ".";
}
if (axis_val < 0) {
axis_val += params_rank;
}
auto calc_shape = [axis_val, &params_shp](const ShapeVector &inp_vec) -> ShapeVector {
auto calc_shape = [axis_val](const ShapeVector &ind_vec, const ShapeVector &params_vec) -> ShapeVector {
ShapeVector out_vec;
std::copy(params_shp.begin(), params_shp.begin() + axis_val, std::back_inserter(out_vec));
copy(inp_vec.begin(), inp_vec.end(), std::back_inserter(out_vec));
copy(params_shp.begin() + axis_val + 1, params_shp.end(), std::back_inserter(out_vec));
std::copy(params_vec.begin(), params_vec.begin() + axis_val, std::back_inserter(out_vec));
copy(ind_vec.begin(), ind_vec.end(), std::back_inserter(out_vec));
copy(params_vec.begin() + axis_val + 1, params_vec.end(), std::back_inserter(out_vec));
return out_vec;
};
ShapeVector out_shape = calc_shape(indices_shp);
if (!indices->shape()->min_shape().empty() && !indices->shape()->max_shape().empty()) {
ShapeVector min_shape = calc_shape(indices->shape()->min_shape());
ShapeVector max_shape = calc_shape(indices->shape()->max_shape());
ShapeVector out_shape = calc_shape(indices_shp, params_shp);
if (ind_dyn || param_dyn) {
ShapeVector min_shape = calc_shape(indices_shp_min, param_shp_min);
ShapeVector max_shape = calc_shape(indices_shp_max, param_shp_max);
return std::make_shared<AbstractTensor>(params->element(),
std::make_shared<Shape>(out_shape, min_shape, max_shape));
}
return std::make_shared<AbstractTensor>(params->element(), std::make_shared<Shape>(out_shape));
}

2
mindspore/core/abstract/prim_others.cc
Unescape View File

@ -535,7 +535,7 @@ AbstractBasePtr InferImplGpuConvertToDynamicShape(const AnalysisEnginePtr &, con
ShapeVector input_shape = input->shape()->shape();
int32_t input_rank = input_shape.size();
ShapeVector inferred_shape(input_rank, Shape::SHP_ANY);
ShapeVector min_shape = {1};
ShapeVector min_shape(input_rank, 1);
ShapeVector max_shape = input_shape;
ShapePtr shape = std::make_shared<Shape>(inferred_shape, min_shape, max_shape);

18
mindspore/ops/operations/array_ops.py
Unescape View File

@ -703,7 +703,6 @@ class GatherV2(PrimitiveWithCheck):
[ 4. 54.]
[ 2. 55.]]
"""
@prim_attr_register
def __init__(self):
"""Initialize index_select"""
@ -713,22 +712,7 @@ class GatherV2(PrimitiveWithCheck):
def __check__(self, params, indices, axis):
validator.check_subclass("params", params['dtype'], mstype.tensor, self.name)
validator.check_tensor_dtype_valid("indices", indices['dtype'], mstype.int_type, self.name)
validator.check_subclass("axis", axis['dtype'], mstype.int_, self.name)
axis_v = axis['value']
params_shp = params['shape']
rank = len(params_shp)
validator.check_int_range(axis_v, -rank, rank, Rel.INC_LEFT, "axis", self.name)
if axis_v < 0:
axis_v += rank
out_shape = params_shp[:axis_v] + indices['shape'] + params_shp[axis_v + 1:]
out = {'shape': out_shape,
'dtype': params['dtype'],
'value': None}
if 'min_shape' in indices and 'max_shape' in indices:
out['min_shape'] = params_shp[:axis_v] + indices['min_shape'] + params_shp[axis_v + 1:]
out['max_shape'] = params_shp[:axis_v] + indices['max_shape'] + params_shp[axis_v + 1:]
return out
validator.check_subclass("axis", axis['dtype'], [mstype.tensor, mstype.int_], self.name)
class SparseGatherV2(GatherV2):

156
tests/st/ops/gpu/test_gatherV2_op.py
Unescape View File

@ -19,6 +19,7 @@ import pytest
import mindspore.context as context
import mindspore.nn as nn
from mindspore import Tensor
from mindspore.ops.operations import _inner_ops as inner
from mindspore.ops import operations as P
@ -937,3 +938,158 @@ def test_gather2():
diff = output.asnumpy() - expect
assert np.all(diff < error)
assert np.all(-diff < error)
# Dynamic Shape testing ahead
class GatherNetDynamic1(nn.Cell):
def __init__(self):
super(GatherNetDynamic1, self).__init__()
self.gather = P.GatherV2()
self.gpu_convert_to_dynamic_shape = inner.GpuConvertToDynamicShape()
def construct(self, x, indices):
# Testing only second input dynamic
indices_dyn = self.gpu_convert_to_dynamic_shape(indices)
return self.gather(x, indices_dyn, 0)
@pytest.mark.level0
@pytest.mark.platform_x86_gpu_training
@pytest.mark.env_onecard
def test_gather_dynamic_1():
x = Tensor(np.array([[4., 5., 4., 1., 5.,],
[4., 9., 5., 6., 4.,],
[9., 8., 4., 3., 6.,],
[0., 4., 2., 2., 8.,],
[1., 8., 6., 2., 8.,],
[8., 1., 9., 7., 3.,],
[7., 9., 2., 5., 7.,],
[9., 8., 6., 8., 5.,],
[3., 7., 2., 7., 4.,],
[4., 2., 8., 2., 9.,]]
).astype(np.float32))
indices = Tensor(np.array([[4000, 1, 300000]]).astype(np.int32))
expect = np.array([[[0., 0., 0., 0., 0.],
[4., 9., 5., 6., 4.],
[0., 0., 0., 0., 0.]]])
context.set_context(mode=context.GRAPH_MODE, device_target="GPU")
gather = GatherNetDynamic1()
output = gather(x, indices)
error = np.ones(shape=output.asnumpy().shape) * 1.0e-6
diff = output.asnumpy() - expect
assert np.all(diff < error)
assert np.all(-diff < error)
class GatherNetDynamic2(nn.Cell):
def __init__(self):
super(GatherNetDynamic2, self).__init__()
self.gather = P.GatherV2()
self.gpu_convert_to_dynamic_shape = inner.GpuConvertToDynamicShape()
def construct(self, x, indices):
# Testing only first input dynamic
x_dyn = self.gpu_convert_to_dynamic_shape(x)
return self.gather(x_dyn, indices, -1)
@pytest.mark.level0
@pytest.mark.platform_x86_gpu_training
@pytest.mark.env_onecard
def test_gather_dynamic_2():
x = Tensor(np.arange(2 * 3 * 4 * 5, dtype=np.float32).reshape(2, 3, 4, 5))
indices = Tensor(np.array([1, 3, 4], dtype='i4'))
expect = np.array([[[[1., 3., 4.],
[6., 8., 9.],
[11., 13., 14.],
[16., 18., 19.]],
[[21., 23., 24.],
[26., 28., 29.],
[31., 33., 34.],
[36., 38., 39.]],
[[41., 43., 44.],
[46., 48., 49.],
[51., 53., 54.],
[56., 58., 59.]]],
[[[61., 63., 64.],
[66., 68., 69.],
[71., 73., 74.],
[76., 78., 79.]],
[[81., 83., 84.],
[86., 88., 89.],
[91., 93., 94.],
[96., 98., 99.]],
[[101., 103., 104.],
[106., 108., 109.],
[111., 113., 114.],
[116., 118., 119.]]]])
context.set_context(mode=context.GRAPH_MODE, device_target="GPU")
gather = GatherNetDynamic2()
output = gather(x, indices)
error = np.ones(shape=output.asnumpy().shape) * 1.0e-6
diff = output.asnumpy() - expect
assert np.all(diff < error)
assert np.all(-diff < error)
class GatherNetDynamic3(nn.Cell):
def __init__(self):
super(GatherNetDynamic3, self).__init__()
self.gather = P.GatherV2()
self.gpu_convert_to_dynamic_shape = inner.GpuConvertToDynamicShape()
def construct(self, x, indices):
# Testing both inputs dynamic shapes
x_dyn = self.gpu_convert_to_dynamic_shape(x)
indices_dyn = self.gpu_convert_to_dynamic_shape(indices)
return self.gather(x_dyn, indices_dyn, -1)
@pytest.mark.level0
@pytest.mark.platform_x86_gpu_training
@pytest.mark.env_onecard
def test_gather_dynamic_3():
x = Tensor(np.arange(2 * 3 * 4 * 5, dtype=np.float32).reshape(2, 3, 4, 5))
indices = Tensor(np.array([1, 3, 4], dtype='i4'))
expect = np.array([[[[1., 3., 4.],
[6., 8., 9.],
[11., 13., 14.],
[16., 18., 19.]],
[[21., 23., 24.],
[26., 28., 29.],
[31., 33., 34.],
[36., 38., 39.]],
[[41., 43., 44.],
[46., 48., 49.],
[51., 53., 54.],
[56., 58., 59.]]],
[[[61., 63., 64.],
[66., 68., 69.],
[71., 73., 74.],
[76., 78., 79.]],
[[81., 83., 84.],
[86., 88., 89.],
[91., 93., 94.],
[96., 98., 99.]],
[[101., 103., 104.],
[106., 108., 109.],
[111., 113., 114.],
[116., 118., 119.]]]])
context.set_context(mode=context.GRAPH_MODE, device_target="GPU")
gather = GatherNetDynamic3()
output = gather(x, indices)
error = np.ones(shape=output.asnumpy().shape) * 1.0e-6
diff = output.asnumpy() - expect
assert np.all(diff < error)
assert np.all(-diff < error)

Write Preview
Loading…
Cancel
Save