60 lines
2.6 KiB
Diff
60 lines
2.6 KiB
Diff
|
From 6972f9dfe325636b3db4e0bc517ee22a159365c0 Mon Sep 17 00:00:00 2001
|
||
|
|
From: Mihai Maruseac <mihaimaruseac@google.com>
|
||
|
|
Date: Thu, 6 May 2021 17:45:51 -0700
|
||
|
|
Subject: [PATCH] Add missing valuidation to FusedBatchNorm.
|
||
|
|
|
||
|
|
---
|
||
|
|
.../core/kernels/fused_batch_norm_op.cc | 29 ++++++++++++++++++-
|
||
|
|
1 file changed, 28 insertions(+), 1 deletion(-)
|
||
|
|
|
||
|
|
diff --git a/tensorflow/core/kernels/fused_batch_norm_op.cc b/tensorflow/core/kernels/fused_batch_norm_op.cc
|
||
|
|
index 59470c8a..bd5dab36 100644
|
||
|
|
--- a/tensorflow/core/kernels/fused_batch_norm_op.cc
|
||
|
|
+++ b/tensorflow/core/kernels/fused_batch_norm_op.cc
|
||
|
|
@@ -1267,6 +1267,33 @@ class FusedBatchNormOpBase : public OpKernel {
|
||
|
|
context, estimated_variance.dims() == 1,
|
||
|
|
errors::InvalidArgument("estimated_variance must be 1-dimensional",
|
||
|
|
estimated_variance.shape().DebugString()));
|
||
|
|
+
|
||
|
|
+ const auto num_channels = GetTensorDim(x, tensor_format_, 'C');
|
||
|
|
+ OP_REQUIRES(
|
||
|
|
+ context, scale.NumElements() == num_channels,
|
||
|
|
+ errors::InvalidArgument("scale must have the same number of elements "
|
||
|
|
+ "as the channels of x, got ",
|
||
|
|
+ scale.NumElements(), " and ", num_channels));
|
||
|
|
+ OP_REQUIRES(
|
||
|
|
+ context, offset.NumElements() == num_channels,
|
||
|
|
+ errors::InvalidArgument("offset must have the same number of elements "
|
||
|
|
+ "as the channels of x, got ",
|
||
|
|
+ offset.NumElements(), " and ", num_channels));
|
||
|
|
+ if (estimated_mean.NumElements() != 0) {
|
||
|
|
+ OP_REQUIRES(context, estimated_mean.NumElements() == num_channels,
|
||
|
|
+ errors::InvalidArgument(
|
||
|
|
+ "mean must be empty or have the same number of "
|
||
|
|
+ "elements as the channels of x, got ",
|
||
|
|
+ estimated_mean.NumElements(), " and ",num_channels));
|
||
|
|
+ }
|
||
|
|
+ if (estimated_variance.NumElements() != 0) {
|
||
|
|
+ OP_REQUIRES(context, estimated_variance.NumElements() == num_channels,
|
||
|
|
+ errors::InvalidArgument(
|
||
|
|
+ "variance must be empty or have the same number of "
|
||
|
|
+ "elements as the channels of x, got ",
|
||
|
|
+ estimated_variance.NumElements(), " and ", num_channels));
|
||
|
|
+ }
|
||
|
|
+
|
||
|
|
if (has_side_input_) {
|
||
|
|
OP_REQUIRES(context, side_input->shape() == x.shape(),
|
||
|
|
errors::InvalidArgument(
|
||
|
|
@@ -1279,7 +1306,7 @@ class FusedBatchNormOpBase : public OpKernel {
|
||
|
|
// NOTE(ezhulenev): This requirement is coming from implementation
|
||
|
|
// details of cudnnBatchNormalizationForwardTrainingEx.
|
||
|
|
OP_REQUIRES(
|
||
|
|
- context, !is_training_ || x.dim_size(3) % 4 == 0,
|
||
|
|
+ context, !is_training_ || num_channels % 4 == 0,
|
||
|
|
errors::InvalidArgument("FusedBatchNorm with activation requires "
|
||
|
|
"channel dimension to be a multiple of 4."));
|
||
|
|
}
|
||
|
|
--
|
||
|
|
2.23.0
|
||
|
|
|