ONNX exported network doesn't have concept of device, assign necessary device option for each op in order to make it runable on GPU runtime.
(
predict_net: caffe2_pb2.NetDef, init_net: caffe2_pb2.NetDef, tensor_inputs: List[torch.Tensor]
)
| 81 | |
| 82 | |
| 83 | def _assign_device_option( |
| 84 | predict_net: caffe2_pb2.NetDef, init_net: caffe2_pb2.NetDef, tensor_inputs: List[torch.Tensor] |
| 85 | ): |
| 86 | """ |
| 87 | ONNX exported network doesn't have concept of device, assign necessary |
| 88 | device option for each op in order to make it runable on GPU runtime. |
| 89 | """ |
| 90 | |
| 91 | def _get_device_type(torch_tensor): |
| 92 | assert torch_tensor.device.type in ["cpu", "cuda"] |
| 93 | assert torch_tensor.device.index == 0 |
| 94 | return torch_tensor.device.type |
| 95 | |
| 96 | def _assign_op_device_option(net_proto, net_ssa, blob_device_types): |
| 97 | for op, ssa_i in zip(net_proto.op, net_ssa): |
| 98 | if op.type in ["CopyCPUToGPU", "CopyGPUToCPU"]: |
| 99 | op.device_option.CopyFrom(core.DeviceOption(caffe2_pb2.CUDA, 0)) |
| 100 | else: |
| 101 | devices = [blob_device_types[b] for b in ssa_i[0] + ssa_i[1]] |
| 102 | assert all(d == devices[0] for d in devices) |
| 103 | if devices[0] == "cuda": |
| 104 | op.device_option.CopyFrom(core.DeviceOption(caffe2_pb2.CUDA, 0)) |
| 105 | |
| 106 | # update ops in predict_net |
| 107 | predict_net_input_device_types = { |
| 108 | (name, 0): _get_device_type(tensor) |
| 109 | for name, tensor in zip(predict_net.external_input, tensor_inputs) |
| 110 | } |
| 111 | predict_net_device_types = infer_device_type( |
| 112 | predict_net, known_status=predict_net_input_device_types, device_name_style="pytorch" |
| 113 | ) |
| 114 | predict_net_ssa, _ = core.get_ssa(predict_net) |
| 115 | _assign_op_device_option(predict_net, predict_net_ssa, predict_net_device_types) |
| 116 | |
| 117 | # update ops in init_net |
| 118 | init_net_ssa, versions = core.get_ssa(init_net) |
| 119 | init_net_output_device_types = { |
| 120 | (name, versions[name]): predict_net_device_types[(name, 0)] |
| 121 | for name in init_net.external_output |
| 122 | } |
| 123 | init_net_device_types = infer_device_type( |
| 124 | init_net, known_status=init_net_output_device_types, device_name_style="pytorch" |
| 125 | ) |
| 126 | _assign_op_device_option(init_net, init_net_ssa, init_net_device_types) |
| 127 | |
| 128 | |
| 129 | def export_caffe2_detection_model(model: torch.nn.Module, tensor_inputs: List[torch.Tensor]): |
no test coverage detected