Rename the op_id-th operator in predict_net, change it's input_id-th input's name to the new_name. It also does automatic re-route and change external_input and init_net if necessary. - It requires the input is only consumed by this op. - This function modifies predict_n
(
predict_net: caffe2_pb2.NetDef,
init_net: caffe2_pb2.NetDef,
op_id: int,
input_id: int,
new_name: str,
from_producer: bool = False,
)
| 660 | |
| 661 | |
| 662 | def rename_op_input( |
| 663 | predict_net: caffe2_pb2.NetDef, |
| 664 | init_net: caffe2_pb2.NetDef, |
| 665 | op_id: int, |
| 666 | input_id: int, |
| 667 | new_name: str, |
| 668 | from_producer: bool = False, |
| 669 | ): |
| 670 | """ |
| 671 | Rename the op_id-th operator in predict_net, change it's input_id-th input's |
| 672 | name to the new_name. It also does automatic re-route and change |
| 673 | external_input and init_net if necessary. |
| 674 | - It requires the input is only consumed by this op. |
| 675 | - This function modifies predict_net and init_net in-place. |
| 676 | - When from_producer is enable, this also updates other operators that consumes |
| 677 | the same input. Be cautious because may trigger unintended behavior. |
| 678 | """ |
| 679 | assert isinstance(predict_net, caffe2_pb2.NetDef) |
| 680 | assert isinstance(init_net, caffe2_pb2.NetDef) |
| 681 | |
| 682 | init_net_ssa, init_net_versions = core.get_ssa(init_net) |
| 683 | predict_net_ssa, predict_net_versions = core.get_ssa( |
| 684 | predict_net, copy.deepcopy(init_net_versions) |
| 685 | ) |
| 686 | |
| 687 | versioned_inputs, versioned_outputs = predict_net_ssa[op_id] |
| 688 | old_name, version = versioned_inputs[input_id] |
| 689 | |
| 690 | if from_producer: |
| 691 | producer_map = get_producer_map(predict_net_ssa) |
| 692 | if not (old_name, version) in producer_map: |
| 693 | raise NotImplementedError( |
| 694 | "Can't find producer, the input {} is probably from" |
| 695 | " init_net, this is not supported yet.".format(old_name) |
| 696 | ) |
| 697 | producer = producer_map[(old_name, version)] |
| 698 | rename_op_output(predict_net, producer[0], producer[1], new_name) |
| 699 | return |
| 700 | |
| 701 | def contain_targets(op_ssa): |
| 702 | return (old_name, version) in op_ssa[0] |
| 703 | |
| 704 | is_consumer = [contain_targets(op_ssa) for op_ssa in predict_net_ssa] |
| 705 | if sum(is_consumer) > 1: |
| 706 | raise IllegalGraphTransformError( |
| 707 | ( |
| 708 | "Input '{}' of operator(#{}) are consumed by other ops, please use" |
| 709 | + " rename_op_output on the producer instead. Offending op: \n{}" |
| 710 | ).format(old_name, op_id, predict_net.op[op_id]) |
| 711 | ) |
| 712 | |
| 713 | # update init_net |
| 714 | _rename_versioned_blob_in_proto( |
| 715 | init_net, old_name, new_name, version, init_net_ssa, {}, init_net_versions |
| 716 | ) |
| 717 | # update predict_net |
| 718 | _rename_versioned_blob_in_proto( |
| 719 | predict_net, |
no test coverage detected