src/operator/nn/convolution.cu - master - githubmirror/incubator-mxnet

src/operator/nn/convolution.cu
208 строк8 KB
Vladimir Cherepanov
Use cuDNN for conv bias and bias grad (#20771)
17 дек 2021, 21:46
17 дек 2021, 21:46f0ef9d8
/*
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you 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.
 */

/*!
 * \file convolution.cu
 * \brief
 * \author Bing Xu, Jun Wu, Da Zheng
 */

#include "./convolution-inl.h"
#include <vector>
#include "./depthwise_convolution-inl.h"
#if MXNET_USE_CUDNN == 1
#include "../cudnn_ops.h"
#include "../tensor/broadcast_reduce_op.h"
#include "../tensor/elemwise_binary_broadcast_op.h"
#include "fully_connected-inl.h"
#endif  // MXNET_USE_CUDNN

namespace mxnet {
namespace op {

template <>
void ConvolutionCompute<gpu>(const nnvm::NodeAttrs& attrs,
                             const OpContext& ctx,
                             const std::vector<TBlob>& inputs,
                             const std::vector<OpReqType>& req,
                             const std::vector<TBlob>& outputs) {
  const ConvolutionParam& param = nnvm::get<ConvolutionParam>(attrs.parsed);
  int dtype                     = inputs[conv::kData].type_flag_;
  CHECK_EQ(req.size(), 1);
  CHECK_EQ(req[conv::kOut], kWriteTo);

#if MXNET_USE_CUDNN == 1
  STATIC_ASSERT_CUDNN_VERSION_GE(8000);
  MSHADOW_REAL_TYPE_SWITCH(dtype, DType, {
    cudnn::ConvParam conv_param(param, false);
    bool ok = !param.cudnn_off &&
              cudnn::Exec<cudnn::Conv>(
                  ctx, conv_param, inputs[conv::kData], inputs[conv::kWeight], outputs[conv::kOut]);
    if (ok && !param.no_bias) {
      CHECK_EQ(inputs[conv::kBias].shape_.ndim(), 1);
      auto layout = static_cast<mshadow::LayoutFlag>(param.layout.value());
      auto li     = cudnn::GetLayoutInfo(layout);
      if (li.channel_last ||
          !cudnn::LegacyAddBias(ctx, li, outputs[conv::kOut], inputs[conv::kBias])) {
        int k  = inputs[conv::kBias].shape_.Size();
        auto b = inputs[conv::kBias].reshape(cudnn::ExpandChannelDims(layout, k));
        BinaryBroadcastRTCCompute{"add"}(  // NOLINT(whitespace/braces)
            attrs,
            ctx,
            {outputs[conv::kOut], b},
            {kWriteInplace},
            {outputs[conv::kOut]});
      }
    }
    if (!ok) {
      if (!param.cudnn_off)
        LOG(WARNING) << "This convolution is not supported by cuDNN, MXNet convolution is applied.";
      ConvolutionOp<gpu, DType> op;
      op.Init(param);
      op.Forward(ctx, inputs, req, outputs);
    }
  })
#else
  if (param.layout.value() != kNCW && param.layout.value() != kNCHW &&
      param.layout.value() != kNCDHW) {
    // Need CuDNN > 5.0 for layout support. use MXNet implementation
    MSHADOW_REAL_TYPE_SWITCH(dtype, DType, {
      ConvolutionOp<gpu, DType> op;
      op.Init(param);
      op.Forward(ctx, inputs, req, outputs);
    })
    return;
  }

  if (param.num_filter == param.num_group && param.layout.value() == mshadow::kNCHW &&
      param.num_filter == inputs[conv::kData].shape_[1] && param.kernel.ndim() == 2 &&
      param.dilate == mshadow::Shape2(1, 1) && dtype == mshadow::kFloat32) {
    mxnet::ShapeVector in_shape(inputs.size());
    mxnet::ShapeVector out_shape(1, outputs[0].shape_);
    for (size_t i = 0; i < in_shape.size(); i++)
      in_shape[i] = inputs[i].shape_;
    DepthwiseConvolutionOp<float> op;
    op.Init(param, in_shape, out_shape);
    op.Forward(ctx, inputs, req, outputs);
    return;
  }

  MSHADOW_REAL_TYPE_SWITCH(dtype, DType, {
    ConvolutionOp<gpu, DType> op;
    op.Init(param);
    op.Forward(ctx, inputs, req, outputs);
  })
#endif  // MXNET_USE_CUDNN
}

template <>
void ConvolutionGradCompute<gpu>(const nnvm::NodeAttrs& attrs,
                                 const OpContext& ctx,
                                 const std::vector<TBlob>& inputs,
                                 const std::vector<OpReqType>& req,
                                 const std::vector<TBlob>& outputs) {
  const ConvolutionParam& param = nnvm::get<ConvolutionParam>(attrs.parsed);
  std::vector<TBlob> in_data(inputs.begin() + 1, inputs.end());
  const TBlob& out_grad             = inputs[0];
  const std::vector<TBlob>& in_grad = outputs;
  int dtype                         = out_grad.type_flag_;
  CHECK_EQ(req.size(), param.no_bias ? 2 : 3);
  CHECK_NE(req[conv::kData], kWriteInplace);
  CHECK_NE(req[conv::kWeight], kWriteInplace);
  if (!param.no_bias)
    CHECK_NE(req[conv::kBias], kWriteInplace);

#if MXNET_USE_CUDNN == 1
  STATIC_ASSERT_CUDNN_VERSION_GE(8000);
  MSHADOW_REAL_TYPE_SWITCH(dtype, DType, {
    cudnn::ConvParam conv_param(param, req[conv::kData] == kAddTo);
    bool ok           = !param.cudnn_off;
    ok                = ok && (req[conv::kData] == kNullOp ||
                cudnn::Exec<cudnn::ConvDgrad>(
                    ctx, conv_param, inputs[1 + conv::kWeight], inputs[0], outputs[conv::kData]));
    conv_param.add_to = req[conv::kWeight] == kAddTo;
    ok                = ok && (req[conv::kWeight] == kNullOp ||
                cudnn::Exec<cudnn::ConvWgrad>(
                    ctx, conv_param, inputs[1 + conv::kData], inputs[0], outputs[conv::kWeight]));
    if (ok && !param.no_bias && req[conv::kBias] != kNullOp) {
      auto li     = cudnn::GetLayoutInfo(static_cast<mshadow::LayoutFlag>(param.layout.value()));
      auto add_to = req[conv::kBias] == kAddTo;
      if (li.channel_last ||
          !cudnn::LegacyBiasGrad(ctx, li, add_to, outputs[conv::kBias], inputs[0])) {
        if (li.channel_last) {
          // This kernel should be faster.
          auto y_grad = FlattenAs2DHead<gpu, DType>(inputs[0], ctx);
          AddBiasGrad(outputs[conv::kBias], y_grad, req[conv::kBias], param.num_filter, ctx);
        } else {
          TShape axes{static_cast<int>(li.ChannelIdx())};
          TShape small = ReduceAxesShapeImpl(
              inputs[0].shape_, dmlc::optional<mxnet::TShape>(axes), true, true);
          ReduceAxesRTCComputeImpl(
              ctx, {inputs[0]}, {req[conv::kBias]}, {outputs[conv::kBias]}, small, "red::sum{}");
        }
      }
    }
    if (!ok) {
      if (!param.cudnn_off)
        LOG(WARNING) << "This convolution backward is not supported by cuDNN, MXNet op is applied.";
      ConvolutionOp<gpu, DType> op;
      op.Init(param);
      op.Backward(ctx, std::vector<TBlob>{out_grad}, in_data, req, in_grad);
    }
  })
#else
  if (param.layout.value() != kNCW && param.layout.value() != kNCHW &&
      param.layout.value() != kNCDHW) {
    MSHADOW_REAL_TYPE_SWITCH(dtype, DType, {
      ConvolutionOp<gpu, DType> op;
      op.Init(param);
      op.Backward(ctx, std::vector<TBlob>{out_grad}, in_data, req, in_grad);
    })
    return;
  }

  if (param.num_filter == param.num_group && param.layout.value() == mshadow::kNCHW &&
      param.num_filter == in_data[conv::kData].shape_[1] && param.kernel.ndim() == 2 &&
      param.dilate == mshadow::Shape2(1, 1) && dtype == mshadow::kFloat32) {
    // The first element stores out grad.
    mxnet::ShapeVector in_shape(in_data.size());
    mxnet::ShapeVector out_shape(1, out_grad.shape_);
    for (size_t i = 0; i < in_shape.size(); i++)
      in_shape[i] = in_data[i].shape_;
    DepthwiseConvolutionOp<float> op;
    op.Init(param, in_shape, out_shape);
    op.Backward(ctx, std::vector<TBlob>{out_grad}, in_data, req, in_grad);
    return;
  }

  MSHADOW_REAL_TYPE_SWITCH(dtype, DType, {
    ConvolutionOp<gpu, DType> op;
    op.Init(param);
    op.Backward(ctx, std::vector<TBlob>{out_grad}, in_data, req, in_grad);
  })
#endif  // MXNET_USE_CUDNN
}

NNVM_REGISTER_OP(Convolution).set_attr<FCompute>("FCompute<gpu>", ConvolutionCompute<gpu>);

NNVM_REGISTER_OP(_backward_Convolution)
    .set_attr<FCompute>("FCompute<gpu>", ConvolutionGradCompute<gpu>);

}  // namespace op
}  // namespace mxnet