saymrwulf/onnxruntime
1
0
Fork 0
mirror of https://github.com/saymrwulf/onnxruntime.git synced 2026-07-30 20:18:08 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Rewrite the softmax op (#1911)

Browse source
This commit is contained in:
Changming Sun 2019-09-25 14:29:33 -07:00 committed by GitHub
parent 4d26f2ce86
commit 76f76251fa
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
9 changed files with 112 additions and 297 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

58
onnxruntime/core/providers/cpu/math/logsoftmax.cc
View file

@ -1,58 +0,0 @@
// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.
#include "core/providers/cpu/math/logsoftmax.h"
#include "core/framework/op_kernel.h"
#include "core/framework/op_kernel_context_internal.h"
#include "core/providers/common.h"
#include "core/providers/cpu/math/softmax_shared.h"
#include "core/util/math.h"
namespace onnxruntime {
template <>
Status LogSoftmax<float>::Compute(OpKernelContext* ctx) const {
auto ctx_internal = static_cast<OpKernelContextInternal*>(ctx);
concurrency::ThreadPool* tp = ctx_internal->GetOperatorThreadPool();
const auto* tensor_pointer = ctx->Input<Tensor>(0);
if (tensor_pointer == nullptr) return Status(common::ONNXRUNTIME, common::FAIL, "input count mismatch");
const Tensor& X = *tensor_pointer;
const TensorShape& input_shape{X.Shape()};
Tensor* Y = ctx->Output(0, input_shape);
const int64_t axis = HandleNegativeAxis(axis_, input_shape.NumDimensions());
size_t N = input_shape.SizeToDimension(axis);
size_t D = input_shape.SizeFromDimension(axis);
auto* Ydata = Y->template MutableData<float>();
std::vector<float> scale_(N);
std::vector<float> rowmax_(N);
std::vector<float> sum_multiplier_(D, 1.f); // initialize all multiplier values to 1.0
const bool logarithmic = true;
auto status = SoftmaxCPU(N, D, X.template Data<float>(), Ydata,
scale_.data(), sum_multiplier_.data(), logarithmic, rowmax_.data(), tp);
return status;
}
ONNX_CPU_OPERATOR_VERSIONED_KERNEL(
LogSoftmax,
1,
10,
KernelDefBuilder().TypeConstraint("T", DataTypeImpl::GetTensorType<float>()),
LogSoftmax<float>);
// Opset 11 starts to support Neg Axis.
ONNX_CPU_OPERATOR_KERNEL(
LogSoftmax,
11,
KernelDefBuilder().TypeConstraint("T", DataTypeImpl::GetTensorType<float>()),
LogSoftmax<float>);
} // namespace onnxruntime

30
onnxruntime/core/providers/cpu/math/logsoftmax.h
View file

@ -1,30 +0,0 @@
// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.
#pragma once
#include "gsl/gsl_util"
#include "core/common/common.h"
#include "core/framework/op_kernel.h"
#include "core/util/math_cpuonly.h"
namespace onnxruntime {
template <typename T>
class LogSoftmax final : public OpKernel {
public:
LogSoftmax(const OpKernelInfo& info) : OpKernel{info}, axis_{1} {
int64_t axis;
Status status = info.GetAttr<int64_t>("axis", &axis);
if (status.IsOK()) {
axis_ = gsl::narrow_cast<int>(axis);
}
}
Status Compute(OpKernelContext* context) const override;
private:
int axis_;
};
} // namespace onnxruntime

59
onnxruntime/core/providers/cpu/math/softmax.cc
View file

