Converts Conv2D or MatMul ops followed by column-wise Muls into equivalent ops with the Mul baked into the convolution weights, to save computation during inference.
| 30 | // ops with the Mul baked into the convolution weights, to save computation |
| 31 | // during inference. |
| 32 | Status FoldBatchNorms(const GraphDef& input_graph_def, |
| 33 | const TransformFuncContext& context, |
| 34 | GraphDef* output_graph_def) { |
| 35 | GraphDef replaced_graph_def; |
| 36 | TF_RETURN_IF_ERROR(ReplaceMatchingOpTypes( |
| 37 | input_graph_def, // clang-format off |
| 38 | {"Mul", // mul_node |
| 39 | { |
| 40 | {"Conv2D|MatMul|DepthwiseConv2dNative", // conv_node |
| 41 | { |
| 42 | {"*"}, // input_node |
| 43 | {"Const"}, // weights_node |
| 44 | } |
| 45 | }, |
| 46 | {"Const"}, // mul_values_node |
| 47 | } |
| 48 | }, // clang-format on |
| 49 | [](const NodeMatch& match, const std::set<string>& input_nodes, |
| 50 | const std::set<string>& output_nodes, |
| 51 | std::vector<NodeDef>* new_nodes) { |
| 52 | // Find all the nodes we expect in the subgraph. |
| 53 | const NodeDef& mul_node = match.node; |
| 54 | const NodeDef& conv_node = match.inputs[0].node; |
| 55 | const NodeDef& input_node = match.inputs[0].inputs[0].node; |
| 56 | const NodeDef& weights_node = match.inputs[0].inputs[1].node; |
| 57 | const NodeDef& mul_values_node = match.inputs[1].node; |
| 58 | |
| 59 | // Check that nodes that we use are not used somewhere else. |
| 60 | for (const auto& node : {conv_node, weights_node, mul_values_node}) { |
| 61 | if (output_nodes.count(node.name())) { |
| 62 | // Return original nodes. |
| 63 | new_nodes->insert(new_nodes->end(), |
| 64 | {mul_node, conv_node, input_node, weights_node, |
| 65 | mul_values_node}); |
| 66 | return Status::OK(); |
| 67 | } |
| 68 | } |
| 69 | |
| 70 | Tensor weights = GetNodeTensorAttr(weights_node, "value"); |
| 71 | Tensor mul_values = GetNodeTensorAttr(mul_values_node, "value"); |
| 72 | |
| 73 | // Make sure all the inputs really are vectors, with as many entries as |
| 74 | // there are columns in the weights. |
| 75 | int64 weights_cols; |
| 76 | if (conv_node.op() == "Conv2D") { |
| 77 | weights_cols = weights.shape().dim_size(3); |
| 78 | } else if (conv_node.op() == "DepthwiseConv2dNative") { |
| 79 | weights_cols = |
| 80 | weights.shape().dim_size(2) * weights.shape().dim_size(3); |
| 81 | } else { |
| 82 | weights_cols = weights.shape().dim_size(1); |
| 83 | } |
| 84 | if ((mul_values.shape().dims() != 1) || |
| 85 | (mul_values.shape().dim_size(0) != weights_cols)) { |
| 86 | return errors::InvalidArgument( |
| 87 | "Mul constant input to batch norm has bad shape: ", |
| 88 | mul_values.shape().DebugString()); |
| 89 | } |