| @@ -34,12 +34,14 @@ OpParameter *PopulateSpaceToBatchNDParameter(const mindspore::lite::PrimitiveC * | |||||
| space_batch_param_nd->m_ = block_sizes.size(); | space_batch_param_nd->m_ = block_sizes.size(); | ||||
| if (block_sizes.size() > std::numeric_limits<size_t>::max() / sizeof(int)) { | if (block_sizes.size() > std::numeric_limits<size_t>::max() / sizeof(int)) { | ||||
| MS_LOG(ERROR) << "The value of block_sizes.size() is too big"; | MS_LOG(ERROR) << "The value of block_sizes.size() is too big"; | ||||
| free(space_batch_param_nd); | |||||
| return nullptr; | return nullptr; | ||||
| } | } | ||||
| memcpy(space_batch_param_nd->block_sizes_, (block_sizes.data()), block_sizes.size() * sizeof(int)); | memcpy(space_batch_param_nd->block_sizes_, (block_sizes.data()), block_sizes.size() * sizeof(int)); | ||||
| auto paddings = ((mindspore::lite::SpaceToBatchND *)primitive)->GetPaddings(); | auto paddings = ((mindspore::lite::SpaceToBatchND *)primitive)->GetPaddings(); | ||||
| if (paddings.size() > std::numeric_limits<size_t>::max() / sizeof(int)) { | if (paddings.size() > std::numeric_limits<size_t>::max() / sizeof(int)) { | ||||
| MS_LOG(ERROR) << "The value of paddings.size() is too big"; | MS_LOG(ERROR) << "The value of paddings.size() is too big"; | ||||
| free(space_batch_param_nd); | |||||
| return nullptr; | return nullptr; | ||||
| } | } | ||||
| memcpy(space_batch_param_nd->paddings_, (paddings.data()), paddings.size() * sizeof(int)); | memcpy(space_batch_param_nd->paddings_, (paddings.data()), paddings.size() * sizeof(int)); | ||||
| @@ -36,12 +36,14 @@ OpParameter *PopulateSpaceToBatchParameter(const mindspore::lite::PrimitiveC *pr | |||||
| space_batch_param->m_ = block_sizes.size(); | space_batch_param->m_ = block_sizes.size(); | ||||
| if (block_sizes.size() > std::numeric_limits<size_t>::max() / sizeof(int)) { | if (block_sizes.size() > std::numeric_limits<size_t>::max() / sizeof(int)) { | ||||
| MS_LOG(ERROR) << "The value of block_sizes.size() is too big"; | MS_LOG(ERROR) << "The value of block_sizes.size() is too big"; | ||||
| free(space_batch_param); | |||||
| return nullptr; | return nullptr; | ||||
| } | } | ||||
| memcpy(space_batch_param->block_sizes_, (block_sizes.data()), block_sizes.size() * sizeof(int)); | memcpy(space_batch_param->block_sizes_, (block_sizes.data()), block_sizes.size() * sizeof(int)); | ||||
| auto paddings = ((mindspore::lite::SpaceToBatch *)primitive)->Paddings(); | auto paddings = ((mindspore::lite::SpaceToBatch *)primitive)->Paddings(); | ||||
| if (paddings.size() > std::numeric_limits<size_t>::max() / sizeof(int)) { | if (paddings.size() > std::numeric_limits<size_t>::max() / sizeof(int)) { | ||||
| MS_LOG(ERROR) << "The value of paddings.size() is too big"; | MS_LOG(ERROR) << "The value of paddings.size() is too big"; | ||||
| free(space_batch_param); | |||||
| return nullptr; | return nullptr; | ||||
| } | } | ||||
| memcpy(space_batch_param->paddings_, (paddings.data()), paddings.size() * sizeof(int)); | memcpy(space_batch_param->paddings_, (paddings.data()), paddings.size() * sizeof(int)); | ||||
| @@ -37,24 +37,28 @@ OpParameter *PopulateStridedSliceParameter(const mindspore::lite::PrimitiveC *pr | |||||
| auto begin = ((lite::StridedSlice *)primitive)->GetBegins(); | auto begin = ((lite::StridedSlice *)primitive)->GetBegins(); | ||||
| if (begin.size() > std::numeric_limits<size_t>::max() / sizeof(int)) { | if (begin.size() > std::numeric_limits<size_t>::max() / sizeof(int)) { | ||||
| MS_LOG(ERROR) << "The value of begin.size() is too big"; | MS_LOG(ERROR) << "The value of begin.size() is too big"; | ||||
| free(strided_slice_param); | |||||
| return nullptr; | return nullptr; | ||||
| } | } | ||||
| memcpy(strided_slice_param->begins_, (begin.data()), begin.size() * sizeof(int)); | memcpy(strided_slice_param->begins_, (begin.data()), begin.size() * sizeof(int)); | ||||
| auto end = ((lite::StridedSlice *)primitive)->GetEnds(); | auto end = ((lite::StridedSlice *)primitive)->GetEnds(); | ||||
| if (end.size() > std::numeric_limits<size_t>::max() / sizeof(int)) { | if (end.size() > std::numeric_limits<size_t>::max() / sizeof(int)) { | ||||
| MS_LOG(ERROR) << "The value of end.size() is too big"; | MS_LOG(ERROR) << "The value of end.size() is too big"; | ||||
| free(strided_slice_param); | |||||
| return nullptr; | return nullptr; | ||||
| } | } | ||||
| memcpy(strided_slice_param->ends_, (end.data()), end.size() * sizeof(int)); | memcpy(strided_slice_param->ends_, (end.data()), end.size() * sizeof(int)); | ||||
| auto stride = ((lite::StridedSlice *)primitive)->GetStrides(); | auto stride = ((lite::StridedSlice *)primitive)->GetStrides(); | ||||
| if (stride.size() > std::numeric_limits<size_t>::max() / sizeof(int)) { | if (stride.size() > std::numeric_limits<size_t>::max() / sizeof(int)) { | ||||
| MS_LOG(ERROR) << "The value of stride.size() is too big"; | MS_LOG(ERROR) << "The value of stride.size() is too big"; | ||||
| free(strided_slice_param); | |||||
| return nullptr; | return nullptr; | ||||
| } | } | ||||
| memcpy(strided_slice_param->strides_, (stride.data()), stride.size() * sizeof(int)); | memcpy(strided_slice_param->strides_, (stride.data()), stride.size() * sizeof(int)); | ||||
| auto in_shape = ((lite::StridedSlice *)primitive)->GetInShape(); | auto in_shape = ((lite::StridedSlice *)primitive)->GetInShape(); | ||||
| if (in_shape.size() > std::numeric_limits<size_t>::max() / sizeof(int)) { | if (in_shape.size() > std::numeric_limits<size_t>::max() / sizeof(int)) { | ||||
| MS_LOG(ERROR) << "The value of in_shape.size() is too big"; | MS_LOG(ERROR) << "The value of in_shape.size() is too big"; | ||||
| free(strided_slice_param); | |||||
| return nullptr; | return nullptr; | ||||
| } | } | ||||
| memcpy(strided_slice_param->in_shape_, (in_shape.data()), in_shape.size() * sizeof(int)); | memcpy(strided_slice_param->in_shape_, (in_shape.data()), in_shape.size() * sizeof(int)); | ||||
| @@ -63,13 +63,13 @@ int SpaceToDepth::InferShape(std::vector<lite::Tensor *> inputs, std::vector<lit | |||||
| MS_ASSERT(this->primitive_ != nullptr); | MS_ASSERT(this->primitive_ != nullptr); | ||||
| if (outputs.size() != kSpaceToDepthOutputNum || inputs.size() != kSpaceToDepthInputNum) { | if (outputs.size() != kSpaceToDepthOutputNum || inputs.size() != kSpaceToDepthInputNum) { | ||||
| MS_LOG(ERROR) << "Invalid output/input size! output size: " << outputs.size() << ",input size: " << inputs.size(); | MS_LOG(ERROR) << "Invalid output/input size! output size: " << outputs.size() << ",input size: " << inputs.size(); | ||||
| return 1; | |||||
| return RET_ERROR; | |||||
| } | } | ||||
| auto input = inputs.at(0); | auto input = inputs.at(0); | ||||
| if (input->format() != schema::Format::Format_NHWC) { | if (input->format() != schema::Format::Format_NHWC) { | ||||
| MS_LOG(ERROR) << "space_to_depth only support NHWC now!"; | MS_LOG(ERROR) << "space_to_depth only support NHWC now!"; | ||||
| return 1; | |||||
| return RET_ERROR; | |||||
| } | } | ||||
| outputs.at(0)->set_format(input->format()); | outputs.at(0)->set_format(input->format()); | ||||
| outputs.at(0)->set_data_type(input->data_type()); | outputs.at(0)->set_data_type(input->data_type()); | ||||
| @@ -79,14 +79,18 @@ int SpaceToDepth::InferShape(std::vector<lite::Tensor *> inputs, std::vector<lit | |||||
| auto input_shape = input->shape(); | auto input_shape = input->shape(); | ||||
| if (input_shape.size() != kDimension_4d) { | if (input_shape.size() != kDimension_4d) { | ||||
| MS_LOG(ERROR) << "input shape dimension size should == " << kDimension_4d; | MS_LOG(ERROR) << "input shape dimension size should == " << kDimension_4d; | ||||
| return 1; | |||||
| return RET_ERROR; | |||||
| } | } | ||||
| int32_t block_size = GetBlockSize(); | int32_t block_size = GetBlockSize(); | ||||
| if (block_size == 0) { | |||||
| MS_LOG(ERROR) << "block_size is zero"; | |||||
| return RET_ERROR; | |||||
| } | |||||
| if (input_shape.at(NHWC_H) % block_size != 0 || input_shape.at(NHWC_H) == 0 || | if (input_shape.at(NHWC_H) % block_size != 0 || input_shape.at(NHWC_H) == 0 || | ||||
| input_shape.at(NHWC_W) % block_size != 0 || input_shape.at(NHWC_W) == 0) { | input_shape.at(NHWC_W) % block_size != 0 || input_shape.at(NHWC_W) == 0) { | ||||
| MS_LOG(ERROR) << "input dimension h or w size error!"; | MS_LOG(ERROR) << "input dimension h or w size error!"; | ||||
| return 1; | |||||
| return RET_ERROR; | |||||
| } | } | ||||
| std::vector<int32_t> output_shape(input_shape.size()); | std::vector<int32_t> output_shape(input_shape.size()); | ||||
| output_shape.at(NHWC_N) = input_shape.at(NHWC_N); | output_shape.at(NHWC_N) = input_shape.at(NHWC_N); | ||||