Fix coreml quantization (#434)

Instead of importing platform, you can check it by `sys.platform`. By the way, `prefix` is a missing value in this script.
pull/449/head
Taka 2022-08-06 05:13:32 +08:00 committed by GitHub
parent 1560034c77
commit cfde076fb4
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 2 additions and 2 deletions

View File

@ -89,12 +89,12 @@ if __name__ == '__main__':
ct_model = ct.convert(ts, inputs=[ct.ImageType('image', shape=img.shape, scale=1 / 255.0, bias=[0, 0, 0])])
bits, mode = (8, 'kmeans_lut') if opt.int8 else (16, 'linear') if opt.fp16 else (32, None)
if bits < 32:
if platform.system() == 'Darwin': # quantization only supported on macOS
if sys.platform == 'Darwin': # quantization only supported on macOS
with warnings.catch_warnings():
warnings.filterwarnings("ignore", category=DeprecationWarning) # suppress numpy==1.20 float warning
ct_model = ct.models.neural_network.quantization_utils.quantize_weights(ct_model, bits, mode)
else:
print(f'{prefix} quantization only supported on macOS, skipping...')
print('quantization only supported on macOS, skipping...')
f = opt.weights.replace('.pt', '.mlmodel') # filename
ct_model.save(f)