| 151 | } |
| 152 | |
| 153 | ::tensorflow::Status IdentifyDilatedConv::Run(Model* model, |
| 154 | std::size_t op_index, |
| 155 | bool* modified) { |
| 156 | *modified = false; |
| 157 | const auto it = model->operators.begin() + op_index; |
| 158 | auto* stb_op = it->get(); |
| 159 | |
| 160 | // 1. IDENTIFY OPERATORS |
| 161 | // *************************************************************************** |
| 162 | // SpaceToBatch Op. |
| 163 | if (stb_op->type != OperatorType::kSpaceToBatchND) { |
| 164 | return ::tensorflow::Status::OK(); |
| 165 | } |
| 166 | if (stb_op->inputs.size() != 3) { |
| 167 | return ::tensorflow::Status::OK(); |
| 168 | } |
| 169 | CHECK_EQ(stb_op->outputs.size(), 1); |
| 170 | // Extract the dilation factor from Input[1] of SpaceToBatch |
| 171 | // TODO(mjmatthews): Support 2D dilation factors. |
| 172 | const auto& block_shape_array = model->GetArray(stb_op->inputs[1]); |
| 173 | if (!block_shape_array.buffer) { |
| 174 | return ::tensorflow::Status::OK(); |
| 175 | } |
| 176 | CHECK_EQ(block_shape_array.shape().dimensions_count(), 1); |
| 177 | int dilation_factor = |
| 178 | block_shape_array.Array::GetBuffer<ArrayDataType::kInt32>().data[0]; |
| 179 | |
| 180 | // Expand Op |
| 181 | auto* post_stb_op = GetOpWithInput(*model, stb_op->outputs[0]); |
| 182 | if (!post_stb_op) { |
| 183 | return ::tensorflow::Status::OK(); |
| 184 | } |
| 185 | bool has_expand_op = false; |
| 186 | if (post_stb_op->type == OperatorType::kExpandDims) { |
| 187 | has_expand_op = true; |
| 188 | CHECK_EQ(post_stb_op->inputs.size(), 2); |
| 189 | CHECK_EQ(post_stb_op->outputs.size(), 1); |
| 190 | } |
| 191 | |
| 192 | // Conv Op |
| 193 | const string& input_of_conv_op = |
| 194 | has_expand_op ? post_stb_op->outputs[0] : stb_op->outputs[0]; |
| 195 | auto* conv_base_op = GetOpWithInput(*model, input_of_conv_op); |
| 196 | bool changed = false; |
| 197 | if (conv_base_op->type == OperatorType::kConv) { |
| 198 | changed = ResolveDilatedConv<ConvOperator>(model, conv_base_op, stb_op, |
| 199 | post_stb_op, has_expand_op, |
| 200 | dilation_factor); |
| 201 | if (changed) { |
| 202 | LOG(INFO) << "Replaced sub-network with Dilated Conv2D op outputting \"" |
| 203 | << conv_base_op->outputs[0] << "\"."; |
| 204 | } |
| 205 | } else if (identify_depthwise_conv_ && |
| 206 | conv_base_op->type == OperatorType::kDepthwiseConv) { |
| 207 | changed = ResolveDilatedConv<DepthwiseConvOperator>( |
| 208 | model, conv_base_op, stb_op, post_stb_op, has_expand_op, |
| 209 | dilation_factor); |
| 210 | if (changed) { |
nothing calls this directly
no test coverage detected