Refine code.

paddle/operators/gemm_conv2d_op.h

@@ -84,8 +84,8 @@ class GemmConv2DKernel : public framework::OpKernel {
       for (int g = 0; g < groups; g++) {
         // im2col
         Tensor in_slice = in_batch.Slice<T>(g * in_step, (g + 1) * in_step);
-        im2col(in_slice, col, strides[0], strides[1], paddings[0], paddings[1],
-               context.device_context());
+        im2col(context.device_context(), in_slice, col, strides[0], strides[1],
+               paddings[0], paddings[1]);
 
         // gemm
         Tensor out_slice = out_batch.Slice<T>(g * out_step, (g + 1) * out_step);
@@ -185,8 +185,8 @@ class GemmConvGrad2DKernel : public framework::OpKernel {
           // col2im
           Tensor in_grad_slice =
               in_grad_batch.Slice<T>(g * in_step, (g + 1) * in_step);
-          col2im(in_grad_slice, col, strides[0], strides[1], paddings[0],
-                 paddings[1], context.device_context());
+          col2im(context.device_context(), in_grad_slice, col, strides[0],
+                 strides[1], paddings[0], paddings[1]);
         }
       }
     }
@@ -207,8 +207,8 @@ class GemmConvGrad2DKernel : public framework::OpKernel {
           Tensor out_grad_slice =
               out_grad_batch.Slice<T>(g * out_step, (g + 1) * out_step);
           Tensor in_slice = in_batch.Slice<T>(g * in_step, (g + 1) * in_step);
-          im2col(in_slice, col, strides[0], strides[1], paddings[0],
-                 paddings[1], context.device_context());
+          im2col(context.device_context(), in_slice, col, strides[0],
+                 strides[1], paddings[0], paddings[1]);
 
           // gemm
           Tensor filter_grad_slice =

paddle/operators/math/im2col.cc

@@ -27,9 +27,10 @@ template <class T>
 class Im2ColFunctor<paddle::operators::math::ColFormat::kCFO,
                     platform::CPUPlace, T> {
  public:
-  void operator()(const framework::Tensor& im, framework::Tensor& col,
+  void operator()(const platform::DeviceContext& context,
+                  const framework::Tensor& im, framework::Tensor& col,
                   int stride_height, int stride_width, int padding_height,
-                  int padding_width, const platform::DeviceContext& context) {
+                  int padding_width) {
     PADDLE_ENFORCE(im.dims().size() == 3);
     PADDLE_ENFORCE(col.dims().size() == 5);
 
@@ -79,9 +80,9 @@ template <class T>
 class Col2ImFunctor<paddle::operators::math::ColFormat::kCFO,
                     platform::CPUPlace, T> {
  public:
-  void operator()(framework::Tensor& im, const framework::Tensor& col,
-                  int stride_height, int stride_width, int padding_height,
-                  int padding_width, const platform::DeviceContext& context) {
+  void operator()(const platform::DeviceContext& context, framework::Tensor& im,
+                  const framework::Tensor& col, int stride_height,
+                  int stride_width, int padding_height, int padding_width) {
     PADDLE_ENFORCE(im.dims().size() == 3);
     PADDLE_ENFORCE(col.dims().size() == 5);
     int input_channels = im.dims()[0];
@@ -137,9 +138,10 @@ template <class T>
 class Im2ColFunctor<paddle::operators::math::ColFormat::kOCF,
                     platform::CPUPlace, T> {
  public:
-  void operator()(const framework::Tensor& im, framework::Tensor& col,
+  void operator()(const platform::DeviceContext& context,
+                  const framework::Tensor& im, framework::Tensor& col,
                   int stride_height, int stride_width, int padding_height,
-                  int padding_width, const platform::DeviceContext& context) {
+                  int padding_width) {
     PADDLE_ENFORCE(im.dims().size() == 3);
     PADDLE_ENFORCE(col.dims().size() == 5);
     int input_channels = im.dims()[0];
@@ -197,9 +199,9 @@ template <class T>
 class Col2ImFunctor<paddle::operators::math::ColFormat::kOCF,
                     platform::CPUPlace, T> {
  public:
-  void operator()(framework::Tensor& im, const framework::Tensor& col,
-                  int stride_height, int stride_width, int padding_height,
-                  int padding_width, const platform::DeviceContext& context) {
+  void operator()(const platform::DeviceContext& context, framework::Tensor& im,
+                  const framework::Tensor& col, int stride_height,
+                  int stride_width, int padding_height, int padding_width) {
     PADDLE_ENFORCE(im.dims().size() == 3);
     PADDLE_ENFORCE(col.dims().size() == 5);
     int input_channels = im.dims()[0];

paddle/operators/math/im2col.cu

@@ -64,9 +64,10 @@ template <class T>
 class Im2ColFunctor<paddle::operators::math::ColFormat::kCFO,
                     platform::GPUPlace, T> {
  public:
-  void operator()(const framework::Tensor& im, framework::Tensor& col,
+  void operator()(const platform::DeviceContext& context,
+                  const framework::Tensor& im, framework::Tensor& col,
                   int stride_height, int stride_width, int padding_height,
-                  int padding_width, const platform::DeviceContext& context) {
+                  int padding_width) {
     PADDLE_ENFORCE(im.dims().size() == 3);
     PADDLE_ENFORCE(col.dims().size() == 5);
 
@@ -149,9 +150,9 @@ template <class T>
 class Col2ImFunctor<paddle::operators::math::ColFormat::kCFO,
                     platform::GPUPlace, T> {
  public:
-  void operator()(framework::Tensor& im, const framework::Tensor& col,
-                  int stride_height, int stride_width, int padding_height,
-                  int padding_width, const platform::DeviceContext& context) {
+  void operator()(const platform::DeviceContext& context, framework::Tensor& im,
+                  const framework::Tensor& col, int stride_height,
+                  int stride_width, int padding_height, int padding_width) {
     PADDLE_ENFORCE(im.dims().size() == 3);
     PADDLE_ENFORCE(col.dims().size() == 5);
 
@@ -235,9 +236,10 @@ template <class T>
 class Im2ColFunctor<paddle::operators::math::ColFormat::kOCF,
                     platform::GPUPlace, T> {
  public:
-  void operator()(const framework::Tensor& im, framework::Tensor& col,
+  void operator()(const platform::DeviceContext& context,
+                  const framework::Tensor& im, framework::Tensor& col,
                   int stride_height, int stride_width, int padding_height,
-                  int padding_width, const platform::DeviceContext& context) {
+                  int padding_width) {
     PADDLE_ENFORCE(im.dims().size() == 3);
     PADDLE_ENFORCE(col.dims().size() == 5);
     int input_channels = im.dims()[0];
@@ -318,9 +320,9 @@ template <class T>
 class Col2ImFunctor<paddle::operators::math::ColFormat::kOCF,
                     platform::GPUPlace, T> {
  public:
-  void operator()(framework::Tensor& im, const framework::Tensor& col,
-                  int stride_height, int stride_width, int padding_height,
-                  int padding_width, const platform::DeviceContext& context) {
+  void operator()(const platform::DeviceContext& context, framework::Tensor& im,
+                  const framework::Tensor& col, int stride_height,
+                  int stride_width, int padding_height, int padding_width) {
     PADDLE_ENFORCE(im.dims().size() == 3);
     PADDLE_ENFORCE(col.dims().size() == 5);
     int input_channels = im.dims()[0];

paddle/operators/math/im2col.h

@@ -72,17 +72,18 @@ enum class ColFormat { kCFO = 0, kOCF = 1 };
 template <ColFormat Format, typename Place, typename T>
 class Im2ColFunctor {
  public:
-  void operator()(const framework::Tensor& im, framework::Tensor& col,
+  void operator()(const platform::DeviceContext& context,
+                  const framework::Tensor& im, framework::Tensor& col,
                   int stride_height, int stride_width, int padding_height,
-                  int padding_width, const platform::DeviceContext& context);
+                  int padding_width);
 };
 
 template <ColFormat Format, typename Place, typename T>
 class Col2ImFunctor {
  public:
-  void operator()(framework::Tensor& im, const framework::Tensor& col,
-                  int stride_height, int stride_width, int padding_height,
-                  int padding_width, const platform::DeviceContext& context);
+  void operator()(const platform::DeviceContext& context, framework::Tensor& im,
+                  const framework::Tensor& col, int stride_height,
+                  int stride_width, int padding_height, int padding_width);
 };
 
 }  // namespace math

paddle/operators/math/im2col_test.cc

@@ -78,8 +78,8 @@ void testIm2col() {
     PADDLE_THROW("no GPU support");
 #endif  // PADDLE_ONLY_CPU
   }
-  im2col(input, output_cfo, stride, stride, padding, padding, *context);
-  im2col_ocf(input, output_ocf, stride, stride, padding, padding, *context);
+  im2col(*context, input, output_cfo, stride, stride, padding, padding);
+  im2col_ocf(*context, input, output_ocf, stride, stride, padding, padding);
 
   float* out_cfo_ptr;
   if (paddle::platform::is_cpu_place(*place)) {