Fix TF Split layer

This commit is contained in:
Liubov Batanina 2019-07-17 15:50:50 +03:00
parent e4e0bb533d
commit 0d2bc7b5fd
4 changed files with 14 additions and 3 deletions

View File

@ -366,6 +366,7 @@ CV__DNN_EXPERIMENTAL_NS_BEGIN
*/
std::vector<std::vector<Range> > sliceRanges;
int axis;
int num_split;
static Ptr<SliceLayer> create(const LayerParams &params);
};

View File

@ -61,6 +61,7 @@ public:
{
setParamsFrom(params);
axis = params.get<int>("axis", 1);
num_split = params.get<int>("num_split", 0);
if (params.has("slice_point"))
{
CV_Assert(!params.has("begin") && !params.has("size") && !params.has("end"));
@ -141,9 +142,10 @@ public:
else // Divide input blob on equal parts by axis.
{
CV_Assert(0 <= axis && axis < inpShape.size());
CV_Assert(requiredOutputs > 0 && inpShape[axis] % requiredOutputs == 0);
inpShape[axis] /= requiredOutputs;
outputs.resize(requiredOutputs, inpShape);
int splits = num_split ? num_split : requiredOutputs;
CV_Assert(splits > 0 && inpShape[axis] % splits == 0);
inpShape[axis] /= splits;
outputs.resize(splits, inpShape);
}
return false;
}

View File

@ -1410,6 +1410,9 @@ void TFImporter::populateNet(Net dstNet)
axis = toNCHW(axis);
layerParams.set("axis", axis);
if (hasLayerAttr(layer, "num_split"))
layerParams.set("num_split", getLayerAttr(layer, "num_split").i());
int id = dstNet.addLayer(name, "Slice", layerParams);
layer_id[name] = id;

View File

@ -350,6 +350,11 @@ TEST_P(Test_TensorFlow_layers, l2_normalize_3d)
runTensorFlowNet("l2_normalize_3d");
}
TEST_P(Test_TensorFlow_layers, Split)
{
runTensorFlowNet("split");
}
class Test_TensorFlow_nets : public DNNTestLayer {};
TEST_P(Test_TensorFlow_nets, MobileNet_SSD)