Skip to content

Commit

Permalink
add log_softmax_op_npu (#35006)
Browse files Browse the repository at this point in the history
* add log_softmax_op_npu

* log_softmax_op_v1

* import test_log_softmax_grad
  • Loading branch information
juneweng authored Sep 3, 2021
1 parent bb63396 commit ba6a312
Show file tree
Hide file tree
Showing 2 changed files with 140 additions and 0 deletions.
45 changes: 45 additions & 0 deletions paddle/fluid/operators/log_softmax_op_npu.cc
Original file line number Diff line number Diff line change
@@ -0,0 +1,45 @@
// Copyright (c) 2021 PaddlePaddle Authors. All Rights Reserved.
//
// 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.

#include "paddle/fluid/operators/log_softmax_op.h"
#include "paddle/fluid/operators/npu_op_runner.h"
namespace paddle {
namespace operators {
template <typename DeviceContext, typename T>
class LogSoftmaxNPUKernel : public framework::OpKernel<T> {
public:
void Compute(const framework::ExecutionContext& ctx) const override {
auto* X = ctx.Input<framework::Tensor>("X");
auto* Out = ctx.Output<framework::Tensor>("Out");
const int rank = X->dims().size();
const int axis = CanonicalAxis(ctx.Attr<int>("axis"), rank);
std::vector<int> axes;
axes.push_back(axis);
framework::NPUAttributeMap attr_input = {{"axes", axes}};
Out->mutable_data<T>(ctx.GetPlace());
const auto& runner = NpuOpRunner("LogSoftmaxV2", {*X}, {*Out}, attr_input);
auto stream =
ctx.template device_context<paddle::platform::NPUDeviceContext>()
.stream();
runner.Run(stream);
}
};
} // namespace operators
} // namespace paddle
namespace ops = paddle::operators;
namespace plat = paddle::platform;

REGISTER_OP_NPU_KERNEL(
log_softmax,
ops::LogSoftmaxNPUKernel<paddle::platform::NPUDeviceContext, float>);
95 changes: 95 additions & 0 deletions python/paddle/fluid/tests/unittests/npu/test_log_softmax_op_npu.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,95 @@
# Copyright (c) 2021 PaddlePaddle Authors. All Rights Reserved.
#
# 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.

from __future__ import print_function
import numpy as np
import unittest
import sys
sys.path.append("..")
from op_test import OpTest
import paddle
import paddle.fluid as fluid
from paddle.fluid import core
import paddle.nn.functional as F
from test_log_softmax import ref_log_softmax, ref_log_softmax_grad
paddle.enable_static()
np.random.seed(10)


class TestLogSoftmaxNPUOp(OpTest):
def setUp(self):
self.set_npu()
self.place = paddle.NPUPlace(0)
self.op_type = "log_softmax"
self.dtype = np.float32
self.shape = [2, 3, 4, 5]
self.axis = -1
self.set_attrs()
self.set_dtype()
x = np.random.uniform(0.1, 1., self.shape).astype(self.dtype)
out = np.apply_along_axis(ref_log_softmax, self.axis, x)
self.x_grad = ref_log_softmax_grad(x, self.axis)
self.inputs = {'X': x}
self.outputs = {'Out': out}
self.attrs = {'axis': self.axis}

def set_npu(self):
self.__class__.use_npu = True
self.__class__.no_need_check_grad = True

def set_attrs(self):
pass

def set_dtype(self):
pass

def test_check_output(self):
self.check_output_with_place(self.place)

def test_check_grad(self):
pass


def test_class(op_type, typename):
class TestLogSoftmaxShape(TestLogSoftmaxNPUOp):
def set_attrs(self):
self.shape = [12, 10]

def set_dtype(self):
self.dtype = typename

cls_name = "{0}_{1}_1".format(op_type, typename)
TestLogSoftmaxShape.__name__ = cls_name
globals()[cls_name] = TestLogSoftmaxShape


def test_class2(op_type, typename):
class TestLogSoftmaxAxis(TestLogSoftmaxNPUOp):
def set_attrs(self):
self.axis = 0

def set_dtype(self):
self.dtype = typename

cls_name = "{0}_{1}_2".format(op_type, typename)

TestLogSoftmaxAxis.__name__ = cls_name
globals()[cls_name] = TestLogSoftmaxAxis


for _typename in {'float32'}:
test_class("logsoftmax", _typename)
test_class2("logsoftmax", _typename)
if __name__ == '__main__':
unittest.main()

0 comments on commit ba6a312

Please sign in to comment.