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

Batchnorm training mode support in a minimal build (#17103)

Browse source
This commit is contained in:
Baiju Meswani 2023-08-10 21:44:50 -07:00 committed by GitHub
parent 344c41fdb9
commit 54153c73f0
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
2 changed files with 7 additions and 12 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

14
onnxruntime/core/providers/cpu/nn/batch_norm.h
View file

@ -29,10 +29,6 @@
namespace onnxruntime {
#if !defined(ORT_MINIMAL_BUILD)
#define BATCHNORM_INCLUDE_TRAINING_SUPPORT
#endif
template <typename T>
class BatchNorm : public OpKernel {
public:
@ -51,7 +47,7 @@ class BatchNorm : public OpKernel {
}
if (is_train_) {
#if defined(BATCHNORM_INCLUDE_TRAINING_SUPPORT)
#ifdef ENABLE_TRAINING_OPS
momentum_ = op_kernel_info.GetAttrOrDefault<float>("momentum", 0.9f);
ORT_ENFORCE(is_spatial_, "Training mode only supports spatial BN");
#else
@ -88,7 +84,7 @@ class BatchNorm : public OpKernel {
// calculate sample_size (including all channels)
size_t sample_size_incl_all_channels = sample_size * C;
#if defined(BATCHNORM_INCLUDE_TRAINING_SUPPORT)
#ifdef ENABLE_TRAINING_OPS
AllocatorPtr alloc;
ORT_RETURN_IF_ERROR(p_op_kernel_context->GetTempSpaceAllocator(&alloc));
@ -115,7 +111,7 @@ class BatchNorm : public OpKernel {
ConstEigenVectorArrayMap<T> scale_arr(scale->Data<T>(), is_spatial_ ? C : sample_size_incl_all_channels);
ConstEigenVectorArrayMap<T> bias_arr(B->Data<T>(), is_spatial_ ? C : sample_size_incl_all_channels);
#if defined(BATCHNORM_INCLUDE_TRAINING_SUPPORT)
#ifdef ENABLE_TRAINING_OPS
// Note that we only support spatial BN for training
if (is_train_) {
EigenVectorArrayMap<T> saved_mean_arr(saved_mean->MutableData<T>(), C);
@ -166,7 +162,7 @@ class BatchNorm : public OpKernel {
ConstEigenVectorArrayMap<T> var_arr(var->Data<T>(), is_spatial_ ? C : sample_size_incl_all_channels);
inv_std = (var_arr + epsilon_).sqrt().inverse();
} else {
#if defined(BATCHNORM_INCLUDE_TRAINING_SUPPORT)
#ifdef ENABLE_TRAINING_OPS
EigenVectorArrayMap<T> saved_inv_std_arr(saved_inv_std->MutableData<T>(), C);
saved_inv_std_arr = (saved_inv_std_arr + epsilon_).inverse().sqrt();
inv_std = saved_inv_std_arr;
@ -175,7 +171,7 @@ class BatchNorm : public OpKernel {
// If we're training, do batch normalization based on computation from this batch
ConstEigenVectorArrayMap<T> mean_arr(
#if defined(BATCHNORM_INCLUDE_TRAINING_SUPPORT)
#ifdef ENABLE_TRAINING_OPS
!is_train_ ? mean->Data<T>() : saved_mean->Data<T>(),
#else
mean->Data<T>(),

5
onnxruntime/test/providers/cpu/nn/batch_norm_op_test.cc
View file

@ -2,7 +2,6 @@
// Licensed under the MIT License.
#include "core/framework/tensor.h"
#include "core/providers/cpu/nn/batch_norm.h" // for BATCHNORM_INCLUDE_TRAINING_SUPPORT
#include "core/session/inference_session.h"
#include "test/common/dnnl_op_test_utils.h"
#include "test/providers/provider_test_utils.h"
@ -847,7 +846,7 @@ TEST(BatchNormTest, BatchNorm2d_bfloat16) {
#endif // USE_DNNL
// TODO fix flaky test for CUDA
#ifdef BATCHNORM_INCLUDE_TRAINING_SUPPORT
#ifdef ENABLE_TRAINING_OPS
TEST(BatchNormTest, ForwardTrainingTestWithSavedOutputsOpset9) {
// TODO: Unskip when fixed #41968513
if (DefaultDmlExecutionProvider().get() != nullptr) {
@ -937,7 +936,7 @@ TEST(BatchNormTest, ForwardTrainingTestOpset15) {
{kCudaExecutionProvider, kRocmExecutionProvider,
kTensorrtExecutionProvider, kOpenVINOExecutionProvider, kDnnlExecutionProvider});
}
#endif // BATCHNORM_INCLUDE_TRAINING_SUPPORT
#endif // ENABLE_TRAINING_OPS
} // namespace test
} // namespace onnxruntime