conv_op.h 16.9 KB
Newer Older
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16
/* Copyright (c) 2016 PaddlePaddle Authors. All Rights Reserve.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License. */

#pragma once

H
hedaoyuan 已提交
17
#include "paddle/framework/eigen.h"
18
#include "paddle/framework/op_registry.h"
Z
zlx 已提交
19
#include "paddle/operators/math/depthwise_conv.h"
20 21
#include "paddle/operators/math/im2col.h"
#include "paddle/operators/math/math_function.h"
C
chengduoZH 已提交
22
#include "paddle/operators/math/vol2col.h"
23 24 25 26 27 28

namespace paddle {
namespace operators {

using Tensor = framework::Tensor;

武毅 已提交
29 30
// Base convolution operator definations for other conv
// like operators to reuse the implementation.
C
chengduoZH 已提交
31
inline int OutputSize(int input_size, int filter_size, int dilation,
C
chengduoZH 已提交
32 33 34
                      int padding, int stride) {
  const int dkernel = dilation * (filter_size - 1) + 1;
  const int output_size = (input_size + 2 * padding - dkernel) / stride + 1;
武毅 已提交
35 36
  return output_size;
}
C
chengduoZH 已提交
37 38 39
inline bool IsExpand(std::vector<int64_t>& filter_dim,
                     std::vector<int>& strides, std::vector<int>& paddings,
                     std::vector<int>& dilations) {
C
chengduoZH 已提交
40 41
  bool filter_1 = true, strides_1 = true, padding_0 = true, dilation_1 = true;
  for (size_t j = 0; j < strides.size(); ++j) {
C
chengduoZH 已提交
42
    filter_1 = filter_1 && (static_cast<int>(filter_dim[j + 2]) == 1);
C
chengduoZH 已提交
43 44 45
    strides_1 = strides_1 && (strides[j] == 1);
    padding_0 = padding_0 && (paddings[j] == 0);
    dilation_1 = dilation_1 && (dilations[j] == 1);
C
chengduoZH 已提交
46
  }
C
chengduoZH 已提交
47
  return !(filter_1 && strides_1 && padding_0 && dilation_1);
C
chengduoZH 已提交
48
}
武毅 已提交
49 50 51 52 53

// Define Op classes in .h file so that other conv
// operator implementations can reuse the code.
class Conv2DOpMaker : public framework::OpProtoAndCheckerMaker {
 public:
54
  Conv2DOpMaker(OpProto* proto, OpAttrChecker* op_checker);
武毅 已提交
55 56
};

C
chengduoZH 已提交
57 58
class Conv3DOpMaker : public framework::OpProtoAndCheckerMaker {
 public:
59
  Conv3DOpMaker(OpProto* proto, OpAttrChecker* op_checker);
C
chengduoZH 已提交
60 61 62
};

class ConvOp : public framework::OperatorWithKernel {
武毅 已提交
63 64 65
 public:
  using framework::OperatorWithKernel::OperatorWithKernel;
  void InferShape(framework::InferShapeContext* ctx) const override;
66 67 68 69

 protected:
  framework::OpKernelType GetExpectedKernelType(
      const framework::ExecutionContext& ctx) const override;
武毅 已提交
70 71
};

C
chengduoZH 已提交
72
class ConvOpGrad : public framework::OperatorWithKernel {
武毅 已提交
73 74 75
 public:
  using framework::OperatorWithKernel::OperatorWithKernel;
  void InferShape(framework::InferShapeContext* ctx) const override;
76 77 78 79

