CoreML自动将int转换为float

Question

我是coremltool的新用户，目前我能够转换我的TF模型，并在python中运行它。然而，我的模型的输出层会自动转换为float，而不是int32。我必须在python中执行回滚转换。

model.predict({"Input": image}).astype(int)

有人知道如何取消这种自动转换吗？并将我的输出保持为多维数组-int？

谢谢你，

++t

Answer 1

您可以使用coremltools Python库将模型输出的数据类型更改为INT32。这将自动将这些数字从浮点转换为整数。

类似这样的东西：

import coremltools
model = coremltools.models.MLModel("YourModel.mlmodel")
spec = model._spec
spec.description.output[0].type.multiArrayType.dataType = coremltools.proto.FeatureTypes_pb2.ArrayFeatureType.INT32
coremltools.models.utils.save_spec(spec, "YourNewModel.mlmodel")

这不是自动完成的原因，是因为通常浮点是正确的数据类型。