remove ops if its output is not used or not in external_output
(net_def: caffe2_pb2.NetDef)
| 1008 | |
| 1009 | |
| 1010 | def remove_dead_end_ops(net_def: caffe2_pb2.NetDef): |
| 1011 | """ remove ops if its output is not used or not in external_output """ |
| 1012 | ssa, versions = core.get_ssa(net_def) |
| 1013 | versioned_external_output = [(name, versions[name]) for name in net_def.external_output] |
| 1014 | consumer_map = get_consumer_map(ssa) |
| 1015 | removed_op_ids = set() |
| 1016 | |
| 1017 | def _is_dead_end(versioned_blob): |
| 1018 | return not ( |
| 1019 | versioned_blob in versioned_external_output |
| 1020 | or ( |
| 1021 | len(consumer_map[versioned_blob]) > 0 |
| 1022 | and all(x[0] not in removed_op_ids for x in consumer_map[versioned_blob]) |
| 1023 | ) |
| 1024 | ) |
| 1025 | |
| 1026 | for i, ssa_i in reversed(list(enumerate(ssa))): |
| 1027 | versioned_outputs = ssa_i[1] |
| 1028 | if all(_is_dead_end(outp) for outp in versioned_outputs): |
| 1029 | removed_op_ids.add(i) |
| 1030 | |
| 1031 | # simply removing those deadend ops should have no effect to external_output |
| 1032 | new_ops = [op for i, op in enumerate(net_def.op) if i not in removed_op_ids] |
| 1033 | del net_def.op[:] |
| 1034 | net_def.op.extend(new_ops) |
no test coverage detected