 protected:
  framework::OpKernelType GetExpectedKernelType(
      const framework::ExecutionContext& ctx) const override;
武毅 已提交
80 81
};

Q
QI JUN 已提交
82
template <typename DeviceContext, typename T>
C
chengduoZH 已提交
83
class GemmConvKernel : public framework::OpKernel<T> {
84 85 86
 public:
  void Compute(const framework::ExecutionContext& context) const override {
    const Tensor* input = context.Input<Tensor>("Input");
H
hedaoyuan 已提交
87 88 89 90
    // The filter will be reshaped in the calculations,
    // so here use an assignment operation,
    // that avoids modifying the variable in the Scope.
    Tensor filter = *context.Input<Tensor>("Filter");
91 92 93
    Tensor* output = context.Output<Tensor>("Output");
    output->mutable_data<T>(context.GetPlace());

C
chengduoZH 已提交
94
    int groups = context.Attr<int>("groups");
95 96
    std::vector<int> strides = context.Attr<std::vector<int>>("strides");
    std::vector<int> paddings = context.Attr<std::vector<int>>("paddings");
C
chengduoZH 已提交
97
    std::vector<int> dilations = context.Attr<std::vector<int>>("dilations");
98

C
chengduoZH 已提交
99 100
    const int batch_size = static_cast<int>(input->dims()[0]);

C
chengduoZH 已提交
101
    // filter_shape_vec: {k_o, k_i, k_h, k_w} or {k_o, k_i, k_d, k_h, k_w}
C
chengduoZH 已提交
102
    std::vector<int64_t> filter_shape_vec(framework::vectorize(filter.dims()));
C
chengduoZH 已提交
103
    // output_shape_vec: {o_n, o_c, o_h, o_w} or {o_n, o_c, o_d, o_h, o_w}
C
chengduoZH 已提交
104
    std::vector<int64_t> output_shape_vec(framework::vectorize(output->dims()));
105

H
hedaoyuan 已提交
106
    // use col_shape in the im2col calculation
C
chengduoZH 已提交
107 108
    // col_shape_vec: {i_c/g, k_h, k_w, o_h, o_w} or {i_c/g, k_d, k_h, k_w, o_d,
    // o_h, o_w}
C
chengduoZH 已提交
109 110 111 112 113 114 115
    size_t data_dim = filter_shape_vec.size() - 2;
    std::vector<int64_t> col_shape_vec(1 + 2 * data_dim);
    col_shape_vec[0] = input->dims()[1] / groups;
    for (size_t j = 0; j < data_dim; ++j) {
      col_shape_vec[j + 1] = filter_shape_vec[j + 2];
      col_shape_vec[j + 1 + data_dim] = output_shape_vec[j + 2];
    }
C
chengduoZH 已提交
116 117
    framework::DDim col_shape(framework::make_ddim(col_shape_vec));

H
hedaoyuan 已提交
118
    // use col_matrix_shape in the gemm calculation
C
chengduoZH 已提交
119 120 121
    // size: (i_c/g * k_h * k_w, o_h * o_w) or (i_c/g * k_d * k_h * k_w, o_d *
    // o_h * o_w)
    framework::DDim col_matrix_shape =
C
chengduoZH 已提交
122
        framework::flatten_to_2d(col_shape, data_dim + 1);
C
chengduoZH 已提交
123

C
chengduoZH 已提交
124
    bool is_expand = IsExpand(filter_shape_vec, strides, paddings, dilations);
H
hedaoyuan 已提交
125
    Tensor col;
H
hedaoyuan 已提交
126 127 128
    // col_matrix shares the same piece of data with col,
    // but will be reshaped into a two-dimensional matrix shape
    // to call the matrix multiplication interface.
C
chengduoZH 已提交
129
    Tensor col_matrix;
C
chengduoZH 已提交
130
    if (is_expand) {
C
chengduoZH 已提交
131 132 133 134
      col.mutable_data<T>(col_shape, context.GetPlace());
      col_matrix.ShareDataWith(col);
      col_matrix.Resize(col_matrix_shape);
    }
135

C
chengduoZH 已提交
136 137 138
    framework::DDim input_shape = framework::slice_ddim(
        input->dims(), 1, static_cast<int>(input->dims().size()));

H
hedaoyuan 已提交
139 140
    framework::DDim filter_matrix_shape = {filter.dims()[0],
                                           filter.numel() / filter.dims()[0]};
H
hedaoyuan 已提交
141 142
    filter.Resize(filter_matrix_shape);

C
chengduoZH 已提交
143 144 145 146 147 148 149 150
    framework::DDim output_matrix_shape = {
        output->dims()[1],
        output->numel() / (output->dims()[0] * output->dims()[1])};

    // convolution operator: im2col(or vol2col) + gemm
    int in_step = static_cast<int>(input->dims()[1]) / groups;
    int out_step = static_cast<int>(output->dims()[1]) / groups;

Q
QI JUN 已提交
151 152
    math::Vol2ColFunctor<DeviceContext, T> vol2col;
    math::Im2ColFunctor<math::ColFormat::kCFO, DeviceContext, T> im2col;
C
chengduoZH 已提交
153

Q
QI JUN 已提交
154
    auto& dev_ctx = context.template device_context<DeviceContext>();
C
chengduoZH 已提交
155 156 157
    for (int i = 0; i < batch_size; i++) {
      Tensor in_batch = input->Slice(i, i + 1).Resize(input_shape);
      Tensor out_batch = output->Slice(i, i + 1).Resize(output_matrix_shape);
C
chengduoZH 已提交
158

C
chengduoZH 已提交
159 160
      for (int g = 0; g < groups; g++) {
        Tensor in_slice = in_batch.Slice(g * in_step, (g + 1) * in_step);
H
hedaoyuan 已提交
161

C
chengduoZH 已提交
162
        if (!is_expand) {
C
chengduoZH 已提交
163 164 165
          col.ShareDataWith(in_slice);
          col_matrix.ShareDataWith(col);
          col_matrix.Resize(col_matrix_shape);
C
chengduoZH 已提交
166
        } else if (data_dim == 2U) {
C
chengduoZH 已提交
167
          // im2col
Q
QI JUN 已提交
168
          im2col(dev_ctx, in_slice, dilations, strides,
C
chengduoZH 已提交
169 170 171
                 std::vector<int>{paddings[0], paddings[1], paddings[0],
                                  paddings[1]},
                 &col);
C
chengduoZH 已提交
172
        } else if (data_dim == 3U) {
C
chengduoZH 已提交
173
          // vol2col
Q
QI JUN 已提交
174
          vol2col(dev_ctx, in_slice, dilations, strides, paddings, &col);
C
chengduoZH 已提交
175
        }
C
chengduoZH 已提交
176 177 178 179

        // gemm
        Tensor out_slice = out_batch.Slice(g * out_step, (g + 1) * out_step);
        Tensor filter_slice = filter.Slice(g * out_step, (g + 1) * out_step);
Q
QI JUN 已提交
180 181
        math::matmul<DeviceContext, T>(dev_ctx, filter_slice, false, col_matrix,
                                       false, T(1.0), &out_slice, T(0.0));
H
hedaoyuan 已提交
182
      }
183 184 185 186
    }
  }
};

Q
QI JUN 已提交
187
template <typename DeviceContext, typename T>
C
chengduoZH 已提交
188
class GemmConvGradKernel : public framework::OpKernel<T> {
189 190
 public:
  void Compute(const framework::ExecutionContext& context) const override {
H
hedaoyuan 已提交
191 192 193 194 195
    const Tensor* input = context.Input<Tensor>("Input");
    const Tensor* output_grad =
        context.Input<Tensor>(framework::GradVarName("Output"));
    Tensor* input_grad =
        context.Output<Tensor>(framework::GradVarName("Input"));
H
hedaoyuan 已提交
196
    Tensor* filter_grad =
H
hedaoyuan 已提交
197
        context.Output<Tensor>(framework::GradVarName("Filter"));
H
hedaoyuan 已提交
198 199 200 201
    // The filter and filter_grad will be reshaped in the calculations,
    // so here use an assignment operation,
    // that avoids modifying the variable in the Scope.
    Tensor filter = *context.Input<Tensor>("Filter");
H
hedaoyuan 已提交
202

C
chengduoZH 已提交
203 204
    if (!input_grad && !filter_grad) return;

C
chengduoZH 已提交
205
    int groups = context.Attr<int>("groups");
H
hedaoyuan 已提交
206 207
    std::vector<int> strides = context.Attr<std::vector<int>>("strides");
    std::vector<int> paddings = context.Attr<std::vector<int>>("paddings");
C
chengduoZH 已提交
208
    std::vector<int> dilations = context.Attr<std::vector<int>>("dilations");
H
hedaoyuan 已提交
209

C
chengduoZH 已提交
210
    const int batch_size = static_cast<int>(input->dims()[0]);
H
hedaoyuan 已提交
211

C
chengduoZH 已提交
212
    // filter_shape_vec: {k_o, k_i, k_h, k_w} or {k_o, k_i, k_d, k_h, k_w}
C
chengduoZH 已提交
213
    std::vector<int64_t> filter_shape_vec(framework::vectorize(filter.dims()));
C
chengduoZH 已提交
214
    // output_shape_vec: {o_n, o_c, o_h, o_w} or {o_n, o_c, o_d, o_h, o_w}
C
chengduoZH 已提交
215 216
    std::vector<int64_t> output_shape_vec(
        framework::vectorize(output_grad->dims()));
C
chengduoZH 已提交
217

C
chengduoZH 已提交
218 219 220
    // use col_shape in the im2col calculation
    // col_shape_vec: {i_c/g, k_h, k_w, o_h, o_w} or {i_c/g, k_d, k_h, k_w, o_d,
    // o_h, o_w}
C
chengduoZH 已提交
221 222 223 224 225 226 227
    size_t data_dim = filter_shape_vec.size() - 2;
    std::vector<int64_t> col_shape_vec(1 + 2 * data_dim);
    col_shape_vec[0] = input->dims()[1] / groups;
    for (size_t j = 0; j < data_dim; ++j) {
      col_shape_vec[j + 1] = filter_shape_vec[j + 2];
      col_shape_vec[j + 1 + data_dim] = output_shape_vec[j + 2];
    }
C
chengduoZH 已提交
228
    framework::DDim col_shape(framework::make_ddim(col_shape_vec));
C
chengduoZH 已提交
229 230

    // use col_matrix_shape in the gemm calculation
C
chengduoZH 已提交
231 232 233 234
    // size: (i_c/g * k_h * k_w, o_h * o_w)
    // or
    // (i_c/g * k_d * k_h * k_w, o_d * o_h * o_w)
    framework::DDim col_matrix_shape =
C
chengduoZH 已提交
235
        framework::flatten_to_2d(col_shape, data_dim + 1);
C
chengduoZH 已提交
236 237 238

    framework::DDim input_shape = framework::slice_ddim(
        input->dims(), 1, static_cast<int>(input->dims().size()));
C
chengduoZH 已提交
239

C
chengduoZH 已提交
240 241
    framework::DDim filter_matrix_shape = {filter.dims()[0],
                                           filter.numel() / filter.dims()[0]};
C
chengduoZH 已提交
242 243 244
    filter.Resize(filter_matrix_shape);

    framework::DDim output_matrix_shape = {
C
chengduoZH 已提交
245 246 247
        output_grad->dims()[1],
        output_grad->numel() /
            (output_grad->dims()[0] * output_grad->dims()[1])};
C
chengduoZH 已提交
248

C
chengduoZH 已提交
249 250 251 252
    // convolution backward input operator:  gemm + col2im(or col2vol)
    // convolution backward weight operator: im2col(or vol2col) + gemm
    int in_step = static_cast<int>(input->dims()[1]) / groups;
    int out_step = static_cast<int>(output_grad->dims()[1]) / groups;
C
chengduoZH 已提交
253

C
chengduoZH 已提交
254
    bool is_expand = IsExpand(filter_shape_vec, strides, paddings, dilations);
C
chengduoZH 已提交
255 256 257 258
    Tensor col;
    // col_matrix shares the same piece of data with col,
    // but will be reshaped into a two-dimensional matrix shape
    // to call the matrix multiplication interface.
C
chengduoZH 已提交
259
    Tensor col_matrix;
C
chengduoZH 已提交
260
    if (is_expand) {
C
chengduoZH 已提交
261 262 263 264
      col.mutable_data<T>(col_shape, context.GetPlace());
      col_matrix.ShareDataWith(col);
      col_matrix.Resize(col_matrix_shape);
    }
C
chengduoZH 已提交
265

Q
QI JUN 已提交
266 267
    math::SetConstant<DeviceContext, T> set_zero;
    auto& dev_ctx = context.template device_context<DeviceContext>();
C
chengduoZH 已提交
268 269 270 271

    if (input_grad) {
      input_grad->mutable_data<T>(context.GetPlace());

C
chengduoZH 已提交
272 273 274
      // if is_expand is false, the operation of set_zero is unnecessary,
      // because math::matmul will reset input_grad.
      if (is_expand) {
C
chengduoZH 已提交
275
        set_zero(dev_ctx, input_grad, static_cast<T>(0));
C
chengduoZH 已提交
276
      }
Q
QI JUN 已提交
277 278
      math::Col2VolFunctor<DeviceContext, T> col2vol;
      math::Col2ImFunctor<math::ColFormat::kCFO, DeviceContext, T> col2im;
C
chengduoZH 已提交
279

C
chengduoZH 已提交
280 281 282 283 284 285 286 287 288 289 290 291 292 293
      for (int i = 0; i < batch_size; i++) {
        Tensor out_grad_batch =
            output_grad->Slice(i, i + 1).Resize(output_matrix_shape);
        Tensor in_grad_batch = input_grad->Slice(i, i + 1).Resize(input_shape);
        for (int g = 0; g < groups; g++) {
          // gemm
          Tensor out_grad_slice =
              out_grad_batch.Slice(g * out_step, (g + 1) * out_step);
          Tensor filter_slice = filter.Slice(g * out_step, (g + 1) * out_step);

          Tensor in_grad_slice =
              in_grad_batch.Slice(g * in_step, (g + 1) * in_step);

          if (!is_expand) {
C
chengduoZH 已提交
294 295
            col_matrix.ShareDataWith(in_grad_slice);
            col_matrix.Resize(col_matrix_shape);
C
chengduoZH 已提交
296
          }
Q
QI JUN 已提交
297 298 299
          math::matmul<DeviceContext, T>(dev_ctx, filter_slice, true,
                                         out_grad_slice, false, T(1.0),
                                         &col_matrix, T(0.0));
C
chengduoZH 已提交
300

C
chengduoZH 已提交
301
          if (is_expand && data_dim == 2U) {
Q
QI JUN 已提交
302
            col2im(dev_ctx, col, dilations, strides,
C
chengduoZH 已提交
303 304 305
                   std::vector<int>{paddings[0], paddings[1], paddings[0],
                                    paddings[1]},
                   &in_grad_slice);
C
chengduoZH 已提交
306
          } else if (is_expand && data_dim == 3U) {
Q
QI JUN 已提交
307
            col2vol(dev_ctx, col, dilations, strides, paddings, &in_grad_slice);
C
chengduoZH 已提交
308
          }
C
chengduoZH 已提交
309 310 311 312 313 314 315 316
        }
      }
    }

    if (filter_grad) {
      filter_grad->mutable_data<T>(context.GetPlace());
      Tensor filter_grad_ = *filter_grad;
      filter_grad_.Resize(filter_matrix_shape);
Q
QI JUN 已提交
317 318 319
      set_zero(dev_ctx, filter_grad, static_cast<T>(0));
      math::Im2ColFunctor<math::ColFormat::kCFO, DeviceContext, T> im2col;
      math::Vol2ColFunctor<DeviceContext, T> vol2col;
C
chengduoZH 已提交
320 321 322 323 324 325 326 327 328
      for (int i = 0; i < batch_size; i++) {
        Tensor out_grad_batch =
            output_grad->Slice(i, i + 1).Resize(output_matrix_shape);
        Tensor in_batch = input->Slice(i, i + 1).Resize(input_shape);
        for (int g = 0; g < groups; g++) {
          // im2col
          Tensor out_grad_slice =
              out_grad_batch.Slice(g * out_step, (g + 1) * out_step);
          Tensor in_slice = in_batch.Slice(g * in_step, (g + 1) * in_step);
C
chengduoZH 已提交
329

C
chengduoZH 已提交
330
          if (!is_expand) {
C
chengduoZH 已提交
331 332 333
            col.ShareDataWith(in_slice);
            col_matrix.ShareDataWith(col);
            col_matrix.Resize(col_matrix_shape);
C
chengduoZH 已提交
334
          } else if (data_dim == 2U) {
Q
QI JUN 已提交
335
            im2col(dev_ctx, in_slice, dilations, strides,
C
chengduoZH 已提交
336 337 338
                   std::vector<int>{paddings[0], paddings[1], paddings[0],
                                    paddings[1]},
                   &col);
C
chengduoZH 已提交
339
          } else if (data_dim == 3U) {
Q
QI JUN 已提交
340
            vol2col(dev_ctx, in_slice, dilations, strides, paddings, &col);
C
chengduoZH 已提交
341
          }
C
chengduoZH 已提交
342 343 344 345

          // gemm
          Tensor filter_grad_slice =
              filter_grad_.Slice(g * out_step, (g + 1) * out_step);
Q
QI JUN 已提交
346 347 348
          math::matmul<DeviceContext, T>(dev_ctx, out_grad_slice, false,
                                         col_matrix, true, T(1.0),
                                         &filter_grad_slice, T(1.0));
C
chengduoZH 已提交
349 350 351 352 353
        }
      }
    }
  }
};
Z
zlx 已提交
354 355 356 357 358 359 360 361 362 363

template <typename DeviceContext, typename T>
class DepthwiseConvKernel : public framework::OpKernel<T> {
 public:
  void Compute(const framework::ExecutionContext& context) const override {
    const Tensor* input = context.Input<Tensor>("Input");
    Tensor filter = *context.Input<Tensor>("Filter");
    Tensor* output = context.Output<Tensor>("Output");
    output->mutable_data<T>(context.GetPlace());

X
xzl 已提交
364 365 366
    PADDLE_ENFORCE_EQ(
        output->dims()[1] % input->dims()[1], 0,
        "The output channels must be a multiple of the input channels");
Z
zlx 已提交
367 368 369 370 371 372 373
    std::vector<int> strides = context.Attr<std::vector<int>>("strides");
    std::vector<int> paddings = context.Attr<std::vector<int>>("paddings");
    std::vector<int> dilations = context.Attr<std::vector<int>>("dilations");

    math::DepthwiseConvFunctor<DeviceContext, T> depthwiseConv;

    auto& dev_ctx = context.template device_context<DeviceContext>();
374 375 376 377 378 379 380 381 382 383 384 385 386 387 388 389 390 391 392 393 394 395 396 397 398 399 400 401 402 403 404 405 406 407 408 409 410 411 412 413 414 415 416 417
    depthwiseConv(dev_ctx, *input, filter, strides, paddings, output);
  }
};

template <typename DeviceContext, typename T>
class DepthwiseConvGradKernel : public framework::OpKernel<T> {
 public:
  void Compute(const framework::ExecutionContext& context) const override {
    const Tensor* input = context.Input<Tensor>("Input");
    const Tensor* output_grad =
        context.Input<Tensor>(framework::GradVarName("Output"));
    Tensor* input_grad =
        context.Output<Tensor>(framework::GradVarName("Input"));
    Tensor* filter_grad =
        context.Output<Tensor>(framework::GradVarName("Filter"));
    Tensor filter = *context.Input<Tensor>("Filter");

    if (!input_grad && !filter_grad) return;

    std::vector<int> strides = context.Attr<std::vector<int>>("strides");
    std::vector<int> paddings = context.Attr<std::vector<int>>("paddings");
    std::vector<int> dilations = context.Attr<std::vector<int>>("dilations");

    math::SetConstant<DeviceContext, T> set_zero;
    auto& dev_ctx = context.template device_context<DeviceContext>();

    math::DepthwiseConvInputGradFunctor<DeviceContext, T>
        depthwiseConvInputGrad;
    math::DepthwiseConvFilterGradFunctor<DeviceContext, T>
        depthwiseConvFilterGrad;

    if (input_grad) {
      input_grad->mutable_data<T>(context.GetPlace());
      set_zero(dev_ctx, input_grad, static_cast<T>(0));
      depthwiseConvInputGrad(dev_ctx, *input, filter, *output_grad, strides,
                             paddings, input_grad);
    }

    if (filter_grad) {
      filter_grad->mutable_data<T>(context.GetPlace());
      set_zero(dev_ctx, filter_grad, static_cast<T>(0));
      depthwiseConvFilterGrad(dev_ctx, *input, *output_grad, strides, paddings,
                              filter_grad);
    }
Z
zlx 已提交
418 419 420
  }
};

421 422
}  // namespace operators
}  // namespace paddle