Skip to content

Commit 55a4479

Browse files
csukuangfjcsukuangfj
authored
Export spleeter model to onnx for source separation (k2-fsa#2237)
1 parent 901b3f0 commit 55a4479

10 files changed

Lines changed: 1100 additions & 0 deletions

File tree

‎.github/workflows/export-spleeter-to-onnx.yaml‎

Lines changed: 117 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,117 @@
1+
name: export-spleeter-to-onnx
2+
3+
on:
4+
push:
5+
branches:
6+
- spleeter-2
7+
workflow_dispatch:
8+
9+
concurrency:
10+
group: export-spleeter-to-onnx-${{ github.ref }}
11+
cancel-in-progress: true
12+
13+
jobs:
14+
export-spleeter-to-onnx:
15+
if: github.repository_owner == 'k2-fsa' || github.repository_owner == 'csukuangfj'
16+
name: export spleeter to ONNX
17+
runs-on: ${{ matrix.os }}
18+
strategy:
19+
fail-fast: false
20+
matrix:
21+
os: [macos-latest]
22+
python-version: ["3.10"]
23+
24+
steps:
25+
- uses: actions/checkout@v4
26+
27+
- name: Setup Python ${{ matrix.python-version }}
28+
uses: actions/setup-python@v5
29+
with:
30+
python-version: ${{ matrix.python-version }}
31+
32+
- name: Install dependencies
33+
shell: bash
34+
run: |
35+
pip install tensorflow torch "numpy<2" onnx==1.17.0 onnxruntime==1.17.1 onnxmltools
36+
37+
- name: Run
38+
shell: bash
39+
run: |
40+
cd scripts/spleeter
41+
./run.sh
42+
43+
echo "---"
44+
ls -lh 2stems
45+
echo "---"
46+
ls -lh 2stems/*.onnx
47+
echo "---"
48+
49+
mv -v 2stems/*.onnx ../..
50+
51+
- name: Collect models
52+
shell: bash
53+
run: |
54+
mkdir sherpa-onnx-spleeter-2stems
55+
mkdir sherpa-onnx-spleeter-2stems-int8
56+
mkdir sherpa-onnx-spleeter-2stems-fp16
57+
58+
mv -v vocals.onnx sherpa-onnx-spleeter-2stems/
59+
mv -v accompaniment.onnx sherpa-onnx-spleeter-2stems/
60+
61+
mv -v vocals.int8.onnx sherpa-onnx-spleeter-2stems-int8/
62+
mv -v accompaniment.int8.onnx sherpa-onnx-spleeter-2stems-int8/
63+
64+
mv -v vocals.fp16.onnx sherpa-onnx-spleeter-2stems-fp16/
65+
mv -v accompaniment.fp16.onnx sherpa-onnx-spleeter-2stems-fp16/
66+
67+
tar cjvf sherpa-onnx-spleeter-2stems.tar.bz2 sherpa-onnx-spleeter-2stems
68+
tar cjvf sherpa-onnx-spleeter-2stems-int8.tar.bz2 sherpa-onnx-spleeter-2stems-int8
69+
tar cjvf sherpa-onnx-spleeter-2stems-fp16.tar.bz2 sherpa-onnx-spleeter-2stems-fp16
70+
71+
ls -lh *.tar.bz2
72+
73+
- name: Release
74+
uses: svenstaro/upload-release-action@v2
75+
with:
76+
file_glob: true
77+
file: ./*.tar.bz2
78+
overwrite: true
79+
repo_name: k2-fsa/sherpa-onnx
80+
repo_token: ${{ secrets.UPLOAD_GH_SHERPA_ONNX_TOKEN }}
81+
tag: source-separation-models
82+
83+
- name: Publish to huggingface
84+
env:
85+
HF_TOKEN: ${{ secrets.HF_TOKEN }}
86+
uses: nick-fields/retry@v3
87+
with:
88+
max_attempts: 20
89+
timeout_seconds: 200
90+
shell: bash
91+
command: |
92+
git config --global user.email "csukuangfj@gmail.com"
93+
git config --global user.name "Fangjun Kuang"
94+
95+
export GIT_LFS_SKIP_SMUDGE=1
96+
export GIT_CLONE_PROTECTION_ACTIVE=false
97+
98+
names=(
99+
sherpa-onnx-spleeter-2stems
100+
sherpa-onnx-spleeter-2stems-int8
101+
sherpa-onnx-spleeter-2stems-fp16
102+
)
103+
for d in ${names[@]}; do
104+
rm -rf huggingface
105+
git clone https://csukuangfj:$HF_TOKEN@huggingface.co/csukuangfj/$d huggingface
106+
cp -v $d/*onnx huggingface
107+
108+
cd huggingface
109+
git lfs track "*.onnx"
110+
git status
111+
git add .
112+
ls -lh
113+
git status
114+
git commit -m "add models"
115+
git push https://csukuangfj:$HF_TOKEN@huggingface.co/csukuangfj/$d main
116+
cd ..
117+
done

‎scripts/spleeter/.gitignore‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,2 @@
1+
2stems.tar.gz
2+
2stems

‎scripts/spleeter/__init__.py‎

Whitespace-only changes.

‎scripts/spleeter/convert_to_pb.py‎

Lines changed: 89 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,89 @@
1+
#!/usr/bin/env python3
2+
3+
# Code in this file is modified from
4+
# https://blog.metaflow.fr/tensorflow-how-to-freeze-a-model-and-serve-it-with-a-python-api-d4f3596b3adc
5+
#
6+
# Please see ./run.sh for usages
7+
import argparse
8+
9+
import tensorflow as tf
10+
11+
12+
def freeze_graph(model_dir, output_node_names, output_filename):
13+
"""Extract the sub graph defined by the output nodes and convert all its
14+
variables into constant
15+
16+
Args:
17+
model_dir:
18+
the root folder containing the checkpoint state file
19+
output_node_names:
20+
a string, containing all the output node's names, comma separated
21+
output_filename:
22+
Filename to save the graph.
23+
"""
24+
if not tf.compat.v1.gfile.Exists(model_dir):
25+
raise AssertionError(
26+
"Export directory doesn't exists. Please specify an export "
27+
"directory: %s" % model_dir
28+
)
29+
30+
if not output_node_names:
31+
print("You need to supply the name of a node to --output_node_names.")
32+
return -1
33+
34+
# We retrieve our checkpoint fullpath
35+
checkpoint = tf.train.get_checkpoint_state(model_dir)
36+
input_checkpoint = checkpoint.model_checkpoint_path
37+
38+
# We precise the file fullname of our freezed graph
39+
output_graph = output_filename
40+
41+
# We clear devices to allow TensorFlow to control on which device it will load operations
42+
clear_devices = True
43+
44+
# We start a session using a temporary fresh Graph
45+
with tf.compat.v1.Session(graph=tf.Graph()) as sess:
46+
# We import the meta graph in the current default Graph
47+
saver = tf.compat.v1.train.import_meta_graph(
48+
input_checkpoint + ".meta", clear_devices=clear_devices
49+
)
50+
51+
# We restore the weights
52+
saver.restore(sess, input_checkpoint)
53+
54+
# We use a built-in TF helper to export variables to constants
55+
output_graph_def = tf.compat.v1.graph_util.convert_variables_to_constants(
56+
sess, # The session is used to retrieve the weights
57+
tf.compat.v1.get_default_graph().as_graph_def(), # The graph_def is used to retrieve the nodes
58+
output_node_names.split(
59+
","
60+
), # The output node names are used to select the usefull nodes
61+
)
62+
63+
# Finally we serialize and dump the output graph to the filesystem
64+
with tf.compat.v1.gfile.GFile(output_graph, "wb") as f:
65+
f.write(output_graph_def.SerializeToString())
66+
print("%d ops in the final graph." % len(output_graph_def.node))
67+
68+
return output_graph_def
69+
70+
71+
if __name__ == "__main__":
72+
parser = argparse.ArgumentParser()
73+
parser.add_argument(
74+
"--model-dir", type=str, default="", help="Model folder to export"
75+
)
76+
parser.add_argument(
77+
"--output-node-names",
78+
type=str,
79+
default="vocals_spectrogram/mul,accompaniment_spectrogram/mul",
80+
help="The name of the output nodes, comma separated.",
81+
)
82+
83+
parser.add_argument(
84+
"--output-filename",
85+
type=str,
86+
)
87+
args = parser.parse_args()
88+
89+
freeze_graph(args.model_dir, args.output_node_names, args.output_filename)

0 commit comments

Comments
 (0)