MossFormer2_SS_16K, Core ML (8 seconds, 16-bit float)

A Core ML program converted from alibabasglab/MossFormer2_SS_16K for two-speaker speech separation at 16 kHz on Apple Silicon. It takes exactly 8 seconds of audio and returns two tracks.

The weights are the upstream ones stored as 16-bit floats. Nothing was retrained.

Source

Weights alibabasglab/MossFormer2_SS_16K, revision 407cb030cd66340918ebb6c8cc63b18f8592cdbe, last_best_checkpoint.pt
Weights SHA-256 00a3a48bda492db1e829b85dd443f8f43a43039a3e90f1a24962ea9caf14a11a
Network code modelscope/ClearerVoice-Studio, commit 6b3774dc79c46ae8bed2a4fa5f706f0ac8c75c61, clearvoice/clearvoice/models/mossformer2_ss
Conversion convert_coreml.py in this repository: coremltools 9.0, PyTorch 2.7.1, ML Program, macOS 13 target

A rerun reproduced weight.bin byte for byte. model.mlmodel holds the conversion date and Manifest.json holds random identifiers, so those two hash differently on each conversion.

What changed in the conversion

These only change how the network is written down, so that coremltools accepts it:

  • Negative widths in pad become slices, and padding a boolean mask becomes a concatenation.
  • The einsum forms in the network and in rotary-embedding-torch become the same products with @ and *.
  • Traced one-element sizes are read as scalars.
  • The two outputs are stacked into one tensor.

This one changes the numbers:

  • Compute precision is 16-bit float, which is what keeps memory low on the GPU. The input and output stay 32-bit float.

Inputs and outputs

Name Shape Type
Input wave [1, 128000] float32, 16 kHz mono
Output tracks [2, 1, 128000] float32, one track per speaker
  • The length is fixed at 8 seconds. Pad a shorter clip with zeros at the end, then cut the tracks back to the clip's length. The length is part of the converted graph, so another length needs another conversion.
  • The order of the two tracks is arbitrary. To join longer audio, overlap the windows and match the order on the overlap.
  • The upstream decoder scales each track to the input's RMS level. This program does not, so do that afterwards, measuring only the samples of the clip and not the padding.
  • Use MLComputeUnits.cpuAndGPU. Core ML warned that the Neural Engine could not compile a 30-second conversion of this graph, and the Neural Engine was not used in the measurements below.
let configuration = MLModelConfiguration()
configuration.computeUnits = .cpuAndGPU
let model = try MLModel(contentsOf: MLModel.compileModel(at: packageURL), configuration: configuration)
let input = MLMultiArray(MLShapedArray<Float>(scalars: paddedWave, shape: [1, 128_000]))
let features = try MLDictionaryFeatureProvider(dictionary: ["wave": input])
let tracks = try model.prediction(from: features).featureValue(for: "tracks")?.multiArrayValue

Measurements

Measured on one Apple Silicon Mac with 48 GB of unified memory. No person checked any of these results: there was no listening test and no human-checked reference. Recall below compares with text transcribed from the mix, which is itself a machine result.

On the GPU (cpuAndGPU), with a light load (load average 4 to 9), in one process per run, with runs of 30, 60, 60 and 240 clips of 8 seconds:

Time per second of audio, steady About 0.18 to 0.19 seconds (all clips: 0.14 to 0.22)
Over a run No continuing worsening up to 240 clips (32 minutes of audio): 0.179 at the end of 240 clips, 0.177 at the end of 30
First clip, with compiling and loading 3.6 to 4.2 seconds
Process peak memory footprint 0.64 to 0.72 GB, not growing over a run
Of the footprint, graphics memory About 0.39 GB of 0.60 GB, in one snapshot during a run
GPU driver's in-use system memory, added 0.58 GB median, 0.77 GB at most
Whole cost in unified memory About 0.9 to 1.1 GB, an estimate
Size on disk 131.8 MB

Two of the four runs began near 0.14 and were at about 0.18 by their end; the other two stayed near 0.18 to 0.19 from the start. The cause is not known. macOS recorded no thermal warning.

The whole cost is the process footprint plus the part of the GPU driver's figure that the footprint does not already count. The driver's figure covers the whole system and moved by hundreds of megabytes by itself, so this is not a precise number.

Difference from the upstream PyTorch model on one 8-second clip of real speech: 64.4 dB and 56.2 dB SNR for the two tracks. Lower than full precision (114 to 122 dB), as 16-bit floats make it.

Recall of the words of nine real overlapping windows of up to 8 seconds, transcribed from each track: 0.691, against 0.683 for PyTorch and 0.674 for the ONNX export. On 24 ranges cut to at most 8 seconds and zero padded: 0.567, against 0.587 for the ONNX export on the CPU. Both are small samples.

Limits

  • macOS on Apple Silicon only. It has not been run on Intel Macs, on iOS, or on other Apple Silicon chips.
  • Other fixed lengths were converted and are not published. They gained nothing: 12 s scored 32 to 41 dB from PyTorch, 16 s 48 to 50 dB and 1.12 GB, 30 s 48 to 51 dB and 1.84 GB.
  • Running on the CPU only works but is less accurate: 37 to 45 dB from PyTorch, at 0.65 to 0.81 seconds per second of audio and 0.44 GB. On the CPU, the ONNX export is closer to PyTorch (64 to 71 dB).

Use

Separating overlapping talkers before transcription. It is intended for Transcribe Captions, which sends only the overlapping passages of a video to this model.

The upstream card says the model was trained on large datasets, including open-sourced and private data, and it separates mixed speakers' speech into individual speakers' speech.

License

Apache-2.0, as upstream. See LICENSE. Upstream ships no NOTICE file. The network code derives from SpeechBrain (Apache-2.0), as its source header states.

Downloads last month
9
Inference Providers NEW
This model isn't deployed by any Inference Provider. ๐Ÿ™‹ Ask for provider support

Model tree for kkk4oru/mossformer2-ss-16k-coreml

Quantized
(6)
this model