(pfor_input)
| 1567 | # each of those iterations. |
| 1568 | @RegisterPFor("FusedBatchNormV3") |
| 1569 | def _convert_fused_batch_norm(pfor_input): |
| 1570 | is_training = pfor_input.get_attr("is_training") |
| 1571 | # When BatchNorm is used with training=False, mean and variance are provided |
| 1572 | # externally and used as is by the op. Thus, we can merge the S and N |
| 1573 | # dimensions as we do for regular operations. |
| 1574 | # When BatchNorm is used with training=True, mean and variance are computed |
| 1575 | # for each channel across the batch dimension (first one). If we merge S and N |
| 1576 | # dimensions, mean and variances will be computed over a larger set. So, we |
| 1577 | # merge the S and C dimensions instead. |
| 1578 | if not is_training: |
| 1579 | # We return zeros for batch_mean and batch_variance output. Note that CPU |
| 1580 | # and GPU seem to have different behavior for those two outputs. CPU outputs |
| 1581 | # zero because these values are not used during inference. GPU outputs |
| 1582 | # something, probably real means and variances. |
| 1583 | inputs = _inputs_with_flattening(pfor_input, [0]) |
| 1584 | outputs = _create_op( |
| 1585 | pfor_input.op_type, |
| 1586 | inputs, [x.dtype for x in pfor_input.outputs], |
| 1587 | attrs=pfor_input.op.node_def.attr).outputs |
| 1588 | y = outputs[0] |
| 1589 | n = pfor_input.pfor.loop_len_vector |
| 1590 | y = _unflatten_first_dim(y, n) |
| 1591 | mean = pfor_input.unstacked_input(3) |
| 1592 | zeros = array_ops.zeros_like(mean) |
| 1593 | return [wrap(y, True)] + [wrap(zeros, False)] * 5 |
| 1594 | |
| 1595 | pfor_input.stack_inputs() |
| 1596 | data_format = pfor_input.get_attr("data_format") |
| 1597 | # We merge the first dimension with the "C" dimension, run FusedBatchNormV3, |
| 1598 | # and then transpose back. |
| 1599 | x = pfor_input.stacked_input(0) |
| 1600 | x, reverse_order, reverse_shape = _channel_flatten_input(x, data_format) |
| 1601 | # Note that we stack all the other inputs as well so that they are the same |
| 1602 | # size as the new size of the channel dimension. |
| 1603 | inputs = [x] + [ |
| 1604 | array_ops.reshape(pfor_input.stacked_input(i), [-1]) |
| 1605 | for i in range(1, pfor_input.num_inputs) |
| 1606 | ] |
| 1607 | outputs = _create_op( |
| 1608 | pfor_input.op_type, |
| 1609 | inputs, [x.dtype for x in pfor_input.outputs], |
| 1610 | attrs=pfor_input.op.node_def.attr).outputs |
| 1611 | y = outputs[0] |
| 1612 | y = array_ops.reshape(y, reverse_shape) |
| 1613 | y = array_ops.transpose(y, reverse_order) |
| 1614 | n = pfor_input.pfor.loop_len_vector |
| 1615 | outputs = [_unflatten_first_dim(x, n) for x in outputs[1:]] |
| 1616 | outputs = [y] + outputs |
| 1617 | return [wrap(x, True) for x in outputs] |
| 1618 | |
| 1619 | |
| 1620 | @RegisterPFor("FusedBatchNormGradV3") |
nothing calls this directly
no test coverage detected