[OpenCL]fix tf_shufflenet model segmentation bug (#5777)

* fix tf_shufflenet model segmentation bug test=develop * add unsupport cue test=develop
PaddlePaddle · Mar 23, 2021 · 9b0436c · 9b0436c
1 parent 22fca16
commit 9b0436c
Showing 1 changed file with 36 additions and 6 deletions.
diff --git a/lite/kernels/opencl/layout_image_compute.cc b/lite/kernels/opencl/layout_image_compute.cc
@@ -80,8 +80,18 @@ class LayoutComputeBufferChwToImageDefault
 
     // out info
     std::vector<size_t> new_dims = {1, 1, 1, 1};
-    for (int tidx = 0; tidx < x_dims.size(); ++tidx) {
-      new_dims[4 - x_dims.size() + tidx] = x_dims[tidx];
+    if (x_dims.size() == 5) {
+      new_dims[4 - x_dims.size() + 1] = x_dims[0] * x_dims[1];
+      for (int tidx = 2; tidx < x_dims.size(); ++tidx) {
+        new_dims[4 - x_dims.size() + tidx] = x_dims[tidx];
+      }
+    } else if (x_dims.size() < 5) {
+      for (int tidx = 0; tidx < x_dims.size(); ++tidx) {
+        new_dims[4 - x_dims.size() + tidx] = x_dims[tidx];
+      }
+    } else {
+      LOG(FATAL) << "unsupported layout tensor dims size, the dims size is:"
+                 << x_dims.size();
     }
     const int out_C = new_dims[1];
     const int out_H = new_dims[2];
@@ -207,8 +217,18 @@ class LayoutComputeImageDefaultToBufferChw
     auto x_image_shape = InitImageDimInfoWith(x_dims);
 
     std::vector<size_t> new_dims = {1, 1, 1, 1};
-    for (int j = 0; j < x_dims.size(); ++j) {
-      new_dims[4 - x_dims.size() + j] = x_dims[j];
+    if (x_dims.size() == 5) {
+      new_dims[4 - x_dims.size() + 1] = x_dims[0] * x_dims[1];
+      for (int j = 2; j < x_dims.size(); ++j) {
+        new_dims[4 - x_dims.size() + j] = x_dims[j];
+      }
+    } else if (x_dims.size() < 5) {
+      for (int j = 0; j < x_dims.size(); ++j) {
+        new_dims[4 - x_dims.size() + j] = x_dims[j];
+      }
+    } else {
+      LOG(FATAL) << "unsupported layout tensor dims size, the dims size is: "
+                 << x_dims.size();
     }
 
 #ifdef LITE_WITH_LOG
@@ -322,8 +342,18 @@ class LayoutComputeBufferChwToImage2DNw
 
     // out info
     std::vector<size_t> new_dims = {1, 1, 1, 1};
-    for (int tidx = 0; tidx < x_dims.size(); ++tidx) {
-      new_dims[4 - x_dims.size() + tidx] = x_dims[tidx];
+    if (x_dims.size() == 5) {
+      new_dims[4 - x_dims.size() + 1] = x_dims[0] * x_dims[1];
+      for (int tidx = 2; tidx < x_dims.size(); ++tidx) {
+        new_dims[4 - x_dims.size() + tidx] = x_dims[tidx];
+      }
+    } else if (x_dims.size() < 5) {
+      for (int tidx = 0; tidx < x_dims.size(); ++tidx) {
+        new_dims[4 - x_dims.size() + tidx] = x_dims[tidx];
+      }
+    } else {
+      LOG(FATAL) << "unsupported layout tensor dims size, the dims size is:"
+                 << x_dims.size();
     }
 
     const int out_N = new_dims[0];