| 596 | } // namespace |
| 597 | |
| 598 | Status DoQuantizeTraining(int32 num_bits, const string& quant_op_type, |
| 599 | Graph* graph) { |
| 600 | if (graph == nullptr) { |
| 601 | return errors::InvalidArgument("Cannot accept empty graph pointer."); |
| 602 | } |
| 603 | |
| 604 | if (num_bits < 1 || num_bits > 63) { |
| 605 | return errors::OutOfRange("num_bits should be in range [1, 63] but is: ", |
| 606 | num_bits); |
| 607 | } |
| 608 | int potential_input = 0; |
| 609 | std::vector<EdgeToConvert> target_edges; |
| 610 | for (Node* node : graph->nodes()) { |
| 611 | if (nodes_to_rewrite->find(node->type_string()) != |
| 612 | nodes_to_rewrite->end() && |
| 613 | !IsGradientNode(graph, node)) { |
| 614 | // Find out which types are the inputs and convert them accordingly. |
| 615 | // 1. Const/Variable OP: This is quantized as signed tensors with no given |
| 616 | // range. |
| 617 | // 2. Activation OP: Set the range accordingly for different types of |
| 618 | // activations. Currently we handle {Relu, Relu6, Sigmoid, Tanh} |
| 619 | // 3. Identity OP: The quantization parameters depend on its input. |
| 620 | // 4. Pooling OPs: various pooling ops. Also depends on its input. |
| 621 | // 5. Reshape OP: Also depends on the first input to this op. |
| 622 | // 6. Not-Listed-Above OP: If there is only 1 such op, consider it as the |
| 623 | // model input. However, if there are >1 unknown ops, then returns an |
| 624 | // error for now to avoid unexpected behavior. |
| 625 | // Note: The list above might not be a complete list. Please let us |
| 626 | // know if you see the error so we can handle your case. |
| 627 | for (const Edge* edge : node->in_edges()) { |
| 628 | if (edge->src_output() == Graph::kControlSlot) { |
| 629 | // Skip the control dependency input. |
| 630 | continue; |
| 631 | } else { |
| 632 | bool signed_input = false; |
| 633 | bool range_given = false; |
| 634 | float input_min = 0; |
| 635 | float input_max = 0; |
| 636 | bool known_op = FindType(graph, edge->src(), &signed_input, |
| 637 | &range_given, &input_min, &input_max); |
| 638 | if (!known_op) { |
| 639 | // Unknown op is considered as input. |
| 640 | potential_input++; |
| 641 | if (potential_input > kAllowedInputs) { |
| 642 | return errors::Unimplemented( |
| 643 | "Found an unknown op: ", edge->src()->name(), |
| 644 | " with type: ", edge->src()->type_string(), |
| 645 | "; Unknown ops are considered as model input for now and " |
| 646 | "only ", |
| 647 | kAllowedInputs, " inputs are supported currently."); |
| 648 | } |
| 649 | } |
| 650 | |
| 651 | target_edges.emplace_back(EdgeToConvert( |
| 652 | edge, num_bits, signed_input, range_given, input_min, input_max)); |
| 653 | } |
| 654 | } |
| 655 | } |