Check CoreML models.train() mode (#3262)
Browse files* Check CoreML models.train() mode
* Update export.py
- models/export.py +1 -0
models/export.py
CHANGED
@@ -131,6 +131,7 @@ if __name__ == '__main__':
|
|
131 |
import coremltools as ct
|
132 |
|
133 |
print(f'{prefix} starting export with coremltools {ct.__version__}...')
|
|
|
134 |
model = ct.convert(ts, inputs=[ct.ImageType('image', shape=img.shape, scale=1 / 255.0, bias=[0, 0, 0])])
|
135 |
f = opt.weights.replace('.pt', '.mlmodel') # filename
|
136 |
model.save(f)
|
|
|
131 |
import coremltools as ct
|
132 |
|
133 |
print(f'{prefix} starting export with coremltools {ct.__version__}...')
|
134 |
+
assert opt.train, 'CoreML exports should be placed in model.train() mode with `python export.py --train`'
|
135 |
model = ct.convert(ts, inputs=[ct.ImageType('image', shape=img.shape, scale=1 / 255.0, bias=[0, 0, 0])])
|
136 |
f = opt.weights.replace('.pt', '.mlmodel') # filename
|
137 |
model.save(f)
|