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

Introduce --enable_training_ops build flag (#6523)

Browse source 
* minimal_build with training ops

* Removing redundant comment from an earlier attempt at a fix

* Fixing a bad merge conflict resolution

* Responding to PR feedback

* tweaking the makefiles based on feedback

* combining two enable_training blocks in CMakeLists.txt
This commit is contained in:
Cian Hayes 2021-02-01 21:54:16 -08:00 committed by GitHub
parent 9a6e71574a
commit 6fc5237d9e
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
7 changed files with 47 additions and 9 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

6
cmake/CMakeLists.txt
View file

@ -142,6 +142,7 @@ option(onnxruntime_FUZZ_TEST "Enable Fuzz testing" OFF)
option(onnxruntime_ENABLE_NVTX_PROFILE "Enable NVTX profile." OFF)
option(onnxruntime_ENABLE_MEMORY_PROFILE "Enable memory profile." OFF)
option(onnxruntime_ENABLE_TRAINING "Enable training functionality." OFF)
option(onnxruntime_ENABLE_TRAINING_OPS "Include training operators but no training session support." OFF)
option(onnxruntime_ENABLE_TRAINING_E2E_TESTS "Enable training end-to-end tests." OFF)
option(onnxruntime_ENABLE_CPU_FP16_OPS "Build with advanced instruction sets" ON)
option(onnxruntime_USE_NCCL "Build with NCCL support" OFF)
@ -1238,8 +1239,13 @@ if (onnxruntime_USE_DML)
include(dml)
endif()
if (onnxruntime_ENABLE_TRAINING_OPS)
add_compile_definitions(ENABLE_TRAINING_OPS)
endif()
if (onnxruntime_ENABLE_TRAINING)
add_compile_definitions(ENABLE_TRAINING)
add_compile_definitions(ENABLE_TRAINING_OPS)
if (UNIX)
# Find MPI

30
cmake/onnxruntime_providers.cmake
View file

@ -129,6 +129,25 @@ if (onnxruntime_USE_FEATURIZERS)
list(APPEND onnxruntime_providers_src ${onnxruntime_cpu_featurizers_cc_srcs})
endif()
if (onnxruntime_ENABLE_TRAINING_OPS)
file(GLOB_RECURSE onnxruntime_cpu_training_ops_srcs CONFIGURE_DEPENDS
"${ORTTRAINING_SOURCE_DIR}/training_ops/cpu/*.h"
"${ORTTRAINING_SOURCE_DIR}/training_ops/cpu/*.cc"
)
source_group(TREE ${ORTTRAINING_ROOT}/ FILES ${onnxruntime_cpu_training_ops_srcs})
list(APPEND onnxruntime_providers_src ${onnxruntime_cpu_training_ops_srcs})
file(GLOB_RECURSE onnxruntime_cpu_full_training_only_srcs
"${ORTTRAINING_SOURCE_DIR}/training_ops/cpu/gist/*.h"
"${ORTTRAINING_SOURCE_DIR}/training_ops/cpu/gist/*.cc"
"${ORTTRAINING_SOURCE_DIR}/training_ops/cpu/tensorboard/*.h"
"${ORTTRAINING_SOURCE_DIR}/training_ops/cpu/tensorboard/*.cc"
)
list(REMOVE_ITEM onnxruntime_providers_src ${onnxruntime_cpu_full_training_only_srcs})
endif()
if (onnxruntime_ENABLE_TRAINING)
file(GLOB_RECURSE onnxruntime_cpu_training_ops_srcs CONFIGURE_DEPENDS
"${ORTTRAINING_SOURCE_DIR}/training_ops/cpu/*.h"
@ -174,8 +193,11 @@ target_include_directories(onnxruntime_providers PRIVATE ${ONNXRUNTIME_ROOT} ${e
add_dependencies(onnxruntime_providers onnx ${onnxruntime_EXTERNAL_DEPENDENCIES})
if (onnxruntime_ENABLE_TRAINING)
if (onnxruntime_ENABLE_TRAINING OR onnxruntime_ENABLE_TRAINING_OPS)
target_include_directories(onnxruntime_providers PRIVATE ${ORTTRAINING_ROOT})
endif()
if (onnxruntime_ENABLE_TRAINING)
add_dependencies(onnxruntime_providers tensorboard)
onnxruntime_add_include_to_target(onnxruntime_providers tensorboard)
@ -208,7 +230,7 @@ if (onnxruntime_USE_CUDA)
list(APPEND onnxruntime_providers_cuda_src ${onnxruntime_cuda_contrib_ops_cc_srcs} ${onnxruntime_cuda_contrib_ops_cu_srcs})
endif()
if (onnxruntime_ENABLE_TRAINING)
if (onnxruntime_ENABLE_TRAINING OR onnxruntime_ENABLE_TRAINING_OPS)
file(GLOB_RECURSE onnxruntime_cuda_training_ops_cc_srcs CONFIGURE_DEPENDS
"${ORTTRAINING_SOURCE_DIR}/training_ops/cuda/*.h"
"${ORTTRAINING_SOURCE_DIR}/training_ops/cuda/*.cc"
@ -253,7 +275,7 @@ if (onnxruntime_USE_CUDA)
target_compile_options(onnxruntime_providers_cuda PRIVATE "$<$<COMPILE_LANGUAGE:CUDA>:SHELL:-Xcompiler /wd4127>")
endif()
onnxruntime_add_include_to_target(onnxruntime_providers_cuda onnxruntime_common onnxruntime_framework onnx onnx_proto protobuf::libprotobuf flatbuffers)
if (onnxruntime_ENABLE_TRAINING)
if (onnxruntime_ENABLE_TRAINING OR onnxruntime_ENABLE_TRAINING_OPS)
onnxruntime_add_include_to_target(onnxruntime_providers_cuda onnxruntime_training)
target_link_libraries(onnxruntime_providers_cuda PRIVATE onnxruntime_training)
endif()
@ -267,7 +289,7 @@ if (onnxruntime_USE_CUDA)
target_include_directories(onnxruntime_providers_cuda PRIVATE ${PROJECT_SOURCE_DIR}/external/cub)
endif()
if (onnxruntime_ENABLE_TRAINING)
if (onnxruntime_ENABLE_TRAINING OR onnxruntime_ENABLE_TRAINING_OPS)
target_include_directories(onnxruntime_providers_cuda PRIVATE ${ORTTRAINING_ROOT} ${MPI_INCLUDE_DIRS})
if (onnxruntime_USE_NCCL)

4
onnxruntime/core/providers/cpu/cpu_execution_provider.cc
View file

@ -13,7 +13,7 @@
#include "featurizers_ops/cpu/cpu_featurizers_kernels.h"
#endif
#ifdef ENABLE_TRAINING
#if defined(ENABLE_TRAINING) || defined(ENABLE_TRAINING_OPS)
#include "orttraining/training_ops/cpu/cpu_training_kernels.h"
#endif
@ -1840,7 +1840,7 @@ static Status RegisterCPUKernels(KernelRegistry& kernel_registry) {
#ifdef ML_FEATURIZERS
ORT_RETURN_IF_ERROR(::onnxruntime::featurizers::RegisterCpuMSFeaturizersKernels(kernel_registry));
#endif
#ifdef ENABLE_TRAINING
#if defined(ENABLE_TRAINING) || defined(ENABLE_TRAINING_OPS)
ORT_RETURN_IF_ERROR(::onnxruntime::contrib::RegisterCpuTrainingKernels(kernel_registry));
#endif
return Status::OK();

4
onnxruntime/core/providers/cpu/tensor/scatter.cc
View file

@ -5,7 +5,7 @@
#include "core/common/common.h"
#include "core/framework/op_kernel.h"
#include "core/providers/common.h"
#ifdef ENABLE_TRAINING
#if defined(ENABLE_TRAINING) || defined(ENABLE_TRAINING_OPS)
#include "orttraining/training_ops/cpu/tensor/gather_elements_grad_impl.h"
#endif
@ -255,7 +255,7 @@ Status Scatter::Compute(OpKernelContext* context) const {
return status;
}
#ifdef ENABLE_TRAINING
#if defined(ENABLE_TRAINING) || defined(ENABLE_TRAINING_OPS)
namespace contrib {

6
orttraining/orttraining/training_ops/cpu/cpu_training_kernels.cc
View file

@ -58,12 +58,14 @@ class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1,
class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, double_float, DropoutGrad);
class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, double_double, DropoutGrad);
#ifdef ENABLE_TRAINING
class ONNX_OPERATOR_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, SummaryScalar);
class ONNX_OPERATOR_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, SummaryHistogram);
class ONNX_OPERATOR_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, SummaryMerge);
class ONNX_OPERATOR_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, SummaryText);
class ONNX_OPERATOR_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, GistBinarizeEncoder);
class ONNX_OPERATOR_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, GistBinarizeDecoder);
#endif
class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, float, LayerNormalizationGrad);
class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, double, LayerNormalizationGrad);
class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, float, SimplifiedLayerNormalizationGrad);
@ -147,18 +149,22 @@ Status RegisterCpuTrainingKernels(KernelRegistry& kernel_registry) {
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, double_float, DropoutGrad)>,
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, double_double, DropoutGrad)>,
#ifdef ENABLE_TRAINING
BuildKernelCreateInfo<ONNX_OPERATOR_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, SummaryScalar)>,
BuildKernelCreateInfo<ONNX_OPERATOR_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, SummaryHistogram)>,
BuildKernelCreateInfo<ONNX_OPERATOR_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, SummaryMerge)>,
BuildKernelCreateInfo<ONNX_OPERATOR_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, SummaryText)>,
#endif
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, float, LayerNormalizationGrad)>,
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, double, LayerNormalizationGrad)>,
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, float, SimplifiedLayerNormalizationGrad)>,
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, double, SimplifiedLayerNormalizationGrad)>,
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, float, InvertibleLayerNormalizationGrad)>,
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, double, InvertibleLayerNormalizationGrad)>,
#ifdef ENABLE_TRAINING
BuildKernelCreateInfo<ONNX_OPERATOR_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, GistBinarizeEncoder)>,
BuildKernelCreateInfo<ONNX_OPERATOR_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, GistBinarizeDecoder)>,
#endif
BuildKernelCreateInfo<ONNX_OPERATOR_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, SliceGrad)>,
BuildKernelCreateInfo<ONNX_OPERATOR_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, FastGeluGrad)>,
BuildKernelCreateInfo<ONNX_OPERATOR_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, BiasGeluGrad_dX)>,