@ -1,60 +1,29 @@
// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.
#include "core/providers/cpu/math/softmax.h"
#include "core/framework/op_kernel.h"
#include "core/framework/op_kernel_context_internal.h"
#include "core/providers/common.h"
#include "core/providers/cpu/math/softmax_shared.h"
#include "core/util/math.h"
#include "core/util/eigen_common_wrapper.h"
namespace onnxruntime {
template <>
Status Softmax<float>::Compute(OpKernelContext* ctx) const {
auto ctx_internal = static_cast<OpKernelContextInternal*>(ctx);
concurrency::ThreadPool* tp = ctx_internal->GetOperatorThreadPool();
const auto* tensor_pointer = ctx->Input<Tensor>(0);
if (tensor_pointer == nullptr) return Status(common::ONNXRUNTIME, common::FAIL, "input count mismatch");
const Tensor& X = *tensor_pointer;
const TensorShape& input_shape{X.Shape()};
VLOGS(ctx->Logger(), 2) << "Input tensor shape: " << input_shape;
Tensor* Y = ctx->Output(0, input_shape);
const int64_t axis = HandleNegativeAxis(axis_, input_shape.NumDimensions());
size_t N = input_shape.SizeToDimension(axis);
size_t D = input_shape.SizeFromDimension(axis);
auto* Ydata = Y->template MutableData<float>();
std::vector<float> scale_(N);
std::vector<float> rowmax_(N);
std::vector<float> sum_multiplier_(D, 1.f); // initialize all multiplier values to 1.0
const bool logarithmic = false;
auto status = SoftmaxCPU(N, D, X.template Data<float>(), Ydata,
scale_.data(), sum_multiplier_.data(), logarithmic, rowmax_.data(), tp);
return status;
}
ONNX_CPU_OPERATOR_VERSIONED_KERNEL(
Softmax,
1,
10,
KernelDefBuilder().TypeConstraint("T", DataTypeImpl::GetTensorType<float>()),
Softmax<float>);
ONNX_CPU_OPERATOR_VERSIONED_KERNEL(Softmax, 1, 10,
KernelDefBuilder().TypeConstraint("T", DataTypeImpl::GetTensorType<float>()),
Softmax<float, false>);
// Opset 11 starts to support Neg Axis.
ONNX_CPU_OPERATOR_KERNEL(
Softmax,
11,
KernelDefBuilder().TypeConstraint("T", DataTypeImpl::GetTensorType<float>()),
Softmax<float>);
ONNX_CPU_OPERATOR_KERNEL(Softmax, 11, KernelDefBuilder().TypeConstraint("T", DataTypeImpl::GetTensorType<float>()),
Softmax<float, false>);
ONNX_CPU_OPERATOR_VERSIONED_KERNEL(LogSoftmax, 1, 10,
KernelDefBuilder().TypeConstraint("T", DataTypeImpl::GetTensorType<float>()),
Softmax<float, true>);
// Opset 11 starts to support Neg Axis.
ONNX_CPU_OPERATOR_KERNEL(LogSoftmax, 11, KernelDefBuilder().TypeConstraint("T", DataTypeImpl::GetTensorType<float>()),
Softmax<float, true>);
} // namespace onnxruntime

99
onnxruntime/core/providers/cpu/math/softmax.h
View file

