gemm_conv2d_op.h 9.0 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

17
#include "paddle/framework/eigen.h"
18 19 20 21 22 23 24 25 26 27
#include "paddle/framework/op_registry.h"
#include "paddle/operators/math/im2col.h"
#include "paddle/operators/math/math_function.h"

namespace paddle {
namespace operators {

using Tensor = framework::Tensor;

template <typename Place, typename T>
H
hedaoyuan 已提交
28
class GemmConv2DKernel : public framework::OpKernel {
29 30 31
 public:
  void Compute(const framework::ExecutionContext& context) const override {
    const Tensor* input = context.Input<Tensor>("Input");
H
hedaoyuan 已提交
32 33 34 35
    // 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");
36 37 38 39 40
    Tensor* output = context.Output<Tensor>("Output");
    output->mutable_data<T>(context.GetPlace());

    std::vector<int> strides = context.Attr<std::vector<int>>("strides");
    std::vector<int> paddings = context.Attr<std::vector<int>>("paddings");
41
    int groups = context.Attr<int>("groups");
42 43 44

    int batch_size = input->dims()[0];
    int input_channels = input->dims()[1];
H
hedaoyuan 已提交
45 46 47
    int filter_height = filter.dims()[filter.dims().size() - 2];
    int filter_width = filter.dims()[filter.dims().size() - 1];
    int output_channels = output->dims()[1];
48 49 50 51 52 53
    int output_height = output->dims()[2];
    int output_width = output->dims()[3];

    paddle::operators::math::Im2ColFunctor<
        paddle::operators::math::ColFormat::kCFO, Place, T>
        im2col;
H
hedaoyuan 已提交
54
    // use col_shape in the im2col calculation
55 56
    framework::DDim col_shape = {input_channels / groups, filter_height,
                                 filter_width, output_height, output_width};
H
hedaoyuan 已提交
57 58
    // use col_matrix_shape in the gemm calculation
    framework::DDim col_matrix_shape = {
59
        input_channels / groups * filter_height * filter_width,
H
hedaoyuan 已提交
60
        output_height * output_width};
61
    Tensor col;
62
    col.mutable_data<T>(col_shape, context.GetPlace());
H
hedaoyuan 已提交
63 64 65 66 67
    // 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.
    Tensor col_matrix = col;
    col_matrix.Resize(col_matrix_shape);
68 69 70

    framework::DDim input_shape = {input->dims()[1], input->dims()[2],
                                   input->dims()[3]};
71 72
    framework::DDim filter_matrix_shape = {filter.dims()[0],
                                           filter.numel() / filter.dims()[0]};
H
hedaoyuan 已提交
73 74 75 76 77 78 79
    filter.Resize(filter_matrix_shape);

    framework::DDim output_matrix_shape = {output_channels,
                                           output_height * output_width};

    auto* device_context =
        const_cast<platform::DeviceContext*>(context.device_context_);
80

81
    // convolution operator: im2col + gemm
82 83
    int in_step = input_channels / groups;
    int out_step = output_channels / groups;
84
    for (int i = 0; i < batch_size; i++) {
85 86
      Tensor in_batch = input->Slice<T>(i, i + 1).Resize(input_shape);
      Tensor out_batch = output->Slice<T>(i, i + 1).Resize(output_matrix_shape);
87 88
      for (int g = 0; g < groups; g++) {
        // im2col
89
        Tensor in_slice = in_batch.Slice<T>(g * in_step, (g + 1) * in_step);
90 91 92 93
        im2col(in_slice, col, strides[0], strides[1], paddings[0], paddings[1],
               device_context);

        // gemm
94
        Tensor out_slice = out_batch.Slice<T>(g * out_step, (g + 1) * out_step);
95 96 97 98
        Tensor filter_slice = filter.Slice<T>(g * out_step, (g + 1) * out_step);
        math::matmul<Place, T>(filter_slice, false, col_matrix, false, T(1.0),
                               &out_slice, T(0.0), device_context);
      }
99 100 101 102 103
    }
  }
};

template <typename Place, typename T>
H
hedaoyuan 已提交
104
class GemmConvGrad2DKernel : public framework::OpKernel {
105 106
 public:
  void Compute(const framework::ExecutionContext& context) const override {
107 108 109 110 111
    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"));
112
    Tensor* filter_grad_ =
113 114
        context.Output<Tensor>(framework::GradVarName("Filter"));
    input_grad->mutable_data<T>(context.GetPlace());
115 116 117 118 119 120 121
    filter_grad_->mutable_data<T>(context.GetPlace());

    // 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");
    Tensor filter_grad = *filter_grad_;
122 123 124

    std::vector<int> strides = context.Attr<std::vector<int>>("strides");
    std::vector<int> paddings = context.Attr<std::vector<int>>("paddings");
125
    int groups = context.Attr<int>("groups");
126 127 128

    int batch_size = input->dims()[0];
    int input_channels = input->dims()[1];
129 130
    int filter_height = filter.dims()[filter.dims().size() - 2];
    int filter_width = filter.dims()[filter.dims().size() - 1];
131
    int output_channels = output_grad->dims()[1];
132 133 134 135 136 137 138 139 140
    int output_height = output_grad->dims()[2];
    int output_width = output_grad->dims()[3];

    paddle::operators::math::Col2ImFunctor<
        paddle::operators::math::ColFormat::kCFO, Place, T>
        col2im;
    paddle::operators::math::Im2ColFunctor<
        paddle::operators::math::ColFormat::kCFO, Place, T>
        im2col;
141
    // use col_shape in the im2col and col2im calculation
142 143
    framework::DDim col_shape = {input_channels / groups, filter_height,
                                 filter_width, output_height, output_width};
144 145
    // use col_matrix_shape in the gemm calculation
    framework::DDim col_matrix_shape = {
146
        input_channels / groups * filter_height * filter_width,
147 148
        output_height * output_width};
    Tensor col;
149
    col.mutable_data<T>(col_shape, context.GetPlace());
150 151 152 153 154
    // 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.
    Tensor col_matrix = col;
    col_matrix.Resize(col_matrix_shape);
155 156 157 158 159 160 161

    framework::DDim input_shape = {input->dims()[1], input->dims()[2],
                                   input->dims()[3]};
    framework::DDim output_matrix_shape = {
        output_grad->dims()[1],
        output_grad->dims()[2] * output_grad->dims()[3]};

162 163
    framework::DDim filter_matrix_shape = {filter.dims()[0],
                                           filter.numel() / filter.dims()[0]};
164 165 166 167
    filter.Resize(filter_matrix_shape);
    filter_grad.Resize(filter_matrix_shape);

    auto t1 = framework::EigenVector<T>::Flatten(filter_grad);
168 169 170 171
    t1.device(context.GetEigenDevice<Place>()) = t1.constant(static_cast<T>(0));
    auto t2 = framework::EigenVector<T>::Flatten(*input_grad);
    t2.device(context.GetEigenDevice<Place>()) = t2.constant(static_cast<T>(0));

172 173 174
    auto* device_context =
        const_cast<platform::DeviceContext*>(context.device_context_);

175 176
    // convolution backward input operator:  gemm + col2im
    // convolution backward weight operator: im2col + gemm
177 178
    int in_step = input_channels / groups;
    int out_step = output_channels / groups;
179
    for (int i = 0; i < batch_size; i++) {
180
      Tensor out_grad_batch =
181
          output_grad->Slice<T>(i, i + 1).Resize(output_matrix_shape);
182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208
      Tensor in_grad_batch = input_grad->Slice<T>(i, i + 1).Resize(input_shape);
      Tensor in_batch = input->Slice<T>(i, i + 1).Resize(input_shape);
      for (int g = 0; g < groups; g++) {
        // gemm
        Tensor out_grad_slice =
            out_grad_batch.Slice<T>(g * out_step, (g + 1) * out_step);
        Tensor filter_slice = filter.Slice<T>(g * out_step, (g + 1) * out_step);
        math::matmul<Place, T>(filter_slice, true, out_grad_slice, false,
                               T(1.0), &col_matrix, T(0.0), device_context);

        // 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], device_context);

        // 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],
               device_context);

        // gemm
        Tensor filter_grad_slice =
            filter_grad.Slice<T>(g * out_step, (g + 1) * out_step);
        math::matmul<Place, T>(out_grad_slice, false, col_matrix, true, T(1.0),
                               &filter_grad_slice, T(1.0), device_context);
      }
209
    }
210 211 212 213 214
  }
};

}  // namespace operators
}  // namespace paddle
新手
引导
客服 返回
顶部