| 1884 | } |
| 1885 | |
| 1886 | bool EstimateArithmeticOpsCount(const Model& model, const Operator& op, |
| 1887 | int64* result) { |
| 1888 | switch (op.type) { |
| 1889 | case OperatorType::kFullyConnected: |
| 1890 | case OperatorType::kConv: |
| 1891 | case OperatorType::kDepthwiseConv: { |
| 1892 | const auto& output_array = model.GetArray(op.outputs[0]); |
| 1893 | const auto& weights_array = model.GetArray(op.inputs[1]); |
| 1894 | if (!output_array.has_shape() || !weights_array.has_shape()) { |
| 1895 | return false; |
| 1896 | } |
| 1897 | int64 cols = 1; |
| 1898 | for (int i = 0; i < output_array.shape().dimensions_count() - 1; i++) { |
| 1899 | cols *= output_array.shape().dims(i); |
| 1900 | } |
| 1901 | const int64 cost_per_col = |
| 1902 | 2 * RequiredBufferSizeForShape(weights_array.shape()); |
| 1903 | *result = cost_per_col * cols; |
| 1904 | if (op.inputs.size() > 2) { |
| 1905 | // There is a bias vector. One more op per output value. |
| 1906 | *result += RequiredBufferSizeForShape(output_array.shape()); |
| 1907 | } |
| 1908 | break; |
| 1909 | } |
| 1910 | case OperatorType::kTransposeConv: { |
| 1911 | const auto& input_array = model.GetArray(op.inputs[2]); |
| 1912 | const auto& weights_array = model.GetArray(op.inputs[1]); |
| 1913 | if (!input_array.has_shape() || !weights_array.has_shape()) { |
| 1914 | return false; |
| 1915 | } |
| 1916 | const Shape& input = input_array.shape(); |
| 1917 | const Shape& weights = weights_array.shape(); |
| 1918 | // Compute op count from the seven nested loops of |
| 1919 | // tflite::reference_ops::TransposeConv(): |
| 1920 | *result = 2 * input.dims(0) * input.dims(1) * input.dims(2) * |
| 1921 | input.dims(3) * weights.dims(1) * weights.dims(2) * |
| 1922 | weights.dims(0); |
| 1923 | // Note that tflite::optimized_ops::TransposeConv() uses an im2col matrix |
| 1924 | // and has a higher op count, by a factor of (output_height*output_width) |
| 1925 | // vs. (input_height*input_width). Yet it generally performs better |
| 1926 | // because of coherent memory access. (At least for 2x2 striding. But not |
| 1927 | // likely for all cases.) |
| 1928 | break; |
| 1929 | } |
| 1930 | case OperatorType::kAdd: |
| 1931 | case OperatorType::kSub: |
| 1932 | case OperatorType::kMul: { |
| 1933 | const auto& output_array = model.GetArray(op.outputs[0]); |
| 1934 | if (!output_array.has_shape()) { |
| 1935 | return false; |
| 1936 | } |
| 1937 | *result = RequiredBufferSizeForShape(output_array.shape()); |
| 1938 | break; |
| 1939 | } |
| 1940 | case OperatorType::kAddN: { |
| 1941 | const auto& output_array = model.GetArray(op.outputs[0]); |
| 1942 | if (!output_array.has_shape()) { |
| 1943 | return false; |
no test coverage detected