mirror of
https://github.com/opencv/opencv.git
synced 2026-07-30 07:43:03 +04:00
Merge pull request #22529 from fengyuentau:scatter_scatternd
DNN: supports Scatter and ScatterND from ONNX
This commit is contained in:
Alexander Smorkalov
GitHub
@@ -181,6 +181,9 @@ void initializeLayerFactory()
|
||||
CV_DNN_REGISTER_LAYER_CLASS(GRU, GRULayer);
|
||||
CV_DNN_REGISTER_LAYER_CLASS(CumSum, CumSumLayer);
|
||||
|
||||
CV_DNN_REGISTER_LAYER_CLASS(Scatter, ScatterLayer);
|
||||
CV_DNN_REGISTER_LAYER_CLASS(ScatterND, ScatterNDLayer);
|
||||
|
||||
CV_DNN_REGISTER_LAYER_CLASS(Quantize, QuantizeLayer);
|
||||
CV_DNN_REGISTER_LAYER_CLASS(Dequantize, DequantizeLayer);
|
||||
CV_DNN_REGISTER_LAYER_CLASS(Requantize, RequantizeLayer);
|
||||
|
||||
@@ -0,0 +1,202 @@
|
||||
// This file is part of OpenCV project.
|
||||
// It is subject to the license terms in the LICENSE file found in the top-level directory
|
||||
// of this distribution and at http://opencv.org/license.html.
|
||||
|
||||
#include "../precomp.hpp"
|
||||
#include "layers_common.hpp"
|
||||
|
||||
#include <algorithm> // for std::max & std::min
|
||||
|
||||
namespace cv { namespace dnn {
|
||||
|
||||
class ScatterNDLayerImpl CV_FINAL : public ScatterNDLayer
|
||||
{
|
||||
public:
|
||||
enum class REDUCTION
|
||||
{
|
||||
NONE = 1,
|
||||
ADD,
|
||||
MUL,
|
||||
MAX,
|
||||
MIN
|
||||
} reduction;
|
||||
|
||||
ScatterNDLayerImpl(const LayerParams& params)
|
||||
{
|
||||
setParamsFrom(params);
|
||||
|
||||
String reduction_name = toLowerCase(params.get<String>("reduction", "none"));
|
||||
if (reduction_name == "none")
|
||||
reduction = REDUCTION::NONE;
|
||||
else if (reduction_name == "add")
|
||||
reduction = REDUCTION::ADD;
|
||||
else if (reduction_name == "mul")
|
||||
reduction = REDUCTION::MUL;
|
||||
else if (reduction_name == "max")
|
||||
reduction = REDUCTION::MAX;
|
||||
else if (reduction_name == "min")
|
||||
reduction = REDUCTION::MIN;
|
||||
else
|
||||
CV_Error(cv::Error::StsBadArg, "Unkown reduction \"" + reduction_name + "\"");
|
||||
}
|
||||
|
||||
virtual bool supportBackend(int backendId) CV_OVERRIDE
|
||||
{
|
||||
return backendId == DNN_BACKEND_OPENCV;
|
||||
}
|
||||
|
||||
virtual bool getMemoryShapes(const std::vector<MatShape> &inputs,
|
||||
const int requiredOutputs,
|
||||
std::vector<MatShape> &outputs,
|
||||
std::vector<MatShape> &internals) const CV_OVERRIDE
|
||||
{
|
||||
CV_CheckEQ(inputs.size(), 3ull, "ScatterND: require three inputs.");
|
||||
|
||||
size_t r = inputs[0].size(), q = inputs[1].size(), p = inputs[2].size(), k = inputs[1].back();
|
||||
CV_CheckEQ(r + q - inputs[1].back() - 1, p, "ScatterND: updates should have rank of data.dims + indices.dims - indices.size[-1] - 1");
|
||||
CV_CheckLE(k, r, "ScatterND: indices.shape[-1] must be less than (or equal to) the rank of input data.");
|
||||
|
||||
for (int i = 0; i < q - 1; i++) // np.ndindex(indices.shape[-1])
|
||||
{
|
||||
CV_CheckEQ(inputs[2][i], inputs[1][i], "ScatterND: updates.shape[0 : rank(indices)-1] must equal to indices.shape[0 : rank(indices)-1].");
|
||||
}
|
||||
for (int i = q - 1, j = k, m = 0; i + m < p; m++)
|
||||
{
|
||||
CV_CheckEQ(inputs[2][i + m], inputs[0][j + m], "ScatterND: updates.shape[rank(indices)-1 : ] must equal to data[indices.shape[-1] : rank(data)-1].");
|
||||
}
|
||||
|
||||
outputs.assign(1, inputs[0]);
|
||||
return false;
|
||||
}
|
||||
|
||||
void forward(InputArrayOfArrays inputs_arr, OutputArrayOfArrays outputs_arr, OutputArrayOfArrays internals_arr) CV_OVERRIDE
|
||||
{
|
||||
CV_TRACE_FUNCTION();
|
||||
CV_TRACE_ARG_VALUE(name, "name", name.c_str());
|
||||
|
||||
std::vector<Mat> inputs, outputs;
|
||||
inputs_arr.getMatVector(inputs);
|
||||
outputs_arr.getMatVector(outputs);
|
||||
|
||||
const Mat& data = inputs[0];
|
||||
const Mat& indices = inputs[1];
|
||||
const Mat& updates = inputs[2];
|
||||
Mat& out = outputs[0];
|
||||
|
||||
typeDispatch(outputs[0].type(), data, indices, updates, out);
|
||||
}
|
||||
|
||||
// NOTE: This impl does not check whether indices have duplicate entries.
|
||||
// The last duplicate entry will overwrite the previous.
|
||||
template<typename T, typename Functor>
|
||||
void forward_impl(const Functor& rd, const Mat& data, const Mat& indices, const Mat& updates, Mat& out)
|
||||
{
|
||||
data.copyTo(out);
|
||||
|
||||
const int* shape = data.size.p;
|
||||
const size_t* step = data.step.p;
|
||||
|
||||
const int ind_ndims = indices.dims;
|
||||
const int* ind_shape = indices.size.p;
|
||||
const T* p_indices = indices.ptr<const T>();
|
||||
|
||||
const int upd_ndims = updates.dims;
|
||||
const int* upd_shape = updates.size.p;
|
||||
const T* p_updates = updates.ptr<const T>();
|
||||
|
||||
T* p_out = out.ptr<T>();
|
||||
|
||||
int k = ind_shape[ind_ndims - 1]; // last dim of indices
|
||||
size_t total = (size_t)(indices.total() / k);
|
||||
|
||||
size_t updates_size = 1;
|
||||
for (int i = ind_ndims - 1; i < upd_ndims; i++)
|
||||
updates_size *= upd_shape[i];
|
||||
|
||||
size_t inp_start_offset = 0;
|
||||
size_t ind_start_offset = 0;
|
||||
size_t upd_start_offset = 0;
|
||||
for (size_t i = 0; i < total; i++, ind_start_offset += k, upd_start_offset += updates_size)
|
||||
{
|
||||
const T* tmp_p_indices = p_indices + ind_start_offset;
|
||||
inp_start_offset = 0;
|
||||
for (int j = 0; j < k; j++)
|
||||
{
|
||||
CV_Assert(tmp_p_indices[j] < shape[j] && tmp_p_indices[j] > -shape[j]);
|
||||
inp_start_offset += (((int)tmp_p_indices[j] + shape[j]) % shape[j]) * step[j];
|
||||
}
|
||||
inp_start_offset /= sizeof(T);
|
||||
|
||||
const T* tmp_p_updates = p_updates + upd_start_offset;
|
||||
T* tmp_p_out = p_out + inp_start_offset;
|
||||
for (int j = 0; j < updates_size; j++)
|
||||
tmp_p_out[j] = rd(tmp_p_out[j], tmp_p_updates[j]);
|
||||
}
|
||||
}
|
||||
|
||||
template<typename... Args>
|
||||
inline void typeDispatch(const int type, Args&&... args)
|
||||
{
|
||||
switch (type)
|
||||
{
|
||||
case CV_8U:
|
||||
reductionDispatch<uint8_t>(std::forward<Args>(args)...);
|
||||
break;
|
||||
case CV_32S:
|
||||
reductionDispatch<int32_t>(std::forward<Args>(args)...);
|
||||
break;
|
||||
case CV_32F:
|
||||
reductionDispatch<float>(std::forward<Args>(args)...);
|
||||
break;
|
||||
default:
|
||||
CV_Error(cv::Error::BadDepth, "Unsupported type.");
|
||||
};
|
||||
}
|
||||
|
||||
template<typename T, typename... Args>
|
||||
inline void reductionDispatch(Args&&... args)
|
||||
{
|
||||
switch (reduction)
|
||||
{
|
||||
case REDUCTION::NONE:
|
||||
{
|
||||
auto rd = [](const T& a, const T& b) { return b; }; // a from input data, b from updates
|
||||
forward_impl<T>(rd, std::forward<Args>(args)...);
|
||||
break;
|
||||
}
|
||||
case REDUCTION::ADD:
|
||||
{
|
||||
auto rd = [](const T& a, const T& b) { return a + b; };
|
||||
forward_impl<T>(rd, std::forward<Args>(args)...);
|
||||
break;
|
||||
}
|
||||
case REDUCTION::MUL:
|
||||
{
|
||||
auto rd = [](const T& a, const T& b) { return a * b; };
|
||||
forward_impl<T>(rd, std::forward<Args>(args)...);
|
||||
break;
|
||||
}
|
||||
case REDUCTION::MAX:
|
||||
{
|
||||
auto rd = [](const T& a, const T& b) { return std::max(a, b); };
|
||||
forward_impl<T>(rd, std::forward<Args>(args)...);
|
||||
break;
|
||||
}
|
||||
case REDUCTION::MIN:
|
||||
{
|
||||
auto rd = [](const T& a, const T& b) { return std::min(a, b); };
|
||||
forward_impl<T>(rd, std::forward<Args>(args)...);
|
||||
break;
|
||||
}
|
||||
default:
|
||||
CV_Error(Error::StsBadArg, "Unsupported reduction.");
|
||||
};
|
||||
}
|
||||
};
|
||||
|
||||
Ptr<ScatterNDLayer> ScatterNDLayer::create(const LayerParams& params)
|
||||
{
|
||||
return makePtr<ScatterNDLayerImpl>(params);
|
||||
}
|
||||
|
||||
}} // namespace cv::dnn
|
||||
@@ -0,0 +1,208 @@
|
||||
// This file is part of OpenCV project.
|
||||
// It is subject to the license terms in the LICENSE file found in the top-level directory
|
||||
// of this distribution and at http://opencv.org/license.html.
|
||||
|
||||
#include "../precomp.hpp"
|
||||
#include "layers_common.hpp"
|
||||
|
||||
#include <algorithm> // for std::max & std::min
|
||||
|
||||
namespace cv { namespace dnn {
|
||||
|
||||
class ScatterLayerImpl CV_FINAL : public ScatterLayer
|
||||
{
|
||||
public:
|
||||
enum class REDUCTION
|
||||
{
|
||||
NONE = 1,
|
||||
ADD,
|
||||
MUL,
|
||||
MAX,
|
||||
MIN
|
||||
} reduction;
|
||||
|
||||
ScatterLayerImpl(const LayerParams& params)
|
||||
{
|
||||
setParamsFrom(params);
|
||||
|
||||
axis = params.get<int>("axis", 0);
|
||||
String reduction_name = toLowerCase(params.get<String>("reduction", "none"));
|
||||
if (reduction_name == "none")
|
||||
reduction = REDUCTION::NONE;
|
||||
else if (reduction_name == "add")
|
||||
reduction = REDUCTION::ADD;
|
||||
else if (reduction_name == "mul")
|
||||
reduction = REDUCTION::MUL;
|
||||
else if (reduction_name == "max")
|
||||
reduction = REDUCTION::MAX;
|
||||
else if (reduction_name == "min")
|
||||
reduction = REDUCTION::MIN;
|
||||
else
|
||||
CV_Error(cv::Error::StsBadArg, "Unkown reduction \"" + reduction_name + "\"");
|
||||
}
|
||||
|
||||
virtual bool supportBackend(int backendId) CV_OVERRIDE
|
||||
{
|
||||
return backendId == DNN_BACKEND_OPENCV;
|
||||
}
|
||||
|
||||
virtual bool getMemoryShapes(const std::vector<MatShape> &inputs,
|
||||
const int requiredOutputs,
|
||||
std::vector<MatShape> &outputs,
|
||||
std::vector<MatShape> &internals) const CV_OVERRIDE
|
||||
{
|
||||
CV_CheckEQ(inputs.size(), 3ull, "Scatter: require three inputs.");
|
||||
CV_CheckEQ(inputs[0].size(), inputs[1].size(), "Scatter: input data should have the same ndim with indices.");
|
||||
CV_CheckEQ(inputs[0].size(), inputs[2].size(), "Scatter: input data should have the same ndim with updates.");
|
||||
for (size_t i = 0; i < inputs[0].size(); i++)
|
||||
{
|
||||
CV_CheckGE(inputs[0][i], inputs[1][i], "Scatter: each dim of input data should be greater than (or equal to) indices'.");
|
||||
CV_CheckEQ(inputs[1][i], inputs[2][i], "Scatter: each dim of indices should be equal to updates'.");
|
||||
}
|
||||
outputs.assign(1, inputs[0]);
|
||||
return false;
|
||||
}
|
||||
|
||||
void forward(InputArrayOfArrays inputs_arr, OutputArrayOfArrays outputs_arr, OutputArrayOfArrays internals_arr) CV_OVERRIDE
|
||||
{
|
||||
CV_TRACE_FUNCTION();
|
||||
CV_TRACE_ARG_VALUE(name, "name", name.c_str());
|
||||
|
||||
std::vector<Mat> inputs, outputs;
|
||||
inputs_arr.getMatVector(inputs);
|
||||
outputs_arr.getMatVector(outputs);
|
||||
|
||||
const Mat& data = inputs[0];
|
||||
const Mat& indices = inputs[1];
|
||||
const Mat& updates = inputs[2];
|
||||
Mat& out = outputs[0];
|
||||
|
||||
typeDispatch(outputs[0].type(), data, indices, updates, out);
|
||||
}
|
||||
|
||||
template<typename T, typename Functor>
|
||||
void forward_impl(const Functor& rd, const Mat& data, const Mat& indices, const Mat& updates, Mat& out)
|
||||
{
|
||||
data.copyTo(out);
|
||||
|
||||
const int ndims = data.dims;
|
||||
const int* shape = data.size.p;
|
||||
const size_t* step = data.step.p;
|
||||
|
||||
const int* ind_shape = indices.size.p;
|
||||
const size_t* ind_step = indices.step.p;
|
||||
|
||||
size_t inp_offset = 0;
|
||||
size_t ind_offset = 0;
|
||||
const T* p_index = indices.ptr<const T>();
|
||||
const T* p_update = updates.ptr<const T>();
|
||||
T* p_out = out.ptr<T>();
|
||||
|
||||
size_t total = indices.total();
|
||||
|
||||
int j, offset_at_idx, index;
|
||||
size_t t, idx;
|
||||
for (size_t i = 0; i < total; i++)
|
||||
{
|
||||
t = i;
|
||||
inp_offset = 0;
|
||||
ind_offset = 0;
|
||||
int offset_at_axis = 0;
|
||||
for (j = ndims - 1; j >= 0; j--)
|
||||
{
|
||||
idx = t / ind_shape[j];
|
||||
offset_at_idx = (int)(t - idx * ind_shape[j]);
|
||||
ind_offset += offset_at_idx * ind_step[j];
|
||||
inp_offset += offset_at_idx * step[j];
|
||||
t = idx;
|
||||
if (j == axis)
|
||||
{
|
||||
offset_at_axis = offset_at_idx * step[j];
|
||||
}
|
||||
}
|
||||
ind_offset /= sizeof(T);
|
||||
|
||||
// get index and overwrite current indices
|
||||
const T* tmp_p_index = p_index + ind_offset;
|
||||
index = (int)(*tmp_p_index);
|
||||
CV_Assert(index < shape[axis] && index > -shape[axis]);
|
||||
|
||||
inp_offset = inp_offset - offset_at_axis + ((index + shape[axis]) % shape[axis]) * step[axis];
|
||||
inp_offset /= sizeof(T);
|
||||
|
||||
const T* tmp_p_update = p_update + ind_offset;
|
||||
T* tmp_p_out = p_out + inp_offset;
|
||||
*tmp_p_out = rd(*tmp_p_out, *tmp_p_update);
|
||||
}
|
||||
}
|
||||
|
||||
template<typename... Args>
|
||||
inline void typeDispatch(const int type, Args&&... args)
|
||||
{
|
||||
switch (type)
|
||||
{
|
||||
case CV_8U:
|
||||
reductionDispatch<uint8_t>(std::forward<Args>(args)...);
|
||||
break;
|
||||
case CV_32S:
|
||||
reductionDispatch<int32_t>(std::forward<Args>(args)...);
|
||||
break;
|
||||
case CV_32F:
|
||||
reductionDispatch<float>(std::forward<Args>(args)...);
|
||||
break;
|
||||
default:
|
||||
CV_Error(cv::Error::BadDepth, "Unsupported type.");
|
||||
};
|
||||
}
|
||||
|
||||
template<typename T, typename... Args>
|
||||
inline void reductionDispatch(Args&&... args)
|
||||
{
|
||||
switch (reduction)
|
||||
{
|
||||
case REDUCTION::NONE:
|
||||
{
|
||||
auto rd = [](const T& a, const T& b) { return b; }; // a from input data, b from updates
|
||||
forward_impl<T>(rd, std::forward<Args>(args)...);
|
||||
break;
|
||||
}
|
||||
case REDUCTION::ADD:
|
||||
{
|
||||
auto rd = [](const T& a, const T& b) { return a + b; };
|
||||
forward_impl<T>(rd, std::forward<Args>(args)...);
|
||||
break;
|
||||
}
|
||||
case REDUCTION::MUL:
|
||||
{
|
||||
auto rd = [](const T& a, const T& b) { return a * b; };
|
||||
forward_impl<T>(rd, std::forward<Args>(args)...);
|
||||
break;
|
||||
}
|
||||
case REDUCTION::MAX:
|
||||
{
|
||||
auto rd = [](const T& a, const T& b) { return std::max(a, b); };
|
||||
forward_impl<T>(rd, std::forward<Args>(args)...);
|
||||
break;
|
||||
}
|
||||
case REDUCTION::MIN:
|
||||
{
|
||||
auto rd = [](const T& a, const T& b) { return std::min(a, b); };
|
||||
forward_impl<T>(rd, std::forward<Args>(args)...);
|
||||
break;
|
||||
}
|
||||
default:
|
||||
CV_Error(Error::StsBadArg, "Unsupported reduction.");
|
||||
};
|
||||
}
|
||||
|
||||
private:
|
||||
// Attributes
|
||||
int axis;
|
||||
};
|
||||
|
||||
Ptr<ScatterLayer> ScatterLayer::create(const LayerParams& params)
|
||||
{
|
||||
return makePtr<ScatterLayerImpl>(params);
|
||||
}
|
||||
|
||||
}} // namespace cv::dnn
|
||||
@@ -188,6 +188,7 @@ private:
|
||||
void parseElementWise (LayerParams& layerParams, const opencv_onnx::NodeProto& node_proto);
|
||||
void parseDepthToSpace (LayerParams& layerParams, const opencv_onnx::NodeProto& node_proto);
|
||||
void parseRange (LayerParams& layerParams, const opencv_onnx::NodeProto& node_proto);
|
||||
void parseScatter (LayerParams& layerParams, const opencv_onnx::NodeProto& node_proto);
|
||||
void parseSimpleLayers (LayerParams& layerParams, const opencv_onnx::NodeProto& node_proto);
|
||||
|
||||
// Domain: com.microsoft
|
||||
@@ -3131,6 +3132,58 @@ void ONNXImporter::parseRange(LayerParams& layerParams, const opencv_onnx::NodeP
|
||||
constBlobsExtraInfo.insert(std::make_pair(node_proto.output(0), TensorInfo(1)));
|
||||
}
|
||||
|
||||
void ONNXImporter::parseScatter(LayerParams& layerParams, const opencv_onnx::NodeProto& node_proto)
|
||||
{
|
||||
CV_CheckEQ(node_proto.input_size(), 3, "Scatter: three inputs are required.");
|
||||
layerParams.type = "Scatter";
|
||||
if (node_proto.op_type() == "ScatterND")
|
||||
layerParams.type = "ScatterND";
|
||||
|
||||
size_t consts = 0;
|
||||
for (size_t i = 0; i < node_proto.input_size(); ++i)
|
||||
if (layer_id.find(node_proto.input(i)) == layer_id.end())
|
||||
++consts;
|
||||
|
||||
if (consts == node_proto.input_size())
|
||||
{
|
||||
std::vector<Mat> inputs, output;
|
||||
for (size_t i = 0; i < node_proto.input_size(); i++)
|
||||
{
|
||||
Mat blob = getBlob(node_proto, i);
|
||||
if (i == 1) // indices
|
||||
blob.convertTo(blob, CV_32F);
|
||||
inputs.push_back(blob);
|
||||
}
|
||||
runLayer(layerParams, inputs, output);
|
||||
CV_Assert(output.size() == 1);
|
||||
addConstant(node_proto.output(0), output[0]);
|
||||
return;
|
||||
}
|
||||
else if (consts > 0)
|
||||
{
|
||||
for (size_t i = 0; i < node_proto.input_size(); i++)
|
||||
{
|
||||
if (layer_id.find(node_proto.input(i)) == layer_id.end())
|
||||
{
|
||||
Mat blob = getBlob(node_proto, i);
|
||||
if (i == 1) // indices, from int32/int64 to float32
|
||||
blob.convertTo(blob, CV_32F);
|
||||
|
||||
LayerParams constParams;
|
||||
constParams.name = node_proto.input(i);
|
||||
constParams.type = "Const";
|
||||
constParams.blobs.push_back(blob);
|
||||
|
||||
opencv_onnx::NodeProto proto;
|
||||
proto.add_output(constParams.name);
|
||||
addLayer(constParams, proto);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
addLayer(layerParams, node_proto);
|
||||
}
|
||||
|
||||
void ONNXImporter::parseSimpleLayers(LayerParams& layerParams, const opencv_onnx::NodeProto& node_proto)
|
||||
{
|
||||
bool is_all_input_const = true;
|
||||
@@ -3785,6 +3838,7 @@ void ONNXImporter::buildDispatchMap_ONNX_AI(int opset_version)
|
||||
dispatch["DetectionOutput"] = &ONNXImporter::parseDetectionOutput;
|
||||
dispatch["CumSum"] = &ONNXImporter::parseCumSum;
|
||||
dispatch["SpaceToDepth"] = dispatch["DepthToSpace"] = &ONNXImporter::parseDepthToSpace;
|
||||
dispatch["ScatterElements"] = dispatch["Scatter"] = dispatch["ScatterND"] = &ONNXImporter::parseScatter;
|
||||
|
||||
dispatch["Equal"] = dispatch["Greater"] = dispatch["Less"] = dispatch["Pow"] = dispatch["Add"] =
|
||||
dispatch["Sub"] = dispatch["Mul"] = dispatch["Div"] = dispatch["GreaterOrEqual"] =
|
||||
|
||||
Reference in New Issue
Block a user