| 1757 | } |
| 1758 | |
| 1759 | Status AddBatchNormNodes(RemapperContext* ctx, const FusedBatchNorm& matched) { |
| 1760 | const GraphDef* graph = ctx->graph_view.graph(); |
| 1761 | const NodeDef& fused_node = graph->node(matched.fused_batch_norm); |
| 1762 | VLOG(2) << "Optimizing fused batch norm node " |
| 1763 | << SummarizeNodeDef(fused_node); |
| 1764 | |
| 1765 | const string& x = fused_node.input(0); |
| 1766 | string scale = fused_node.input(1); |
| 1767 | string offset = fused_node.input(2); |
| 1768 | string mean = fused_node.input(3); |
| 1769 | string variance = fused_node.input(4); |
| 1770 | |
| 1771 | utils::Mutation* mutation = ctx->graph_view.GetMutationBuilder(); |
| 1772 | Status status; |
| 1773 | |
| 1774 | string x_format = fused_node.attr().at(kDataFormat).s(); |
| 1775 | if (x_format == "NCHW" || x_format == "NCDHW") { |
| 1776 | // Need to reshape the last 4 inputs |
| 1777 | NodeDef new_shape; |
| 1778 | const string new_shape_name = |
| 1779 | AddPrefixToNodeName(x_format + "Shape", fused_node.name()); |
| 1780 | new_shape.set_name(new_shape_name); |
| 1781 | new_shape.set_op("Const"); |
| 1782 | new_shape.set_device(fused_node.device()); |
| 1783 | *new_shape.add_input() = AsControlDependency(scale); |
| 1784 | (*new_shape.mutable_attr())["dtype"].set_type(DT_INT32); |
| 1785 | if (x_format == "NCHW") { |
| 1786 | Tensor t(DT_INT32, {4}); |
| 1787 | t.flat<int32>()(0) = 1; |
| 1788 | t.flat<int32>()(1) = -1; |
| 1789 | t.flat<int32>()(2) = 1; |
| 1790 | t.flat<int32>()(3) = 1; |
| 1791 | t.AsProtoTensorContent( |
| 1792 | (*new_shape.mutable_attr())["value"].mutable_tensor()); |
| 1793 | } else { |
| 1794 | Tensor t(DT_INT32, {5}); |
| 1795 | t.flat<int32>()(0) = 1; |
| 1796 | t.flat<int32>()(1) = -1; |
| 1797 | t.flat<int32>()(2) = 1; |
| 1798 | t.flat<int32>()(3) = 1; |
| 1799 | t.flat<int32>()(4) = 1; |
| 1800 | t.AsProtoTensorContent( |
| 1801 | (*new_shape.mutable_attr())["value"].mutable_tensor()); |
| 1802 | } |
| 1803 | mutation->AddNode(std::move(new_shape), &status); |
| 1804 | TF_RETURN_IF_ERROR(status); |
| 1805 | |
| 1806 | NodeDef reshaped_scale; |
| 1807 | reshaped_scale.set_name( |
| 1808 | AddPrefixToNodeName(x_format + "ShapedScale", fused_node.name())); |
| 1809 | reshaped_scale.set_op("Reshape"); |
| 1810 | reshaped_scale.set_device(fused_node.device()); |
| 1811 | *reshaped_scale.add_input() = scale; |
| 1812 | *reshaped_scale.add_input() = new_shape_name; |
| 1813 | (*reshaped_scale.mutable_attr())["T"] = fused_node.attr().at("T"); |
| 1814 | (*reshaped_scale.mutable_attr())["Tshape"].set_type(DT_INT32); |
| 1815 | scale = reshaped_scale.name(); |
| 1816 | mutation->AddNode(std::move(reshaped_scale), &status); |
no test coverage detected