[nnc] Fix strided slice shrink axis mask (#2725)
authorПавел Ильютченко/AI Tools Lab /SRR/Engineer/삼성전자 <p.iliutchenk@samsung.com>
Wed, 19 Dec 2018 12:17:20 +0000 (15:17 +0300)
committerEfimov Alexander/AI Tools Lab/./Samsung Electronics <a.efimov@samsung.com>
Wed, 19 Dec 2018 12:17:20 +0000 (15:17 +0300)
* Fix processing shrink axis mask on StridedSlice operation creation on TFLite importer

Signed-off-by: Pavel Iliutchenko <p.iliutchenk@samsung.com>
contrib/nnc/passes/tflite_frontend/tflite_op_creator.cpp

index 66c0943..02e2e48 100644 (file)
@@ -480,21 +480,24 @@ TFLiteOpCreator::createStridedSlice(InputOps& inputs, const InputParams& params,
 
   Shape start(num_dims);
   Shape size(num_dims);
+  std::vector<int32_t> squeeze_dims;
   for (int axis = 0; axis < num_dims; axis++) {
     if (begin_mask & (1 << axis))
       start.dim(axis) = 0;
     else
       start.dim(axis) = begin.at(Index{axis});
+
     if (end_mask & (1 << axis))
       size.dim(axis) = inputs[0]->getOutputShape(0).dim(axis) - start.dim(axis);
     else
       size.dim(axis) = end.at(Index{axis}) - start.dim(axis);
+    
+    if (shrink_axis_mask & (1 << axis))
+      squeeze_dims.push_back(axis);
   }
 
-  std::vector<int32_t> squeeze_dims{shrink_axis_mask - 1};
-
   auto slice_outputs = createOp<ops::SliceOp>(ActivationFunctionType_NONE,
-                                               inputs[0]->getOutput(0), start, size);
+                                              inputs[0]->getOutput(0), start, size);
   return createOp<ops::SqueezeOp>(ActivationFunctionType_NONE,
                                   slice_outputs[0]->getOutput(0), squeeze_dims);
 }