2
orttraining/orttraining/training_ops/cpu/nn/dropout_7.cc
View file

@ -6,7 +6,7 @@
namespace onnxruntime {
#ifdef ENABLE_TRAINING
#if defined(ENABLE_TRAINING) || defined(ENABLE_TRAINING_OPS)
ONNX_CPU_OPERATOR_KERNEL(
Dropout,
7,

4
tools/ci_build/build.py
View file

@ -159,6 +159,8 @@ def parse_arguments():
"--enable_memory_profile", action='store_true', help="Enable memory profile in ORT.")
parser.add_argument(
"--enable_training", action='store_true', help="Enable training in ORT.")
parser.add_argument(
"--enable_training_ops", action='store_true', help="Enable training ops in inference graph.")
parser.add_argument(
"--disable_nccl", action='store_true', help="Disable Nccl.")
parser.add_argument(
@ -711,6 +713,8 @@ def generate_build_tree(cmake_path, source_dir, build_dir, cuda_home, cudnn_home
"ON" if args.enable_nvtx_profile else "OFF"),
"-Donnxruntime_ENABLE_TRAINING=" + (
"ON" if args.enable_training else "OFF"),
"-Donnxruntime_ENABLE_TRAINING_OPS=" + (
"ON" if args.enable_training_ops else "OFF"),
# Enable advanced computations such as AVX for some traininig related ops.
"-Donnxruntime_ENABLE_CPU_FP16_OPS=" + (
"ON" if args.enable_training else "OFF"),