Revert BatchNormalization change for now, falling back to CPU on mixed types until a more advanced solution is written

This commit is contained in:
Dwayne Robinson 2022-06-15 21:49:18 -07:00
parent babd6e3fcd
commit fe7b8b80ae
2 changed files with 99 additions and 7 deletions

View file

@ -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

View file

@ -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)},