Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Kwargs propagation from _get_model in quantize_weights to MLModel in _get_model #2244

Open
wants to merge 1 commit into
base: main
Choose a base branch
from
+3 −3
Open

Kwargs propagation from _get_model in quantize_weights to MLModel in _get_model #2244

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion coremltools/models/neural_network/quantization_utils.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1688,7 +1688,7 @@ def quantize_weights(
raise Exception("updatable models cannot get quantized to FP16.")

qspec = _quantize_spec_weights(spec, nbits, qmode, **kwargs)
quantized_model = _get_model(qspec, compute_units=full_precision_model.compute_unit)
quantized_model = _get_model(qspec, compute_units=full_precision_model.compute_unit, **kwargs)

if _macos_version() >= (10, 14) and sample_data:
compare_models(full_precision_model, quantized_model, sample_data)
Expand Down
4 changes: 2 additions & 2 deletions coremltools/models/utils.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -356,7 +356,7 @@ def _convert_neural_network_weights_to_fp16(full_precision_model):
return _get_model(_convert_neural_network_spec_weights_to_fp16(spec))


def _get_model(spec, compute_units=_ComputeUnit.ALL):
def _get_model(spec, compute_units=_ComputeUnit.ALL, **kwargs):
"""
Utility to get the model and the data.
"""
Expand All @@ -365,7 +365,7 @@ def _get_model(spec, compute_units=_ComputeUnit.ALL):
if isinstance(spec, MLModel):
return spec
else:
return MLModel(spec, compute_units=compute_units)
return MLModel(spec, compute_units=compute_units, **kwargs)

def evaluate_regressor(model, data, target="target", verbose=False):
"""
Expand Down