| 24 | namespace graph_transforms { |
| 25 | |
| 26 | Status FlattenAtrousConv(const GraphDef& input_graph_def, |
| 27 | const TransformFuncContext& context, |
| 28 | GraphDef* output_graph_def) { |
| 29 | GraphDef replaced_graph_def; |
| 30 | TF_RETURN_IF_ERROR(ReplaceMatchingOpTypes( |
| 31 | input_graph_def, // clang-format off |
| 32 | {"BatchToSpaceND", |
| 33 | { |
| 34 | {"Conv2D|DepthwiseConv2dNative", |
| 35 | { |
| 36 | {"SpaceToBatchND", |
| 37 | { |
| 38 | {"*"}, // Input to the flattened op. |
| 39 | {"*"}, // block_shape |
| 40 | {"*"} // paddings |
| 41 | } |
| 42 | }, |
| 43 | {"*"} // filter |
| 44 | } |
| 45 | }, |
| 46 | {"*"}, // block_shape |
| 47 | {"*"} // crops |
| 48 | } |
| 49 | }, // clang-format on |
| 50 | [](const NodeMatch& match, const std::set<string>& input_nodes, |
| 51 | const std::set<string>& output_nodes, |
| 52 | std::vector<NodeDef>* new_nodes) { |
| 53 | // Find all the nodes we expect in the subgraph. |
| 54 | const NodeDef& batch_to_space_node = match.node; |
| 55 | const NodeDef& conv_node = match.inputs[0].node; |
| 56 | const NodeDef& filter_node = match.inputs[0].inputs[1].node; |
| 57 | const NodeDef& input_node = match.inputs[0].inputs[0].inputs[0].node; |
| 58 | const NodeDef& space_to_batch_block_shape_node = |
| 59 | match.inputs[0].inputs[0].inputs[1].node; |
| 60 | |
| 61 | // The atrous rate value is inferred from the block shape. |
| 62 | Tensor block_shape = |
| 63 | GetNodeTensorAttr(space_to_batch_block_shape_node, "value"); |
| 64 | const int32 block_height = block_shape.flat<int32>()(0); |
| 65 | const int32 block_width = block_shape.flat<int32>()(1); |
| 66 | |
| 67 | // Compute the upsampled filter. |
| 68 | const Tensor& filter = GetNodeTensorAttr(filter_node, "value"); |
| 69 | const int32 filter_height = filter.dim_size(0); |
| 70 | const int32 filter_width = filter.dim_size(1); |
| 71 | const int32 in_channels = filter.dim_size(2); |
| 72 | const int32 out_channels = filter.dim_size(3); |
| 73 | |
| 74 | const int32 upsampled_filter_height = |
| 75 | (filter_height - 1) * block_height + 1; |
| 76 | const int32 upsampled_filter_width = |
| 77 | (filter_width - 1) * block_width + 1; |
| 78 | Tensor upsampled_filter( |
| 79 | DT_FLOAT, |
| 80 | TensorShape({upsampled_filter_height, upsampled_filter_width, |
| 81 | in_channels, out_channels})); |
| 82 | |
| 83 | auto filter_eigen = filter.tensor<float, 4>(); |