(ML) Fix softmax after converting CoreML model becomes softmaxND

8 min readAug 27, 2023

在轉換模型為 CoreML 時，你可能會遇到一個問題，即原本的 softmax 被轉換為 softmaxND。本篇文章旨在解決這個問題，不論你是遇到 softmaxND 的情況，或是使用 PyTorch 轉換模型後沒有出現 softmax，都可以參考這篇文章來解決問題。

環境：Macbook Air M1 = 13.2 + Coremltools = 6.3.0

情況一

將 Pytorch MobileNetV2 模型轉換為 CoreML，轉換後的模型出現 softmaxND

先將模型加載進來，移除 softmaxND 後，並改變最後輸出的名稱

import coremltools

# Load the original model
spec = coremltools.models.utils.load_spec("mobile_coreML.mlmodel")
nn = spec.neuralNetworkClassifier

# Remove last layers: softmaxND
del nn.layers[-1]

# Get the index of the second-to-last layer
second_to_last_index = len(spec.neuralNetworkClassifier.layers) - 1

# Rename the output of the second-to-last layer
spec.neuralNetworkClassifier.layers[second_to_last_index].output[0] = "output_r"

# save spec
coremltools.utils.save_spec(spec, "new_mobile_coreML.mlmodel")

手動加上 softmax 層

# Load the spec of the new_mobile_coreML.mlmodel 
spec = coremltools.models.utils.load_spec("new_mobile_coreML.mlmodel")

# Add a softmax layer to the neural network
softmax_layer = spec.neuralNetworkClassifier.layers.add()
softmax_layer.name = "softmax"
softmax_layer.softmax.MergeFromString(b"")
softmax_layer.input.append("output_r")  # Change the output name to the name of the last layer.
softmax_layer.output.append("classLabelProbs")  # Specify the output name of the softmax layer

# Save the modified spec as new_mobile_coreML.mlmodel
coremltools.utils.save_spec(spec, "new_mobile_coreML.mlmodel")

output name 非常重要，這取決於原本模型的 stringClassLabels 的 input name，若原本模型與官方的輸出相同，僅需輸入 classLabelProbs 即可

若是使用錯誤的 output name classLabelProbs:labelProbabilityLayerName進行轉換，轉換後雖然看起來正常，但是在 Xcode 會出現下方的錯誤

最終轉換後的結果

Xcode 模型測試

情況二

使用 Pytorch 轉換後的模型，沒有 softmax 層，因為它由損失函數處理，雖然模型看起來能正常預測，但是輸出的值不代表有效的機率，模型的預測結果可能不可靠。

先將模型加載進來，並改變最後輸出的名稱

import coremltools

# Load the original model
spec = coremltools.models.utils.load_spec("mobile_coreML_withoutsoftmax.mlmodel")
nn = spec.neuralNetworkClassifier

# Get the index of the second-to-last layer
second_to_last_index = len(spec.neuralNetworkClassifier.layers) - 1

# Rename the output of the second-to-last layer
spec.neuralNetworkClassifier.layers[second_to_last_index].output[0] = "output_r"

# save spec
coremltools.utils.save_spec(spec, "mobile_coreML_withoutsoftmax.mlmodel")

手動加上 softmax 層

# Load the spec of the new_mobile_coreML.mlmodel 
spec = coremltools.models.utils.load_spec("mobile_coreML_withoutsoftmax.mlmodel")

# Add a softmax layer to the neural network
softmax_layer = spec.neuralNetworkClassifier.layers.add()
softmax_layer.name = "softmax"
softmax_layer.softmax.MergeFromString(b"")
softmax_layer.input.append("output_r")  # Change the output name to the name of the last layer.
softmax_layer.output.append("classLabelProbs")  # Specify the output name of the softmax layer

# Save the modified spec as new_mobile_coreML.mlmodel
coremltools.utils.save_spec(spec, "mobile_coreML_withoutsoftmax.mlmodel")

最終轉換後的結果