|
1 | 1 | fromfunctoolsimportpartial |
2 | | -fromtypingimportAny,Dict,List,Optional,Sequence,Tuple,Union |
| 2 | +fromtypingimportAny,Dict,List,Optional,Sequence,Set,Tuple,Union |
3 | 3 | importnumpyasnp |
4 | 4 | importonnx.helperasoh |
5 | 5 | importonnx.numpy_helperasonh |
@@ -604,14 +604,56 @@ def to_onnx( |
604 | 604 | model=oh.make_model(graph,opset_imports=opsets) |
605 | 605 | returnmodel |
606 | 606 |
|
607 | | -defoptimize(self): |
| 607 | +def_check_order_node(self,ind:int,node:NodeProto,existing:Set[str]): |
| 608 | +foriinnode.input: |
| 609 | +ifinotinexisting: |
| 610 | +raiseRuntimeError( |
| 611 | +f"Unknown input{i!r} from node{ind}:{node.op_type}:{node.name}. " |
| 612 | +f"Known:{existing}." |
| 613 | + ) |
| 614 | +forattinnode.attribute: |
| 615 | +ifatt.type==AttributeProto.GRAPHandatt.g: |
| 616 | +g_existing=existing.copy() |
| 617 | +foriinatt.g.input: |
| 618 | +g_existing.add(i.name) |
| 619 | +forind2,node2inenumerate(att.g.node): |
| 620 | +self._check_order_node((ind,ind2),node2,g_existing) |
| 621 | +foroinatt.g.output: |
| 622 | +ifo.namenoting_existing: |
| 623 | +raiseRuntimeError( |
| 624 | +f"Unknown output{o.name!r}. Known:{g_existing}." |
| 625 | + ) |
| 626 | +foroinnode.output: |
| 627 | +existing.add(o) |
| 628 | + |
| 629 | +defcheck_order(self): |
| 630 | +existing=set(self.initializers_dict) |
| 631 | +foriinself.inputs: |
| 632 | +existing.add(i.name) |
| 633 | +forind,nodeinenumerate(self.nodes): |
| 634 | +self._check_order_node(ind,node,existing) |
| 635 | +foroinself.outputs: |
| 636 | +ifo.namenotinexisting: |
| 637 | +raiseRuntimeError(f"Unknown output{o.name!r}. Known:{existing}.") |
| 638 | + |
| 639 | +defoptimize(self,check_order:bool=False): |
| 640 | +ifcheck_order: |
| 641 | +self.check_order() |
608 | 642 | self.remove_identity_nodes() |
| 643 | +ifcheck_order: |
| 644 | +self.check_order() |
609 | 645 | ifself.optimization_options.remove_unused: |
610 | 646 | self.remove_unused() |
| 647 | +ifcheck_order: |
| 648 | +self.check_order() |
611 | 649 | ifself.optimization_options.constant_folding: |
612 | 650 | self.constant_folding() |
| 651 | +ifcheck_order: |
| 652 | +self.check_order() |
613 | 653 | ifself.optimization_options.remove_unused: |
614 | 654 | self.remove_unused() |
| 655 | +ifcheck_order: |
| 656 | +self.check_order() |
615 | 657 |
|
616 | 658 | defremove_unused(self): |
617 | 659 | """ |
|