add a few comments on layout after im2col (#12429)

Summary: Pull Request resolved: #12429 Comments to clarify layout after NHWC im2col for group convolution. i-am-not-moving-c2-to-c10 Reviewed By: houseroad Differential Revision: D10233284 fbshipit-source-id: 996a69f2f932e02c978abaade7571b00741b6ae8
pytorch · Nov 4, 2018 · f000101 · f000101
1 parent 6b578cd
commit f000101
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 0 deletions.
diff --git a/caffe2/operators/conv_op_impl.h b/caffe2/operators/conv_op_impl.h
@@ -264,6 +264,8 @@ bool ConvOp<T, Context>::RunOnDeviceWithOrderNHWC() {
           group_);
       // Weight term
       for (int group_id = 0; group_id < group_; ++group_id) {
+        // col_buffer_data in G (H W) (R S C/G) layout
+        // filter_data in G K/G (R S C/G) layout
         math::GemmEx<T, Context>(
             CblasNoTrans,
             CblasTrans,

diff --git a/caffe2/utils/math_cpu.cc b/caffe2/utils/math_cpu.cc
@@ -3339,6 +3339,8 @@ C10_EXPORT void Im2Col<float, CPUContext, StorageOrder::NHWC>(
       h_pad += stride_h;
     } // h
   } else {
+    // img_data in N H W G C/G layout
+    // col_data in N G H W R S C/G layout
     const int C_per_G = C / groups;
     for (int h = 0; h < output_h; ++h) {
       int w_pad = -pad_l;