[GitHub] sandeep-krishnamurthy commented on a change in pull request #13802: Image normalize operator - GPU support, 3D/4D inputs

[GitBox]

sandeep-krishnamurthy commented on a change in pull request #13802: Image 
normalize operator - GPU support, 3D/4D inputs
URL: https://github.com/apache/incubator-mxnet/pull/13802#discussion_r251134919
 
 

 ##########
 File path: src/operator/image/image_random-inl.h
 ##########
 @@ -123,28 +146,157 @@ inline bool NormalizeShape(const nnvm::NodeAttrs& attrs,
   return true;
 }
 
-void Normalize(const nnvm::NodeAttrs &attrs,
+// Type Inference
+inline bool NormalizeOpType(const nnvm::NodeAttrs& attrs,
+                          std::vector<int>* in_attrs,
+                          std::vector<int>* out_attrs) {
+  CHECK_EQ(in_attrs->size(), 1U);
+  CHECK_EQ(out_attrs->size(), 1U);
+
+  // Normalized Tensor will be a float
+  TYPE_ASSIGN_CHECK(*out_attrs, 0, mshadow::kFloat32);
+  return out_attrs->at(0) != -1;
+}
+
+template<int req>
+struct normalize_forward {
+    template<typename DType>
+    MSHADOW_XINLINE static void Map(int j, DType* out_data, const DType* 
in_data,
+                                    const int i, const int length, const int 
step,
+                                    const DType mean, const DType std_dev) {
+        KERNEL_ASSIGN(out_data[step + i*length + j], req,
+                      (in_data[step + i*length + j] - mean) / std_dev);
+    }
+};
+
+template<typename xpu>
+void NormalizeImpl(const OpContext &ctx,
+                          const std::vector<TBlob> &inputs,
+                          const std::vector<TBlob> &outputs,
+                          const std::vector<OpReqType> &req,
+                          const NormalizeParam &param,
+                          const int length,
+                          const int channel,
+                          const int step = 0) {
+    mshadow::Stream<xpu> *s = ctx.get_stream<xpu>();
+
+    MSHADOW_TYPE_SWITCH(outputs[0].type_flag_, DType, {
+      MXNET_ASSIGN_REQ_SWITCH(req[0], req_type, {
+        DType* input = inputs[0].dptr<DType>();
+        DType* output = outputs[0].dptr<DType>();
+
+        for (int i = 0; i < channel; ++i) {
+            DType mean = param.mean[param.mean.ndim() > 1 ? i : 0];
+            DType std_dev = param.std[param.std.ndim() > 1 ? i : 0];
+            mxnet_op::Kernel<normalize_forward<req_type>, xpu>::Launch(
+                s, length, output, input,
+                i, length, step, mean, std_dev);
+        }
+      });
+    });
+}
+
+template<typename xpu>
+void NormalizeOpForward(const nnvm::NodeAttrs &attrs,
                       const OpContext &ctx,
                       const std::vector<TBlob> &inputs,
                       const std::vector<OpReqType> &req,
                       const std::vector<TBlob> &outputs) {
+  CHECK_EQ(inputs.size(), 1U);
+  CHECK_EQ(outputs.size(), 1U);
+  CHECK_EQ(req.size(), 1U);
+
   const NormalizeParam &param = nnvm::get<NormalizeParam>(attrs.parsed);
 
-  int nchannels = inputs[0].shape_[0];
-  int length = inputs[0].shape_[1] * inputs[0].shape_[2];
+  // 3D input (c, h, w)
+  if (inputs[0].ndim() == 3) {
+    const int length = inputs[0].shape_[1] * inputs[0].shape_[2];
+    const int channel = inputs[0].shape_[0];
+    NormalizeImpl<xpu>(ctx, inputs, outputs, req, param, length, channel);
+  } else if (inputs[0].ndim() == 4) {
+    // 4D input (n, c, h, w)
+    const int batch_size = inputs[0].shape_[0];
+    const int length = inputs[0].shape_[2] * inputs[0].shape_[3];
+    const int channel = inputs[0].shape_[1];
+    const int step = channel * length;
+
+    #pragma omp parallel for
+    for (auto n = 0; n < batch_size; ++n) {
+      NormalizeImpl<xpu>(ctx, inputs, outputs, req, param, length, channel, 
n*step);
+    }
+  }
+}
 
-  MSHADOW_TYPE_SWITCH(outputs[0].type_flag_, DType, {
-    DType* input = inputs[0].dptr<DType>();
-    DType* output = outputs[0].dptr<DType>();
+// Backward function
+template<int req>
+struct normalize_backward {
+  template<typename DType>
+  MSHADOW_XINLINE static void Map(int j, DType* in_grad, const DType* out_grad,
+                                  const DType* in_data, const int i, const int 
length,
+                                  const int step, const DType std_dev) {
+    // d/dx{(x - mean) / std_dev} => (1 / std_dev)
+    KERNEL_ASSIGN(in_grad[step + i*length + j], req,
+                  out_grad[step + i*length + j] * (1.0 / std_dev));
+  }
+};
 
-    for (int i = 0; i < nchannels; ++i) {
-      DType mean = param.mean[param.mean.ndim() > 1 ? i : 0];
-      DType std = param.std[param.std.ndim() > 1 ? i : 0];
-      for (int j = 0; j < length; ++j) {
-        output[i*length + j] = (input[i*length + j] - mean) / std;
-      }
+template<typename xpu>
+void NormalizeBackwardImpl(const OpContext &ctx,
+                           const std::vector<TBlob> &inputs,
+                           const std::vector<TBlob> &outputs,
+                           const std::vector<OpReqType> &req,
+                           const NormalizeParam &param,
+                           const int length,
+                           const int channel,
+                           const int step = 0) {
+    mshadow::Stream<xpu> *s = ctx.get_stream<xpu>();
+    const TBlob& out_grad = inputs[0];
+    const TBlob& in_data = inputs[1];
+    const TBlob& in_grad = outputs[0];
+    MSHADOW_TYPE_SWITCH(outputs[0].type_flag_, DType, {
+      MXNET_ASSIGN_REQ_SWITCH(req[0], req_type, {
+        for (int i = 0; i < channel; ++i) {
+            DType std_dev = param.std[param.std.ndim() > 1 ? i : 0];
 
 Review comment:
   same as above

----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on GitHub and use the
URL above to go to the specific comment.
 
For queries about this service, please contact Infrastructure at:
us...@infra.apache.org


With regards,
Apache Git Services