| |
| """Widen SAM 3D Body's decoder from bfloat16 to float32. |
| |
| ONNX Runtime's CPU provider has no MatMul kernel for bfloat16, so the decoder as |
| shipped loads and then refuses at its first matrix multiply -- which is what |
| `peasanttide_portrait` hit and is the only reason this script exists. Every |
| bfloat16 is exactly representable as a float32 (it *is* a float32 with the low |
| sixteen bits of the mantissa cleared), so this is a widening and not a |
| conversion: no retraining, no quantisation, and the same numbers to the bit. |
| |
| Three things have to move together or the graph will not type-check: the |
| initialisers, the tensors that live inside node attributes, and every `Cast` |
| whose target is bfloat16. The stale `value_info` is dropped and re-inferred |
| rather than edited, because editing it by hand is how a graph ends up half |
| converted. |
| |
| uv run --with onnx --with numpy python tools/widen_decoder.py in.onnx out.onnx |
| """ |
|
|
| import sys |
|
|
| import numpy as np |
| import onnx |
| from onnx import TensorProto, numpy_helper |
|
|
|
|
| def widen(tensor: TensorProto) -> bool: |
| """Rewrites one tensor in place. Answers whether it was a bfloat16 one.""" |
| if tensor.data_type != TensorProto.BFLOAT16: |
| return False |
| if tensor.raw_data: |
| raw = np.frombuffer(tensor.raw_data, dtype=np.uint16) |
| else: |
| raw = np.array(tensor.int32_data, dtype=np.uint32).astype(np.uint16) |
| wide = (raw.astype(np.uint32) << 16).view(np.float32) |
| shape = [dim for dim in tensor.dims] |
| tensor.CopyFrom(numpy_helper.from_array(wide.reshape(shape), tensor.name)) |
| return True |
|
|
|
|
| def main() -> None: |
| source, into = sys.argv[1], sys.argv[2] |
| model = onnx.load(source) |
|
|
| initialisers = sum(widen(one) for one in model.graph.initializer) |
| attributes = 0 |
| casts = 0 |
| for node in model.graph.node: |
| for attribute in node.attribute: |
| if attribute.type == onnx.AttributeProto.TENSOR: |
| attributes += widen(attribute.t) |
| elif attribute.type == onnx.AttributeProto.TENSORS: |
| for one in attribute.tensors: |
| attributes += widen(one) |
| elif attribute.name == "to" and attribute.i == TensorProto.BFLOAT16: |
| attribute.i = TensorProto.FLOAT |
| casts += 1 |
| for value in list(model.graph.input) + list(model.graph.output): |
| if value.type.tensor_type.elem_type == TensorProto.BFLOAT16: |
| value.type.tensor_type.elem_type = TensorProto.FLOAT |
| del model.graph.value_info[:] |
|
|
| model = onnx.shape_inference.infer_shapes(model, strict_mode=False) |
| onnx.save(model, into) |
| print( |
| f"{initialisers} initialisers, {attributes} attribute tensors and " |
| f"{casts} casts widened -> {into}" |
| ) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