@ -1,3 +1,15 @@
/* Copyright 2015 The TensorFlow 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.
==============================================================================*/
// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.
@ -7,10 +19,60 @@
#include "core/common/common.h"
#include "core/framework/op_kernel.h"
#include "core/framework/op_kernel_context_internal.h"
#include "core/util/math_cpuonly.h"
#include "core/util/eigen_common_wrapper.h"
#include "core/providers/common.h"
namespace onnxruntime {
template <typename T>
// copied from tensorflow/core/kernels/softmax_op.cc
template <typename Device, typename T>
static void ComputeSoftMax(
const Device& d,
typename Eigen::TensorMap<Eigen::Tensor<const T, 2, Eigen::RowMajor, Eigen::DenseIndex>, Eigen::Aligned> logits,
typename Eigen::TensorMap<Eigen::Tensor<T, 2, Eigen::RowMajor, Eigen::DenseIndex>, Eigen::Aligned> softmax,
const bool log) {
const int kBatchDim = 0;
const int kClassDim = 1;
const int batch_size = (int)logits.dimension(kBatchDim);
const int num_classes = (int)logits.dimension(kClassDim);
// These arrays are used to reduce along the class dimension, and broadcast
// the resulting value to all classes.
#if !defined(EIGEN_HAS_INDEX_LIST)
Eigen::DSizes<int, 1> along_class(kClassDim);
Eigen::DSizes<int, 2> batch_by_one(batch_size, 1);
Eigen::DSizes<int, 2> one_by_class(1, num_classes);
#else
Eigen::IndexList<Eigen::type2index<kClassDim> > along_class;
Eigen::IndexList<int, Eigen::type2index<1> > batch_by_one;
batch_by_one.set(0, batch_size);
Eigen::IndexList<Eigen::type2index<1>, int> one_by_class;
one_by_class.set(1, num_classes);
#endif
// shifted_logits = logits - max(logits along classes);
auto shifted_logits = (logits - logits.maximum(along_class).eval().reshape(batch_by_one).broadcast(one_by_class));
if (log) {
// Calculate the log of the softmax
// softmax = logits - max(logits along classes);
softmax.device(d) = shifted_logits;
// softmax = softmax - log(sum(exp(softmax along classes)));
softmax.device(d) =
(softmax - softmax.exp().sum(along_class).log().eval().reshape(batch_by_one).broadcast(one_by_class));
} else {
// NOTE(touts): If you modify this implementation please run
// the BM_ImageNetSoftmaxFwd benchmark in nn_ops_test.cc.
//
// softmax = exp(logits - max(logits along classes));
softmax.device(d) = shifted_logits.exp();
// softmax = softmax * (1 / sum(softmax along classes));
softmax.device(d) =
(softmax * softmax.sum(along_class).inverse().eval().reshape(batch_by_one).broadcast(one_by_class));
}
}
template <typename T, bool use_log>
class Softmax final : public OpKernel {
public:
Softmax(const OpKernelInfo& info) : OpKernel{info}, axis_{1} {
@ -22,7 +84,40 @@ class Softmax final : public OpKernel {
}
}
Status Compute(OpKernelContext* context) const override;
Status Compute(OpKernelContext* ctx) const override {
#ifndef USE_OPENMP
auto ctx_internal = static_cast<OpKernelContextInternal*>(ctx);
concurrency::ThreadPool* tp = ctx_internal->GetOperatorThreadPool();
#endif
const auto* tensor_pointer = ctx->Input<Tensor>(0);
if (tensor_pointer == nullptr)
return Status(common::ONNXRUNTIME, common::FAIL, "input count mismatch");
const Tensor& X = *tensor_pointer;
const TensorShape& input_shape{X.Shape()};
VLOGS(ctx->Logger(), 2) << "Input tensor shape: " << input_shape;
Tensor* Y = ctx->Output(0, input_shape);
const int64_t axis = HandleNegativeAxis(axis_, input_shape.NumDimensions());
int N = static_cast<int>(input_shape.SizeToDimension(axis));
int D = static_cast<int>(input_shape.SizeFromDimension(axis));
Eigen::TensorMap<Eigen::Tensor<const float, 2, Eigen::RowMajor, Eigen::DenseIndex>, Eigen::Aligned> X_tensor(
X.Data<float>(), N, D);
Eigen::TensorMap<Eigen::Tensor<float, 2, Eigen::RowMajor, Eigen::DenseIndex>, Eigen::Aligned> Y_tensor(
Y->MutableData<float>(), N, D);
#ifndef USE_OPENMP
if (tp == nullptr)
#endif
ComputeSoftMax(Eigen::DefaultDevice(), X_tensor, Y_tensor, use_log);
#ifndef USE_OPENMP
else
ComputeSoftMax(Eigen::ThreadPoolDevice(&tp->GetHandler(), tp->NumThreads()), X_tensor, Y_tensor, use_log);
#endif
return Status::OK();
}
private:
int axis_;

93
onnxruntime/core/providers/cpu/math/softmax_shared.cc
View file

@ -1,93 +0,0 @@
/**
* Derived from caffe2, need copy right announcement here.
*/
/**
* Copyright (c) 2016-present, Facebook, Inc.
*
* 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.
*/
// disable noisy warning about use of std::copy_n by gsl::copy. gsl_algorithm does the same thing
// but that's too late for the disable to work
#ifdef _MSC_VER
#pragma warning(push)
#pragma warning(disable : 4996)
#endif
#include <algorithm>
#include <cmath>
#ifdef _MSC_VER
#pragma warning(pop)
#endif
#include "core/providers/cpu/math/softmax_shared.h"
#include "core/util/math.h"
#include "core/util/math_cpuonly.h"
#include "gsl/gsl_algorithm"
#include "gsl/gsl_util"
namespace onnxruntime {
common::Status SoftmaxCPU(const int64_t N,
const int64_t D,
const float* Xdata,
float* Ydata,
float* scale,
const float* sum_multiplier,
bool logarithmic,
float* rowmax, onnxruntime::concurrency::ThreadPool* tp) {
// the Math functions SoftmaxCPU uses only support int32_t as input, so enforce that
if (N * D > INT32_MAX || N > INT32_MAX || D > INT32_MAX) {
std::ostringstream ss;
ss << "SoftmaxCPU inputs N, D and N * D must be < " << INT32_MAX << ". N=" << N << ", D=" << D;
std::string msg = ss.str();
return Status(common::ONNXRUNTIME, common::INVALID_ARGUMENT, msg);
}
const int n = gsl::narrow_cast<int>(N);
const int d = gsl::narrow_cast<int>(D);
const int nd = gsl::narrow_cast<int>(N * D);
math::RowwiseMax<float, CPUMathUtil>(n, d, Xdata, rowmax, nullptr);
// Put the intermediate result X - max(X) into Y by first copying X to Y, and then subtracting max from each entry
gsl::copy(gsl::make_span(Xdata, nd), gsl::make_span(Ydata, nd));
math::Gemm<float>(CblasNoTrans, CblasNoTrans, n, d, 1, -1, rowmax, sum_multiplier, 1, Ydata, tp);
// Exponentiation
math::Exp<float, CPUMathUtil>(nd, Ydata, Ydata, nullptr);
math::Gemv<float, CPUMathUtil>(CblasNoTrans, n, d, 1, Ydata, sum_multiplier, 0, scale, nullptr);
// Do division
if (!logarithmic) {
for (int i = 0; i < N; ++i) {
for (int j = 0; j < D; ++j) {
Ydata[i * D + j] /= scale[i];
}
}
} else {
for (int i = 0; i < N; ++i) {
auto log_fmaxf_scale_i = std::log(fmaxf(scale[i], 1e-20f));
for (int j = 0; j < D; ++j) {
Ydata[i * D + j] = Xdata[i * D + j] - rowmax[i] - log_fmaxf_scale_i;
}
}
}
return Status::OK();
}
} // namespace onnxruntime

25
onnxruntime/core/providers/cpu/math/softmax_shared.h
View file

@ -1,25 +0,0 @@
// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.
#pragma once
#include "core/common/status.h"
namespace onnxruntime {
namespace concurrency {
class ThreadPool;
}
/**
Calculate Softmax using CPU memory.
@param N Number of rows
@param D Number of elements in each row
@param Xdata Source data
@param Ydata Output data
@param scale Storage for scale calculation. Size must be >= N.
@param sum_multiplier Weights for each element. Size must be >= D.
@param logarithmic If true, compute LogSoftmax. If false compute Softmax.
@param rowmax Storage for calculation of maximum in each row. Size must be >= N.
*/
common::Status SoftmaxCPU(int64_t N, int64_t D, const float* Xdata, float* Ydata, float* scale,
const float* sum_multiplier, bool logarithmic, float* rowmax, concurrency::ThreadPool* tp);
} // namespace onnxruntime

1
onnxruntime/test/providers/cpu/math/hardmax_test.cc
View file

@ -1,7 +1,6 @@
// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.
#include "core/providers/cpu/math/softmax_shared.h"
#include "gmock/gmock.h"
#include "gtest/gtest.h"
#include "test/providers/provider_test_utils.h"

2
onnxruntime/test/providers/cpu/math/logsoftmax_test.cc
View file

@ -1,7 +1,7 @@
// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.
#include "core/providers/cpu/math/logsoftmax.h"
#include "core/providers/cpu/math/softmax.h"
#include "gmock/gmock.h"
#include "gtest/gtest.h"
#include "test/providers/provider_test_utils.h"

42
onnxruntime/test/providers/cpu/math/softmax_test.cc
View file

@ -1,7 +1,6 @@
// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.
#include "core/providers/cpu/math/softmax_shared.h"
#include "gmock/gmock.h"
#include "gtest/gtest.h"
#include "test/providers/provider_test_utils.h"
@ -192,46 +191,5 @@ TEST(SoftmaxOperator, InvalidAxis) {
"-10 is not in valid range [-2,1]");
}
TEST(SoftmaxOperator, TestInputTooLarge) {
float* ignored = nullptr;
concurrency::ThreadPool tp("", 1);
// N > INT32_MAX
int64_t N = int64_t(INT32_MAX) + 1;
int64_t D = 1;
auto status = SoftmaxCPU(N, D, ignored, ignored, ignored, ignored, true, ignored, &tp);
EXPECT_EQ(status.Code(), common::INVALID_ARGUMENT);
// D > INT32_MAX
N = 1;
D = int64_t(INT32_MAX) + 1;
status = SoftmaxCPU(N, D, ignored, ignored, ignored, ignored, true, ignored, &tp);
EXPECT_EQ(status.Code(), common::INVALID_ARGUMENT);
// N * D > INT32_MAX
N = int64_t(INT32_MAX) / 2;
D = 3;
status = SoftmaxCPU(N, D, ignored, ignored, ignored, ignored, true, ignored, &tp);
EXPECT_EQ(status.Code(), common::INVALID_ARGUMENT);
/*
common::Status SoftmaxCPU(const int64_t N,
const int64_t D,
const float* Xdata,
float* Ydata,
float* scale,
const float* sum_multiplier,
bool logarithmic,
float* rowmax)
{
// the Math functions SoftmaxCPU uses only support int32_t as input, so enforce that
if (N * D > INT32_MAX || N > INT32_MAX || D > INT32_MAX)
{
std::ostringstream ss;
ss << "SoftmaxCPU inputs N, D and N * D must be < " << INT32_MAX << ". N=" << N << ", D=" << D;
std::string msg = ss.str();
return Status(common::ONNXRUNTIME, common::INVALID_ARGUMENT, msg);
}*/
}
} // namespace test
} // namespace onnxruntime