Populates a layer match object containing ops/tensors for folding BNs. Args: match_result: Matched result from graph matcher Returns: layer_op: Matching conv/fc op prior to batch norm BatchNormMatch: _BatchNormMatch containing all required batch norm parameters.
(match_result)
| 221 | moving_average_mul_pattern) |
| 222 | |
| 223 | def _GetLayerMatch(match_result): |
| 224 | """Populates a layer match object containing ops/tensors for folding BNs. |
| 225 | |
| 226 | Args: |
| 227 | match_result: Matched result from graph matcher |
| 228 | |
| 229 | Returns: |
| 230 | layer_op: Matching conv/fc op prior to batch norm |
| 231 | BatchNormMatch: _BatchNormMatch containing all required batch norm |
| 232 | parameters. |
| 233 | """ |
| 234 | moving_mean_tensor = None |
| 235 | moving_variance_tensor = None |
| 236 | bn_decay_mean_tensor = None |
| 237 | bn_decay_var_tensor = None |
| 238 | batch_to_space_op = None |
| 239 | layer_op = match_result.get_op(layer_pattern) |
| 240 | layer_tensor = match_result.get_tensor(layer_pattern) |
| 241 | bn_id_op = match_result.get_op(batch_norm_identity_pattern) |
| 242 | bn_op = match_result.get_op(batch_norm_pattern) |
| 243 | if bn_id_op is None: |
| 244 | bn_id_op = bn_op |
| 245 | |
| 246 | batch_epsilon = bn_op.get_attr('epsilon') |
| 247 | |
| 248 | # In the MatMul case, the output of batch norm is reshaped back into a |
| 249 | # 2D tensor, so the output_tensor is the output of the Reshape op. |
| 250 | output_tensor = bn_op.outputs[0] |
| 251 | if layer_op.type == 'MatMul': |
| 252 | output_reshape_op = match_result.get_op(matmul_bn_output_reshape_pattern) |
| 253 | # If the matcher didn't match matmul_bn_output_reshape, there will be |
| 254 | # another match for this 'MatMul' later, so we can skip this one. |
| 255 | if output_reshape_op is None: |
| 256 | return None, None |
| 257 | output_tensor = output_reshape_op.outputs[0] |
| 258 | |
| 259 | # Ensure that the output tensor has consumers, otherwise this is a dangling |
| 260 | # node and not a match. |
| 261 | if not output_tensor.consumers(): |
| 262 | return None, None |
| 263 | |
| 264 | batch_to_space_op = match_result.get_op(batch_to_space_pattern) |
| 265 | input_tensor = match_result.get_tensor(input_pattern) |
| 266 | weight_tensor = match_result.get_tensor(weight_pattern) |
| 267 | gamma_tensor = match_result.get_tensor(gamma_pattern) |
| 268 | beta_tensor = match_result.get_tensor(beta_pattern) |
| 269 | # FusedBatchNorm in training is different from that in inference. It takes |
| 270 | # empty 'mean' and empty 'variance', and produces the mean and the variance |
| 271 | # of the batch. Therefore, when is_training is true, mean_tensor and |
| 272 | # variance_tensor point to 1st and 2nd (0-based) output of bn_op, |
| 273 | # respectively; when is_training is false, they point to bn_op's inputs. |
| 274 | is_training = bn_op.get_attr('is_training') |
| 275 | if is_training: |
| 276 | # FusedBatchNormGrad doesn't compute gradients of the batch_mean and |
| 277 | # batch_variance outputs, so we need to substitute our own custom |
| 278 | # gradient. |
| 279 | # TODO(suharshs, raghuramank): Find a way to avoid needing this hack. |
| 280 | # pylint: disable=protected-access |
no test coverage detected