Onnx converter improvements#52
Conversation
…ers. cleaned more code
| to_np, | ||
| ) | ||
|
|
||
| _ACTIVATION_OPS = {"relu": "Relu", "sigmoid": "Sigmoid", "tanh": "Tanh"} |
There was a problem hiding this comment.
can be moved within constants.py?
|
|
||
| def _add_standard_activation(layer, prefix, current, nodes): | ||
| """keras.layers.ReLU or keras.layers.Activation with a supported activation.""" | ||
| activation = ( |
There was a problem hiding this comment.
it would be better to keep some activation functions registry to avoid these if -> else conditions;
global redesign of modules can be a good idea
| ): | ||
| """Emit ONNX nodes for a single Keras layer. Returns the ONNX output name.""" | ||
|
|
||
| if isinstance(layer, PQMultiheadAttention): |
There was a problem hiding this comment.
all these if conditions can be avoided via polymorphism
| if isinstance(layer, keras.layers.Dropout): | ||
| return current # identity at inference | ||
|
|
||
| raise TypeError(f"Unsupported Keras layer type for ONNX export: {type(layer).__name__!r}") |
There was a problem hiding this comment.
would be nice to have a separate layers class that would reimplement all these functions in the way it's needed
|
|
||
| The model must have apply_final_compression() called on all PQ layers | ||
| before calling this function. Only inference-mode semantics are exported. | ||
|
|
There was a problem hiding this comment.
would be useful to stick to one type of functions documentation, for example, reST or Google Style within the modules
| assert out_scale == 2.0**output_scale_log2 | ||
|
|
||
| # ----- numerical check via onnxruntime ----- | ||
| for expected in ("Mul", "Sub", "Div", "Sigmoid"): |
There was a problem hiding this comment.
should you put this set in the beginning of the file?
| return self.a(x), self.b(x) | ||
|
|
||
|
|
||
| def test_multi_output_onnx(cfg_quant, tmp_path): |
There was a problem hiding this comment.
cfg is not how we name config usually
| with torch.no_grad(): | ||
| t0, t1 = (t.cpu().numpy() for t in model(x)) | ||
|
|
||
| path = str(tmp_path / "multi_output.onnx") |
There was a problem hiding this comment.
maybe multi_output_model instead?
| with torch.no_grad(): | ||
| torch_output = model(x).cpu().numpy() | ||
|
|
||
| path = str(tmp_path / "pqlayernorm.onnx") |
There was a problem hiding this comment.
pqlayernorm_model or chkpt/cp?
| path = str(tmp_path / f"int_{integer_ops}.onnx") | ||
| kwargs = {"integer_ops": True} if integer_ops else {"store_integer_weights": True} | ||
| convert_to_onnx(model, input_shape=(16,), output_path=path, **kwargs) | ||
| sess = ort.InferenceSession(path) |
|
Thanks for the PR, @nroope. I've left a few minor comments in the code, mostly regarding naming and a couple of places where functionality appears to be duplicated or misleading. |
2 participants
Split ONNX conversion logic to multiple files, cleanup some unecessary code