| 136 | } |
| 137 | |
| 138 | Status FusePadAndConv(const GraphDef& input_graph_def, |
| 139 | const TransformFuncContext& context, |
| 140 | GraphDef* output_graph_def) { |
| 141 | GraphDef replaced_graph_def; |
| 142 | TF_RETURN_IF_ERROR(ReplaceMatchingOpTypes( |
| 143 | input_graph_def, // clang-format off |
| 144 | {"Conv2D", |
| 145 | { |
| 146 | {"MirrorPad", |
| 147 | { |
| 148 | {"*"}, |
| 149 | {"*"}, |
| 150 | } |
| 151 | }, |
| 152 | {"*"} |
| 153 | } |
| 154 | }, // clang-format on |
| 155 | [](const NodeMatch& match, const std::set<string>& input_nodes, |
| 156 | const std::set<string>& output_nodes, |
| 157 | std::vector<NodeDef>* new_nodes) { |
| 158 | // Find all the nodes we expect in the subgraph. |
| 159 | const NodeDef& conv_node = match.node; |
| 160 | CHECK_EQ("Conv2D", conv_node.op()); |
| 161 | const NodeDef& mirror_pad_node = match.inputs[0].node; |
| 162 | CHECK_EQ("MirrorPad", mirror_pad_node.op()); |
| 163 | const NodeDef& weights_node = match.inputs[1].node; |
| 164 | const NodeDef& input_node = match.inputs[0].inputs[0].node; |
| 165 | const NodeDef& pad_dims_node = match.inputs[0].inputs[1].node; |
| 166 | |
| 167 | // We'll be reusing the old weights and pad dimensions. |
| 168 | new_nodes->push_back(weights_node); |
| 169 | new_nodes->push_back(input_node); |
| 170 | new_nodes->push_back(pad_dims_node); |
| 171 | |
| 172 | // Set up the new fused version of the convolution op. |
| 173 | NodeDef fused_conv; |
| 174 | fused_conv.set_op("FusedPadConv2D"); |
| 175 | fused_conv.set_name(match.node.name()); |
| 176 | AddNodeInput(mirror_pad_node.input(0), &fused_conv); |
| 177 | AddNodeInput(mirror_pad_node.input(1), &fused_conv); |
| 178 | AddNodeInput(conv_node.input(1), &fused_conv); |
| 179 | CopyNodeAttr(mirror_pad_node, "mode", "mode", &fused_conv); |
| 180 | CopyNodeAttr(conv_node, "T", "T", &fused_conv); |
| 181 | CopyNodeAttr(conv_node, "padding", "padding", &fused_conv); |
| 182 | CopyNodeAttr(conv_node, "strides", "strides", &fused_conv); |
| 183 | new_nodes->push_back(fused_conv); |
| 184 | |
| 185 | return Status::OK(); |
| 186 | }, |
| 187 | {}, &replaced_graph_def)); |
| 188 | *output_graph_def = replaced_graph_def; |
| 189 | return Status::OK(); |
| 190 | } |
| 191 | |
| 192 | REGISTER_GRAPH_TRANSFORM("fuse_resize_pad_and_conv", FuseResizePadAndConv); |
| 193 | |