Support NCHW in forward and backward convolution op.

Test both layouts in tests.
Change: 115096872

1 files changed, 32 insertions, 0 deletions

diff --git a/tensorflow/core/util/tensor_format.cc b/tensorflow/core/util/tensor_format.cc
new file mode 100644
index 0000000000..9e29775628
--- /dev/null
+++ b/tensorflow/core/util/tensor_format.cc
@@ -0,0 +1,32 @@
+#include "tensorflow/core/util/tensor_format.h"
+
+namespace tensorflow {
+
+string GetConvnetDataFormatAttrString() {
+  return "data_format: { 'NHWC', 'NCHW' } = 'NHWC' ";
+}
+
+string ToString(TensorFormat format) {
+  switch (format) {
+    case FORMAT_NHWC:
+      return "NHWC";
+    case FORMAT_NCHW:
+      return "NCHW";
+    default:
+      LOG(FATAL) << "Invalid Format: " << static_cast<int32>(format);
+      return "INVALID_FORMAT";
+  }
+}
+
+bool FormatFromString(const string& format_str, TensorFormat* format) {
+  if (format_str == "NHWC") {
+    *format = FORMAT_NHWC;
+    return true;
+  } else if (format_str == "NCHW") {
+    *format = FORMAT_NCHW;
+    return true;
+  }
+  return false;
+}
+
+}  // namespace tensorflow