clip_by_norm_op.h 3.1 KB
Newer Older
1
/* Copyright (c) 2016 PaddlePaddle Authors. All Rights Reserved.
W
wwhu 已提交
2

L
Luo Tao 已提交
3 4 5
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
W
wwhu 已提交
6

L
Luo Tao 已提交
7
    http://www.apache.org/licenses/LICENSE-2.0
W
wwhu 已提交
8

L
Luo Tao 已提交
9 10 11 12 13
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. */
W
wwhu 已提交
14 15 16

#pragma once

Y
Yi Wang 已提交
17 18
#include "paddle/fluid/framework/eigen.h"
#include "paddle/fluid/framework/op_registry.h"
19
#include "paddle/fluid/operators/math/selected_rows_functor.h"
Y
Yi Wang 已提交
20
#include "paddle/fluid/platform/transform.h"
W
wwhu 已提交
21 22 23 24 25 26 27 28 29

namespace paddle {
namespace operators {

using Tensor = framework::Tensor;
template <typename T, int MajorType = Eigen::RowMajor,
          typename IndexType = Eigen::DenseIndex>
using EigenVector = framework::EigenVector<T, MajorType, IndexType>;

Q
QI JUN 已提交
30
template <typename DeviceContext, typename T>
W
wwhu 已提交
31 32 33 34
class ClipByNormKernel : public framework::OpKernel<T> {
 public:
  void Compute(const framework::ExecutionContext& context) const override {
    auto max_norm = context.Attr<T>("max_norm");
35
    auto in_var = context.InputVar("X");
W
wwhu 已提交
36

M
minqiyang 已提交
37
    Tensor* output = nullptr;
38 39 40
    const Tensor* input = nullptr;
    if (in_var->IsType<framework::LoDTensor>()) {
      input = context.Input<Tensor>("X");
M
minqiyang 已提交
41 42 43

      output = context.Output<Tensor>("Out");
      output->mutable_data<T>(context.GetPlace());
44 45 46 47 48 49 50 51 52 53 54
    } else if (in_var->IsType<framework::SelectedRows>()) {
      auto* x = context.Input<framework::SelectedRows>("X");

      // merge ids in selected rows first
      math::scatter::MergeAdd<DeviceContext, T> merge_func;
      auto* merged_input = const_cast<framework::Scope&>(context.scope())
                               .Var()
                               ->GetMutable<framework::SelectedRows>();
      merge_func(context.template device_context<DeviceContext>(), *x,
                 merged_input);
      input = &(merged_input->value());
M
minqiyang 已提交
55 56 57 58 59

      auto* output_selected_rows = context.Output<SelectedRows>("Out");
      output_selected_rows->set_rows(merged_input.rows());
      output = output_selected_rows->mutable_data();
      output->Resize(framework::make_ddim(merged_input.value().dims()));
60 61 62 63 64 65 66
    } else {
      PADDLE_THROW("Unexpected branch, input variable type is %s",
                   in_var->Type().name());
    }

    PADDLE_ENFORCE_NOT_NULL(input);

W
wwhu 已提交
67 68 69
    auto x = EigenVector<T>::Flatten(*input);
    auto out = EigenVector<T>::Flatten(*output);
    auto x_norm = x.square().sum().sqrt();
Q
QI JUN 已提交
70 71
    auto& place =
        *context.template device_context<DeviceContext>().eigen_device();
W
wwhu 已提交
72 73 74 75 76 77 78 79 80 81 82

    auto temp = (x_norm <= max_norm).template cast<T>().eval();
    auto scaling = temp + (static_cast<T>(1) - temp) * max_norm / x_norm;
    Eigen::array<int, 1> one_dim{{1}};
    Eigen::DSizes<int, 1> m_dsize(input->numel());
    out.device(place) = x * scaling.reshape(one_dim).broadcast(m_dsize);
  }
};

}  // namespace operators
}  // namespace paddle