mirror of
https://github.com/saymrwulf/onnxruntime.git
synced 2026-07-07 17:15:29 +00:00
Revert BatchNormalization change for now, falling back to CPU on mixed types until a more advanced solution is written
This commit is contained in:
parent
babd6e3fcd
commit
fe7b8b80ae
2 changed files with 99 additions and 7 deletions
|
|
@ -21,6 +21,67 @@ class DmlOperatorBatchNormalization : public DmlOperator, BatchNormalizationHelp
|
|||
|
||||
public:
|
||||
DmlOperatorBatchNormalization(const MLOperatorKernelCreationContext& kernelCreationContext)
|
||||
: DmlOperator(kernelCreationContext),
|
||||
BatchNormalizationHelper(kernelCreationContext, kernelCreationContext.GetTensorShapeDescription())
|
||||
{
|
||||
std::vector<std::optional<uint32_t>> kernelInputIndices = {X, Mean, Variance, Scale, Bias};
|
||||
DmlOperator::Initialize(kernelCreationContext, kernelInputIndices);
|
||||
|
||||
ML_CHECK_VALID_ARGUMENT(m_inputTensorDescs.size() == 5);
|
||||
ML_CHECK_VALID_ARGUMENT(m_outputTensorDescs.size() >= 1);
|
||||
|
||||
const float epsilon = kernelCreationContext.GetOptionalAttribute<float>(AttrName::Epsilon, 0.0f);
|
||||
const int spatial = kernelCreationContext.GetOptionalAttribute<int>(AttrName::Spatial, 1);
|
||||
const std::optional<ActivationOperatorDesc> fusedActivation = FusionHelpers::TryGetFusedActivationDesc(kernelCreationContext);
|
||||
DML_OPERATOR_DESC fusedActivationDmlDesc = fusedActivation ? fusedActivation->GetDmlDesc() : DML_OPERATOR_DESC();
|
||||
|
||||
m_inputTensorDescs[0] = CreateTensorDescFromInput(kernelCreationContext, 0, TensorAxis::DoNotCoerce, TensorAxis::N, TensorAxis::LeftAligned);
|
||||
|
||||
// Massage each of these 1D tensors (of length C) into ND tensors of the form [1,C,1,1,...].
|
||||
for (uint32_t i = Scale; i < OnnxInputIndex::Count; ++i)
|
||||
{
|
||||
m_inputTensorDescs[i] = CreateTensorDescFromInput(kernelCreationContext, i, TensorAxis::DoNotCoerce, TensorAxis::C, TensorAxis::LeftAligned, std::nullopt, m_inputTensorDescs[0].GetDimensionCount());
|
||||
}
|
||||
|
||||
m_outputTensorDescs[0] = CreateTensorDescFromOutput(kernelCreationContext, 0, TensorAxis::DoNotCoerce, TensorAxis::N, TensorAxis::LeftAligned, std::nullopt, m_inputTensorDescs[0].GetDimensionCount());
|
||||
|
||||
ML_CHECK_VALID_ARGUMENT(m_inputTensorDescs.size() == 5);
|
||||
ML_CHECK_VALID_ARGUMENT(m_outputTensorDescs.size() >= 1);
|
||||
|
||||
std::vector<DML_TENSOR_DESC> inputDescs = GetDmlInputDescs();
|
||||
std::vector<DML_TENSOR_DESC> outputDescs = GetDmlOutputDescs();
|
||||
|
||||
DML_BATCH_NORMALIZATION_OPERATOR_DESC operatorDesc = {};
|
||||
operatorDesc.InputTensor = &inputDescs[X];
|
||||
operatorDesc.MeanTensor = &inputDescs[Mean];
|
||||
operatorDesc.VarianceTensor = &inputDescs[Variance];
|
||||
operatorDesc.ScaleTensor = &inputDescs[Scale];
|
||||
operatorDesc.BiasTensor = &inputDescs[Bias];
|
||||
operatorDesc.OutputTensor = &outputDescs[0];
|
||||
operatorDesc.Spatial = static_cast<BOOL>(spatial);
|
||||
operatorDesc.Epsilon = epsilon;
|
||||
operatorDesc.FusedActivation = fusedActivation ? &fusedActivationDmlDesc : nullptr;
|
||||
|
||||
DML_OPERATOR_DESC opDesc = { DML_OPERATOR_BATCH_NORMALIZATION, &operatorDesc };
|
||||
SetDmlOperatorDesc(opDesc, kernelCreationContext);
|
||||
}
|
||||
};
|
||||
|
||||
class DmlOperatorBatchNormalization15 : public DmlOperator, BatchNormalizationHelper
|
||||
{
|
||||
// This order matches the ONNX schema.
|
||||
enum OnnxInputIndex
|
||||
{
|
||||
X, // Input
|
||||
Scale,
|
||||
Bias,
|
||||
Mean,
|
||||
Variance,
|
||||
Count,
|
||||
};
|
||||
|
||||
public:
|
||||
DmlOperatorBatchNormalization15(const MLOperatorKernelCreationContext& kernelCreationContext)
|
||||
: DmlOperator(kernelCreationContext),
|
||||
BatchNormalizationHelper(kernelCreationContext, kernelCreationContext.GetTensorShapeDescription())
|
||||
{
|
||||
|
|
@ -101,15 +162,46 @@ public:
|
|||
|
||||
void CALLBACK QueryBatchNormalization(IMLOperatorSupportQueryContextPrivate* context, /*out*/ bool* isSupported)
|
||||
{
|
||||
// training_mode=1 is unsupported as it isn't needed for inference (https://github.com/onnx/onnx/pull/3333).
|
||||
*isSupported = false;
|
||||
|
||||
// training_mode=1 is unsupported as it isn't needed for inference (https://github.com/onnx/onnx/pull/3333).
|
||||
MLOperatorAttributes attributes(context);
|
||||
int32_t trainingMode = attributes.GetOptionalAttribute<int32_t>(AttrName::TrainingMode, 0);
|
||||
*isSupported = (trainingMode == 0);
|
||||
if (trainingMode != 0)
|
||||
{
|
||||
return;
|
||||
}
|
||||
|
||||
if (context->GetInputCount() < 5)
|
||||
{
|
||||
return;
|
||||
}
|
||||
|
||||
// Get the data type of each tensor.
|
||||
MLOperatorEdgeDescription operatorEdgeDescription[5];
|
||||
for (uint32_t i = 0; i < 5; ++i)
|
||||
{
|
||||
if (FAILED(context->GetInputEdgeDescription(i, &operatorEdgeDescription[i]))
|
||||
|| operatorEdgeDescription[i].edgeType != MLOperatorEdgeType::Tensor)
|
||||
{
|
||||
return;
|
||||
}
|
||||
}
|
||||
|
||||
// Fall back if the data types of the mean/variance or scale/bias differ from the input.
|
||||
MLOperatorTensorDataType inputTensorDataType = operatorEdgeDescription[0].tensorDataType;
|
||||
for (uint32_t i = 1; i < 5; ++i)
|
||||
{
|
||||
if (operatorEdgeDescription[i].tensorDataType != inputTensorDataType)
|
||||
{
|
||||
return;
|
||||
}
|
||||
}
|
||||
|
||||
*isSupported = true;
|
||||
}
|
||||
|
||||
DML_OP_DEFINE_CREATION_FUNCTION(BatchNormalization, DmlOperatorBatchNormalization);
|
||||
DML_OP_DEFINE_CREATION_FUNCTION(BatchNormalization15, DmlOperatorBatchNormalization);
|
||||
DML_OP_DEFINE_CREATION_FUNCTION(FusedBatchNormalization, DmlOperatorBatchNormalization);
|
||||
|
||||
} // namespace Dml
|
||||
|
|
|
|||
|
|
@ -397,10 +397,10 @@ constexpr static OperatorRegistrationInformation operatorRegistrationInformation
|
|||
{REG_INFO( 7, MaxRoiPool, typeNameListDefault, supportedTypeListFloat16to32, DmlGraphSupport::Supported)},
|
||||
{REG_INFO_VER( 10, RoiAlign, typeNameListTwo, supportedTypeListRoiAlign, DmlGraphSupport::Supported)},
|
||||
{REG_INFO( 7, InstanceNormalization, typeNameListDefault, supportedTypeListFloat16to32, DmlGraphSupport::Supported)},
|
||||
{REG_INFO( 7, BatchNormalization, typeNameListDefault, supportedTypeListFloat16to32, DmlGraphSupport::NotSupported)},
|
||||
{REG_INFO( 9, BatchNormalization, typeNameListDefault, supportedTypeListFloat16to32, DmlGraphSupport::NotSupported)}, // v9 just removes 'spatial' attribute.
|
||||
{REG_INFO( 14, BatchNormalization, typeNameListDefault, supportedTypeListFloat16to32, DmlGraphSupport::NotSupported, requiredConstantCpuInputs(), std::nullopt, QueryBatchNormalization)}, // v14 adds training_mode attribute
|
||||
{REG_INFO_VER( 15, BatchNormalization, typeNameListDefault, supportedTypeListFloat16to32, DmlGraphSupport::NotSupported, requiredConstantCpuInputs(), std::nullopt, QueryBatchNormalization)}, // v15 adds differing types for scale and bias vs input.
|
||||
{REG_INFO( 7, BatchNormalization, typeNameListDefault, supportedTypeListFloat16to32, DmlGraphSupport::Supported)},
|
||||
{REG_INFO( 9, BatchNormalization, typeNameListDefault, supportedTypeListFloat16to32, DmlGraphSupport::Supported)}, // v9 just removes 'spatial' attribute.
|
||||
{REG_INFO( 14, BatchNormalization, typeNameListDefault, supportedTypeListFloat16to32, DmlGraphSupport::Supported, requiredConstantCpuInputs(), std::nullopt, QueryBatchNormalization)}, // v14 adds training_mode attribute
|
||||
{REG_INFO( 15, BatchNormalization, typeNameListDefault, supportedTypeListFloat16to32, DmlGraphSupport::Supported, requiredConstantCpuInputs(), std::nullopt, QueryBatchNormalization)}, // v15 adds differing types for scale and bias vs input.
|
||||
{REG_INFO( 7, LRN, typeNameListDefault, supportedTypeListFloat16to32, DmlGraphSupport::Supported)},
|
||||
{REG_INFO( 13, LRN, typeNameListDefault, supportedTypeListFloat16to32, DmlGraphSupport::Supported)},
|
||||
{REG_INFO( 7, MeanVarianceNormalization, typeNameListDefault, supportedTypeListFloat16to32, DmlGraphSupport::Supported)},
|
||||
|
|
|
|||
Loading…
Reference in a new issue