|
| 1 | +import importlib.util |
| 2 | +import numpy as np |
| 3 | +import os |
| 4 | +import sys |
| 5 | +import torch |
| 6 | + |
| 7 | + |
| 8 | +def get_device_index(device): |
| 9 | + '''Returns device index from a device''' |
| 10 | + |
| 11 | + if type(device) == str: |
| 12 | + # Could be 'cuda:0', 'cuda:1', or 'cpu'. with cpu, set index=0 |
| 13 | + device = torch.device(device) |
| 14 | + return 0 if device.index is None else device.index |
| 15 | + |
| 16 | + |
| 17 | +def get_device_index_from_input(input): |
| 18 | + '''Returns device index from a input PyTorch Tensor''' |
| 19 | + |
| 20 | + if isinstance(input, (list, tuple)): |
| 21 | + device_index = get_device_index(input[0].device) |
| 22 | + else: |
| 23 | + device_index = get_device_index(input.device) |
| 24 | + return device_index |
| 25 | + |
| 26 | + |
| 27 | +def get_all_gradients_finite_name_from_session(session): |
| 28 | + '''Find all_gradients_finite node on Session graph and return its name''' |
| 29 | + |
| 30 | + nodes = [x for x in session._outputs_meta if 'all_gradients_finite' in x.name] |
| 31 | + if len(nodes) != 1: |
| 32 | + raise RuntimeError("'all_gradients_finite' node not found within training session") |
| 33 | + return nodes[0].name |
| 34 | + |
| 35 | + |
| 36 | +def get_gradient_accumulation_name_from_session(session): |
| 37 | + '''Find Group_Accumulated_Gradients node on Session graph and return its name''' |
| 38 | + |
| 39 | + nodes = [x for x in session._outputs_meta if 'Group_Accumulated_Gradients' in x.name] |
| 40 | + if len(nodes) != 1: |
| 41 | + raise RuntimeError("'Group_Accumulated_Gradients' node not found within training session") |
| 42 | + return nodes[0].name |
| 43 | + |
| 44 | + |
| 45 | +def dtype_torch_to_numpy(torch_dtype): |
| 46 | + '''Converts PyTorch types to Numpy types |
| 47 | +
|
| 48 | + Also must map to types accepted by: |
| 49 | + MLDataType NumpyTypeToOnnxRuntimeType(int numpy_type) |
| 50 | +
|
| 51 | + References: |
| 52 | + https://docs.scipy.org/doc/numpy-1.13.0/user/basics.types.html |
| 53 | + https://pytorch.org/docs/stable/tensors.html |
| 54 | + ''' |
| 55 | + if torch_dtype == torch.float64 or torch_dtype == torch.double: |
| 56 | + return np.float64 |
| 57 | + elif torch_dtype == torch.float32 or torch_dtype == torch.float: |
| 58 | + return np.float32 |
| 59 | + elif torch_dtype == torch.float16 or torch_dtype == torch.half or torch_dtype == torch.bfloat16: |
| 60 | + # NOTE: numpy doesn't support bfloat16 |
| 61 | + return np.float16 |
| 62 | + elif torch_dtype == torch.int64 or torch_dtype == torch.long: |
| 63 | + return np.longlong # np.int64 doesn't work!? |
| 64 | + elif torch_dtype == torch.int32 or torch_dtype == torch.int: |
| 65 | + return np.int32 |
| 66 | + elif torch_dtype == torch.int16 or torch_dtype == torch.short: |
| 67 | + return np.int16 |
| 68 | + elif torch_dtype == torch.int8: |
| 69 | + return np.int8 |
| 70 | + elif torch_dtype == torch.uint8: |
| 71 | + return np.uint8 |
| 72 | + elif torch_dtype == torch.complex32 or torch_dtype == torch.complex64: |
| 73 | + # NOTE: numpy doesn't support complex32 |
| 74 | + return np.complex64 |
| 75 | + elif torch_dtype == torch.complex128 or torch_dtype == torch.cdouble: |
| 76 | + return np.complex128 |
| 77 | + elif torch_dtype == torch.bool: |
| 78 | + return np.bool_ |
| 79 | + else: |
| 80 | + raise ValueError( |
| 81 | + f'torch_dtype ({str(torch_dtype)}) type is not supported by Numpy') |
| 82 | + |
| 83 | + |
| 84 | +def dtype_onnx_to_torch(onnx_type): |
| 85 | + '''Converts ONNX types to PyTorch types |
| 86 | +
|
| 87 | + Reference: https://github.com/onnx/onnx/blob/master/onnx/onnx.in.proto (enum DataType) |
| 88 | + https://pytorch.org/docs/stable/tensors.html |
| 89 | + ''' |
| 90 | + onnx_types = ['UNDEFINED', 'FLOAT', 'UINT8', 'INT8', 'UINT16', 'INT16', 'INT32', 'INT64', 'STRING', |
| 91 | + 'BOOL', 'FLOAT16', 'DOUBLE', 'UINT32', 'UINT64', 'COMPLEX64', 'COMPLEX128', 'BFLOAT16'] |
| 92 | + |
| 93 | + if isinstance(onnx_type, int): |
| 94 | + assert onnx_type < len(onnx_types), "Invalid onnx_type integer" |
| 95 | + elif isinstance(onnx_type, str): |
| 96 | + onnx_type = onnx_type.upper() |
| 97 | + assert onnx_type in onnx_types, "Invalid onnx_type string" |
| 98 | + onnx_type = onnx_types.index(onnx_type) |
| 99 | + else: |
| 100 | + raise ValueError( |
| 101 | + "'onnx_type' must be an ONNX type represented by either a string or integer") |
| 102 | + |
| 103 | + if onnx_type == 0: |
| 104 | + return None |
| 105 | + elif onnx_type == 1: |
| 106 | + return torch.float |
| 107 | + elif onnx_type >= 2 and onnx_type <= 3: |
| 108 | + # NOTE: Pytorch doesn't support uint8 |
| 109 | + return torch.int8 |
| 110 | + elif onnx_type >= 4 and onnx_type <= 5: |
| 111 | + # NOTE: Pytorch doesn't support int16 |
| 112 | + return torch.int16 |
| 113 | + elif onnx_type == 6 or onnx_type == 12: |
| 114 | + # NOTE: Pytorch doesn't support uint32 |
| 115 | + return torch.int32 |
| 116 | + elif onnx_type == 7 or onnx_type == 13: |
| 117 | + # NOTE: Pytorch doesn't support uint64 |
| 118 | + return torch.int64 |
| 119 | + elif onnx_type == 8: |
| 120 | + return str |
| 121 | + elif onnx_type == 9: |
| 122 | + return torch.bool |
| 123 | + elif onnx_type == 10: |
| 124 | + return torch.float16 |
| 125 | + elif onnx_type == 11: |
| 126 | + return torch.double |
| 127 | + elif onnx_type == 14: |
| 128 | + return torch.complex64 |
| 129 | + elif onnx_type == 15: |
| 130 | + return torch.complex128 |
| 131 | + elif onnx_type == 16: |
| 132 | + return torch.bfloat |
| 133 | + |
| 134 | + |
| 135 | +def static_vars(**kwargs): |
| 136 | + r'''Decorator to add :py:attr:`kwargs` as static vars to 'func' |
| 137 | +
|
| 138 | + Example: |
| 139 | +
|
| 140 | + .. code-block:: python |
| 141 | +
|
| 142 | + >>> @static_vars(counter=0) |
| 143 | + ... def myfync(): |
| 144 | + ... myfync.counter += 1 |
| 145 | + ... return myfync.counter |
| 146 | + ... |
| 147 | + >>> print(myfunc()) |
| 148 | + 1 |
| 149 | + >>> print(myfunc()) |
| 150 | + 2 |
| 151 | + >>> print(myfunc()) |
| 152 | + 3 |
| 153 | + >>> myfunc.counter = 100 |
| 154 | + >>> print(myfunc()) |
| 155 | + 101 |
| 156 | + ''' |
| 157 | + def decorate(func): |
| 158 | + for k in kwargs: |
| 159 | + setattr(func, k, kwargs[k]) |
| 160 | + return func |
| 161 | + return decorate |
| 162 | + |
| 163 | + |
| 164 | +def import_module_from_file(file_path, module_name=None): |
| 165 | + '''Import a Python module from a file into interpreter''' |
| 166 | + |
| 167 | + assert isinstance(file_path, str) and os.path.exists(file_path),\ |
| 168 | + "'file_path' must be a full path string with the python file to load" |
| 169 | + assert module_name is None or isinstance(module_name, str) and module_name,\ |
| 170 | + "'module_name' must be a string with the python module name to load" |
| 171 | + |
| 172 | + if not module_name: |
| 173 | + module_name = os.path.basename(file_path).split('.')[0] |
| 174 | + |
| 175 | + spec = importlib.util.spec_from_file_location(module_name, file_path) |
| 176 | + module = importlib.util.module_from_spec(spec) |
| 177 | + sys.modules[module_name] = module |
| 178 | + spec.loader.exec_module(module) |
| 179 | + return module |
0 commit comments