For ONNX exported model, GroupNorm will be represented as ATen op, this can be a drop in replacement from ATen to GroupNorm
(predict_net: caffe2_pb2.NetDef)
| 561 | |
| 562 | |
| 563 | def group_norm_replace_aten_with_caffe2(predict_net: caffe2_pb2.NetDef): |
| 564 | """ |
| 565 | For ONNX exported model, GroupNorm will be represented as ATen op, |
| 566 | this can be a drop in replacement from ATen to GroupNorm |
| 567 | """ |
| 568 | count = 0 |
| 569 | for op in predict_net.op: |
| 570 | if op.type == "ATen": |
| 571 | op_name = get_pb_arg_vals(op, "operator", None) # return byte in py3 |
| 572 | if op_name and op_name.decode() == "group_norm": |
| 573 | op.arg.remove(get_pb_arg(op, "operator")) |
| 574 | |
| 575 | if get_pb_arg_vali(op, "cudnn_enabled", None): |
| 576 | op.arg.remove(get_pb_arg(op, "cudnn_enabled")) |
| 577 | |
| 578 | num_groups = get_pb_arg_vali(op, "num_groups", None) |
| 579 | if num_groups is not None: |
| 580 | op.arg.remove(get_pb_arg(op, "num_groups")) |
| 581 | check_set_pb_arg(op, "group", "i", num_groups) |
| 582 | |
| 583 | op.type = "GroupNorm" |
| 584 | count += 1 |
| 585 | if count > 1: |
| 586 | logger.info("Replaced {} ATen operator to GroupNormOp".format(count)) |
| 587 | |
| 588 | |
| 589 | # ==== torch/utils_toffee/alias.py ============================================= |
no test coverage detected