| 31 | |
| 32 | |
| 33 | class MGXModel: |
| 34 | def __init__(self, |
| 35 | model, |
| 36 | input_shapes=None, |
| 37 | fp16=False, |
| 38 | bf16=True, |
| 39 | exhaustive_tune=False, |
| 40 | config=None): |
| 41 | |
| 42 | if isinstance(model, mgx.program): |
| 43 | self.model = model |
| 44 | |
| 45 | elif isinstance(model, str) and os.path.isfile(model): |
| 46 | |
| 47 | if model.endswith(".mxr"): |
| 48 | self.model = mgx.load(model, format="msgpack") |
| 49 | |
| 50 | elif model.endswith(".onnx"): |
| 51 | if not input_shapes: |
| 52 | raise ValueError( |
| 53 | "input_shapes need to be specified for loading a .onnx file" |
| 54 | ) |
| 55 | self.model = mgx.parse_onnx(model, map_input_dims=input_shapes) |
| 56 | if fp16: |
| 57 | mgx.quantize_fp16(self.model) |
| 58 | elif bf16: |
| 59 | mgx.quantize_bf16(self.model) |
| 60 | self.model.compile(mgx.get_target("gpu"), |
| 61 | exhaustive_tune=exhaustive_tune, |
| 62 | offload_copy=False) |
| 63 | else: |
| 64 | raise ValueError( |
| 65 | f"File type not recognized (should eend with .mxr or .onnx): {model}" |
| 66 | ) |
| 67 | else: |
| 68 | raise ValueError( |
| 69 | "model should be a migraphx.program object or path to .mxr/.onnx file" |
| 70 | ) |
| 71 | self.config = config |
| 72 | |
| 73 | self.mgx_to_torch_dtype_dict = { |
| 74 | "bool_type": torch.bool, |
| 75 | "uint8_type": torch.uint8, |
| 76 | "int8_type": torch.int8, |
| 77 | "int16_type": torch.int16, |
| 78 | "int32_type": torch.int32, |
| 79 | "int64_type": torch.int64, |
| 80 | "float_type": torch.float32, |
| 81 | "double_type": torch.float64, |
| 82 | "half_type": torch.float16, |
| 83 | } |
| 84 | self.torch_to_mgx_dtype_dict = { |
| 85 | v: k |
| 86 | for k, v in self.mgx_to_torch_dtype_dict.items() |
| 87 | } |
| 88 | |
| 89 | self.input_names = [] |
| 90 | self.output_names = [] |
no outgoing calls
no test coverage detected