From 3cc6da4a0de2525e17d14c9dadffb036ab607240 Mon Sep 17 00:00:00 2001 From: AlumKal Date: Thu, 26 Oct 2023 19:27:02 +0800 Subject: [PATCH 1/2] Fix setup.py --- README.rst | 8 ++++++-- setup.py | 4 ++-- 2 files changed, 8 insertions(+), 4 deletions(-) diff --git a/README.rst b/README.rst index dd99bbf42..8b06d72af 100644 --- a/README.rst +++ b/README.rst @@ -95,8 +95,12 @@ Package Index) `_:: pip install madmom -This includes the latest code and trained models and will install all -dependencies automatically. +This includes the latest stable release and will install all dependencies +automatically. + +Alternatively, if you prefer the latest code that might be unstable:: + + pip install git+https://github.com/CPJKU/madmom You might need higher privileges (use su or sudo) to install the package, model files and scripts globally. Alternatively you can install the package locally diff --git a/setup.py b/setup.py index c28ff1238..0e1df270f 100755 --- a/setup.py +++ b/setup.py @@ -43,11 +43,11 @@ package_data = [ 'models/LICENSE', 'models/README.rst', - 'models/beats/201[56]/*', + 'models/beats/*/*', 'models/chords/*/*', 'models/chroma/*/*', 'models/downbeats/*/*', - 'models/key/2018/*', + 'models/key/*/*', 'models/notes/*/*', 'models/onsets/*/*', 'models/patterns/*/*', From 83777460eeb1872fdb63660f0a62c4c7ff4e4607 Mon Sep 17 00:00:00 2001 From: AlumKal Date: Tue, 24 Feb 2026 13:04:04 +0800 Subject: [PATCH 2/2] use ONNX models --- .gitignore | 1 + .gitmodules | 2 +- MANIFEST.in | 5 + README.rst | 55 +- docs/conf.py | 76 +- docs/installation.rst | 14 +- madmom/__init__.py | 27 +- madmom/features/notes.py | 142 +- madmom/ml/nn/__init__.py | 135 +- madmom/ml/nn/onnx_runtime.py | 239 +++ madmom/models | 2 +- setup.py | 105 +- tests/test_ml_nn.py | 704 +++--- tests/test_ml_nn_onnx_parity.py | 158 ++ tests/test_ml_nn_onnx_primitives.py | 927 ++++++++ tests/test_ml_nn_onnx_tcn.py | 537 +++++ tools/convert_models_to_onnx.py | 3090 +++++++++++++++++++++++++++ 17 files changed, 5767 insertions(+), 452 deletions(-) create mode 100644 madmom/ml/nn/onnx_runtime.py create mode 100644 tests/test_ml_nn_onnx_parity.py create mode 100644 tests/test_ml_nn_onnx_primitives.py create mode 100644 tests/test_ml_nn_onnx_tcn.py create mode 100644 tools/convert_models_to_onnx.py diff --git a/.gitignore b/.gitignore index 020099b95..72efb9a3c 100644 --- a/.gitignore +++ b/.gitignore @@ -16,6 +16,7 @@ # Python egg /madmom.egg-info/ +/madmom_onnx.egg-info/ # Vim swap files *.swp diff --git a/.gitmodules b/.gitmodules index 6f4cca1c5..f31433c50 100644 --- a/.gitmodules +++ b/.gitmodules @@ -1,4 +1,4 @@ [submodule "madmom/models"] path = madmom/models - url = https://github.com/CPJKU/madmom_models.git + url = https://github.com/alumkal/madmom_models.git branch = master diff --git a/MANIFEST.in b/MANIFEST.in index 9e446c6ab..3cd7ce449 100644 --- a/MANIFEST.in +++ b/MANIFEST.in @@ -3,3 +3,8 @@ prune tests include CHANGES.rst recursive-include madmom/ *.pyx recursive-exclude madmom/ *.c +recursive-exclude madmom/ *.pkl +include madmom/models/chords/2016/chords_cnncrf.pkl +include madmom/models/chords/2016/chords_dccrf.pkl +include madmom/models/patterns/2013/ballroom_pattern_3_4.pkl +include madmom/models/patterns/2013/ballroom_pattern_4_4.pkl diff --git a/README.rst b/README.rst index 8b06d72af..221821f1c 100644 --- a/README.rst +++ b/README.rst @@ -19,6 +19,47 @@ It includes reference implementations for some music information retrieval algorithms, please see the `References`_ section. +Changes from upstream madmom +============================ + +This is ``madmom-onnx``, a fork of `madmom `_ +that replaces the legacy pickle-based NumPy neural network inference with +`ONNX Runtime `_. + +Key differences: + +- All pre-trained neural network models are shipped as ``.onnx`` files instead + of ``.pkl``. Pickle-based model loading is removed. +- Neural network inference uses ONNX Runtime (``onnxruntime`` package) instead + of the custom NumPy-based forward pass. + +Benchmark +--------- + +Measured on a 3m45s audio file, 10 runs each (mean +/- std): + +======================== ================ ================ ========= +Program madmom (pkl) madmom-onnx Speedup +======================== ================ ================ ========= +CNNOnsetDetector 2.397 +/- 0.029s 1.442 +/- 0.046s 1.66x +BeatTracker 5.326 +/- 0.071s 4.298 +/- 0.087s 1.24x +TCNTempoDetector 5.081 +/- 0.321s 1.765 +/- 0.052s 2.88x +======================== ================ ================ ========= + +Generating ONNX model files +---------------------------- + +The ``.onnx`` model files are pre-generated and included in the repository. If +you need to regenerate them from the original ``.pkl`` files (e.g. after +modifying the converter), run:: + + python tools/convert_models_to_onnx.py --convert + +This requires the ``onnx`` and ``scipy`` packages to be installed. The converter +reads each ``.pkl`` file, extracts the neural network layers, and writes an +equivalent ``.onnx`` graph to the same directory. + + Documentation ============= @@ -59,8 +100,8 @@ you choose, please make sure that all prerequisites are installed. Prerequisites ------------- -To install the ``madmom`` package, you must have either Python 2.7 or Python -3.5 or newer and the following packages installed: +To install the ``madmom-onnx`` package, you must have either Python +3.9 or newer and the following packages installed: - `numpy `_ - `scipy `_ @@ -93,7 +134,7 @@ please follow the steps in the next section. The easiest way to install the package is via ``pip`` from the `PyPI (Python Package Index) `_:: - pip install madmom + pip install madmom-onnx This includes the latest stable release and will install all dependencies automatically. @@ -106,7 +147,7 @@ You might need higher privileges (use su or sudo) to install the package, model files and scripts globally. Alternatively you can install the package locally (i.e. only for you) by adding the ``--user`` argument:: - pip install --user madmom + pip install --user madmom-onnx This will also install the executable programs to a common place (e.g. ``/usr/local/bin``), which should be in your ``$PATH`` already. If you @@ -153,13 +194,13 @@ Upgrade a package Simply upgrade the package via pip:: - pip install --upgrade madmom [--user] + pip install --upgrade madmom-onnx [--user] If some of the provided programs or models changed (please refer to the CHANGELOG) you should first uninstall the package and then reinstall:: - pip uninstall madmom - pip install madmom [--user] + pip uninstall madmom-onnx + pip install madmom-onnx [--user] Upgrade from source ~~~~~~~~~~~~~~~~~~~ diff --git a/docs/conf.py b/docs/conf.py index e048acb92..703e195c7 100644 --- a/docs/conf.py +++ b/docs/conf.py @@ -20,7 +20,7 @@ # If extensions (or modules to document with autodoc) are in another directory, # add these directories to sys.path here. If the directory is relative to the # documentation root, use os.path.abspath to make it absolute, like shown here. -sys.path.insert(0, os.path.abspath('../madmom')) +sys.path.insert(0, os.path.abspath("../madmom")) # -- General configuration ------------------------------------------------ @@ -31,44 +31,44 @@ # extensions coming with Sphinx (named 'sphinx.ext.*') or your custom # ones. extensions = [ - 'sphinx.ext.autodoc', - 'sphinx.ext.autosummary', - 'sphinx.ext.doctest', - 'sphinx.ext.todo', - 'sphinx.ext.coverage', - 'sphinx.ext.mathjax', - 'sphinx.ext.viewcode', - 'numpydoc', + "sphinx.ext.autodoc", + "sphinx.ext.autosummary", + "sphinx.ext.doctest", + "sphinx.ext.todo", + "sphinx.ext.coverage", + "sphinx.ext.mathjax", + "sphinx.ext.viewcode", + "numpydoc", ] # see http://stackoverflow.com/q/12206334/562769 numpydoc_show_class_members = False # Add any paths that contain templates here, relative to this directory. -templates_path = ['_templates'] +templates_path = ["_templates"] # The suffix(es) of source filenames. # You can specify multiple suffix as a list of string: # source_suffix = ['.rst', '.md'] -source_suffix = '.rst' +source_suffix = ".rst" # The encoding of source files. # source_encoding = 'utf-8-sig' # The master toctree document. -master_doc = 'index' +master_doc = "index" # General information about the project. -project = u'madmom' -copyright = u'2015, madmom development team' -author = u'madmom development team' +project = "madmom" +copyright = "2015, madmom development team" +author = "madmom development team" # The version info for the project you're documenting, acts as replacement for # |version| and |release|, also used in various other places throughout the # built documents. # # The short X.Y version. -version = pkg_resources.get_distribution("madmom").version +version = pkg_resources.get_distribution("madmom-onnx").version # The full version, including alpha/beta/rc tags. release = version @@ -87,7 +87,7 @@ # List of patterns, relative to source directory, that match files and # directories to ignore when looking for source files. -exclude_patterns = ['_build'] +exclude_patterns = ["_build"] # The reST default role (used for this markup: `text`) to use for all # documents. @@ -105,7 +105,7 @@ # show_authors = False # The name of the Pygments (syntax highlighting) style to use. -pygments_style = 'sphinx' +pygments_style = "sphinx" # A list of ignored prefixes for module index sorting. # modindex_common_prefix = [] @@ -119,9 +119,10 @@ # -- Options for HTML output ---------------------------------------------- # only import and set the theme if we're building docs locally -if os.environ.get('READTHEDOCS') != 'True': +if os.environ.get("READTHEDOCS") != "True": import sphinx_rtd_theme - html_theme = 'sphinx_rtd_theme' + + html_theme = "sphinx_rtd_theme" html_theme_path = [sphinx_rtd_theme.get_html_theme_path()] # The theme to use for HTML and HTML Help pages. See the documentation for @@ -155,7 +156,7 @@ # Add any paths that contain custom static files (such as style sheets) here, # relative to this directory. They are copied after the builtin static files, # so a file named "default.css" will overwrite the builtin "default.css". -html_static_path = ['_static'] +html_static_path = ["_static"] # Add any extra paths that contain custom files (such as robots.txt or # .htaccess) here, relative to this directory. These files are copied @@ -218,20 +219,17 @@ # html_search_scorer = 'scorer.js' # Output file base name for HTML help builder. -htmlhelp_basename = 'madmomdoc' +htmlhelp_basename = "madmomdoc" # -- Options for LaTeX output --------------------------------------------- latex_elements = { # The paper size ('letterpaper' or 'a4paper'). # 'papersize': 'letterpaper', - # The font size ('10pt', '11pt' or '12pt'). # 'pointsize': '10pt', - # Additional stuff for the LaTeX preamble. # 'preamble': '', - # Latex figure (float) alignment # 'figure_align': 'htbp', } @@ -240,8 +238,13 @@ # (source start file, target name, title, # author, documentclass [howto, manual, or own class]). latex_documents = [ - (master_doc, 'madmom.tex', u'madmom Documentation', - u'madmom development team', 'manual'), + ( + master_doc, + "madmom.tex", + "madmom Documentation", + "madmom development team", + "manual", + ), ] # The name of an image file (relative to this directory) to place at the top of @@ -269,10 +272,7 @@ # One entry per manual page. List of tuples # (source start file, name, description, authors, manual section). -man_pages = [ - (master_doc, 'madmom', u'madmom Documentation', - [author], 1) -] +man_pages = [(master_doc, "madmom", "madmom Documentation", [author], 1)] # If true, show URL addresses after external links. # man_show_urls = False @@ -284,9 +284,15 @@ # (source start file, target name, title, author, # dir menu entry, description, category) texinfo_documents = [ - (master_doc, 'madmom', u'madmom Documentation', - author, 'madmom', 'One line description of project.', - 'Miscellaneous'), + ( + master_doc, + "madmom", + "madmom Documentation", + author, + "madmom", + "One line description of project.", + "Miscellaneous", + ), ] # Documents to append as an appendix to all manuals. @@ -302,4 +308,4 @@ # texinfo_no_detailmenu = False # -- other options -------------------------------------------------------- -autodoc_member_order = 'bysource' +autodoc_member_order = "bysource" diff --git a/docs/installation.rst b/docs/installation.rst index 2cb4b08cc..a3f724d88 100644 --- a/docs/installation.rst +++ b/docs/installation.rst @@ -12,10 +12,10 @@ all :ref:`prerequisites ` are installed. Prerequisites ------------- -To install the ``madmom`` package, you must have either Python 2.7 or Python +To install the ``madmom-onnx`` package, you must have either Python 2.7 or Python 3.3 or newer and the following packages installed: -To install the ``madmom`` package, you must have either Python 2.7 or Python +To install the ``madmom-onnx`` package, you must have either Python 2.7 or Python 3.5 or newer and the following packages installed: - `numpy `_ @@ -51,7 +51,7 @@ please follow the steps in :ref:`the next section `. The easiest way to install the package is via ``pip`` from the `PyPI (Python Package Index) `_:: - pip install madmom + pip install madmom-onnx This includes the latest code and trained models and will install all dependencies automatically. @@ -60,7 +60,7 @@ You might need higher privileges (use su or sudo) to install the package, model files and scripts globally. Alternatively you can install the package locally (i.e. only for you) by adding the ``--user`` argument:: - pip install --user madmom + pip install --user madmom-onnx This will also install the executable programs to a common place (e.g. ``/usr/local/bin``), which should be in your ``$PATH`` already. If you @@ -111,13 +111,13 @@ Upgrade a package Simply upgrade the package via pip:: - pip install --upgrade madmom [--user] + pip install --upgrade madmom-onnx [--user] If some of the provided programs or models changed (please refer to the CHANGELOG) you should first uninstall the package and then reinstall:: - pip uninstall madmom - pip install madmom [--user] + pip uninstall madmom-onnx + pip install madmom-onnx [--user] Upgrade from source ~~~~~~~~~~~~~~~~~~~ diff --git a/madmom/__init__.py b/madmom/__init__.py index 3e4e43f9c..ecb0f4d8d 100644 --- a/madmom/__init__.py +++ b/madmom/__init__.py @@ -23,7 +23,7 @@ from . import audio, evaluation, features, io, ml, models, processors, utils # define a version variable -__version__ = distribution("madmom") .version +__version__ = distribution("madmom-onnx").version # Create a doctest output checker that optionally ignores the unicode string # literal. @@ -70,25 +70,26 @@ def check_output(self, want, got, optionflags): """ import re + if optionflags & _NORMALIZE_ARRAYS: # in different versions of numpy arrays sometimes are displayed as # 'array([ 0. ,' or 'array([0.0,', thus correct both whitespace # after parenthesis and before commas as well as .0 decimals - got = re.sub(r'\( ', '(', got) - got = re.sub(r'\[ ', '[', got) - got = re.sub(r'0\.0', '0.', got) - got = re.sub(r'\s*,', ',', got) - want = re.sub(r'\( ', '(', want) - want = re.sub(r'\[ ', '[', want) - want = re.sub(r'0\.0', '0.', want) - want = re.sub(r'\s*,', ',', want) + got = re.sub(r"\( ", "(", got) + got = re.sub(r"\[ ", "[", got) + got = re.sub(r"0\.0", "0.", got) + got = re.sub(r"\s*,", ",", got) + want = re.sub(r"\( ", "(", want) + want = re.sub(r"\[ ", "[", want) + want = re.sub(r"0\.0", "0.", want) + want = re.sub(r"\s*,", ",", want) if optionflags & _NORMALIZE_FFT: # in different versions of numpy arrays, FFT results can be ±0.j # and the unwrapped phase ±pi - got = re.sub(r'-0.j', '+0.j', got) - want = re.sub(r'-0.j', '+0.j', want) - got = re.sub(r'-3.14159', ' 3.14159', got) - want = re.sub(r'-3.14159', ' 3.14159', want) + got = re.sub(r"-0.j", "+0.j", got) + want = re.sub(r"-0.j", "+0.j", want) + got = re.sub(r"-3.14159", " 3.14159", got) + want = re.sub(r"-3.14159", " 3.14159", want) super_check_output = _DoctestOutputChecker.check_output return super_check_output(self, want, got, optionflags) diff --git a/madmom/features/notes.py b/madmom/features/notes.py index 5c2b86383..628ead967 100644 --- a/madmom/features/notes.py +++ b/madmom/features/notes.py @@ -59,8 +59,10 @@ def __init__(self, **kwargs): from ..audio.signal import SignalProcessor, FramedSignalProcessor from ..audio.stft import ShortTimeFourierTransformProcessor from ..audio.spectrogram import ( - FilteredSpectrogramProcessor, LogarithmicSpectrogramProcessor, - SpectrogramDifferenceProcessor) + FilteredSpectrogramProcessor, + LogarithmicSpectrogramProcessor, + SpectrogramDifferenceProcessor, + ) from ..models import NOTES_BRNN from ..ml.nn import NeuralNetwork @@ -72,10 +74,12 @@ def __init__(self, **kwargs): frames = FramedSignalProcessor(frame_size=frame_size, fps=100) stft = ShortTimeFourierTransformProcessor() # caching FFT window filt = FilteredSpectrogramProcessor( - num_bands=12, fmin=30, fmax=17000, norm_filters=True) + num_bands=12, fmin=30, fmax=17000, norm_filters=True + ) spec = LogarithmicSpectrogramProcessor(mul=5, add=1) diff = SpectrogramDifferenceProcessor( - diff_ratio=0.5, positive_diffs=True, stack_diffs=np.hstack) + diff_ratio=0.5, positive_diffs=True, stack_diffs=np.hstack + ) # process each frame size with spec and diff sequentially multi.append(SequentialProcessor((frames, stft, filt, spec, diff))) # stack the features and processes everything sequentially @@ -149,24 +153,42 @@ class NoteOnsetPeakPickingProcessor(OnsetPeakPickingProcessor): [ 3.37, 75. ]]) """ + THRESHOLD = 0.5 # binary threshold - SMOOTH = 0. - PRE_AVG = 0. - POST_AVG = 0. - PRE_MAX = 0. - POST_MAX = 0. + SMOOTH = 0.0 + PRE_AVG = 0.0 + POST_AVG = 0.0 + PRE_MAX = 0.0 + POST_MAX = 0.0 COMBINE = 0.03 - DELAY = 0. - - def __init__(self, threshold=THRESHOLD, smooth=SMOOTH, pre_avg=PRE_AVG, - post_avg=POST_AVG, pre_max=PRE_MAX, post_max=POST_MAX, - combine=COMBINE, delay=DELAY, fps=None, pitch_offset=0, - **kwargs): + DELAY = 0.0 + + def __init__( + self, + threshold=THRESHOLD, + smooth=SMOOTH, + pre_avg=PRE_AVG, + post_avg=POST_AVG, + pre_max=PRE_MAX, + post_max=POST_MAX, + combine=COMBINE, + delay=DELAY, + fps=None, + pitch_offset=0, + **kwargs, + ): # pylint: disable=unused-argument super(NoteOnsetPeakPickingProcessor, self).__init__( - threshold=threshold, smooth=smooth, pre_avg=pre_avg, - post_avg=post_avg, pre_max=pre_max, post_max=post_max, - combine=combine, delay=delay, fps=fps) + threshold=threshold, + smooth=smooth, + pre_avg=pre_avg, + post_avg=post_avg, + pre_max=pre_max, + post_max=post_max, + combine=combine, + delay=delay, + fps=fps, + ) self.pitch_offset = pitch_offset def process(self, activations, **kwargs): @@ -186,8 +208,12 @@ def process(self, activations, **kwargs): """ # convert timing information to frames and set default values # TODO: use at least 1 frame if any of these values are > 0? - timings = np.array([self.smooth, self.pre_avg, self.post_avg, - self.pre_max, self.post_max]) * self.fps + timings = ( + np.array( + [self.smooth, self.pre_avg, self.post_avg, self.pre_max, self.post_max] + ) + * self.fps + ) timings = np.round(timings).astype(int) # detect the peaks (function returns int indices) onsets, pitches = peak_picking(activations, self.threshold, *timings) @@ -208,7 +234,7 @@ def process(self, activations, **kwargs): # get all onsets for this pitch onsets_ = onsets[pitches == pitch] # combine onsets - onsets_ = combine_events(onsets_, self.combine, 'left') + onsets_ = combine_events(onsets_, self.combine, "left") # zip onsets and pitches and add them to list of detections notes.extend(list(zip(onsets_, [pitch] * len(onsets_)))) else: @@ -229,7 +255,8 @@ class NotePeakPickingProcessor(NoteOnsetPeakPickingProcessor): def __init__(self, fps=100, pitch_offset=21, **kwargs): # pylint: disable=unused-argument super(NotePeakPickingProcessor, self).__init__( - fps=fps, pitch_offset=pitch_offset, **kwargs) + fps=fps, pitch_offset=pitch_offset, **kwargs + ) def _cnn_pad(data): @@ -239,6 +266,12 @@ def _cnn_pad(data): return np.concatenate((pad_start, data, pad_stop)) +def _stack_note_channels(data): + if isinstance(data, tuple): + return np.dstack(data) + return data + + class CNNPianoNoteProcessor(SequentialProcessor): """ Processor to get piano note activations from a CNN in a multi-task fashion @@ -293,10 +326,13 @@ class CNNPianoNoteProcessor(SequentialProcessor): def __init__(self, **kwargs): from ..audio.signal import SignalProcessor, FramedSignalProcessor from ..audio.stft import ShortTimeFourierTransformProcessor - from ..audio.spectrogram import (FilteredSpectrogramProcessor, - LogarithmicSpectrogramProcessor) + from ..audio.spectrogram import ( + FilteredSpectrogramProcessor, + LogarithmicSpectrogramProcessor, + ) from ..models import NOTES_CNN from ..ml.nn import NeuralNetworkEnsemble + # define pre-processing chain sig = SignalProcessor(num_channels=1, sample_rate=44100) frames = FramedSignalProcessor(frame_size=4096, fps=50) @@ -304,12 +340,13 @@ def __init__(self, **kwargs): filt = FilteredSpectrogramProcessor(num_bands=24, fmin=30, fmax=10000) spec = LogarithmicSpectrogramProcessor(add=1) # pre-processes everything sequentially - pre_processor = SequentialProcessor( - (sig, frames, stft, filt, spec, _cnn_pad)) + pre_processor = SequentialProcessor((sig, frames, stft, filt, spec, _cnn_pad)) # process the pre-processed signal with a NN nn = NeuralNetworkEnsemble.load(NOTES_CNN) # instantiate a SequentialProcessor - super(CNNPianoNoteProcessor, self).__init__((pre_processor, nn)) + super(CNNPianoNoteProcessor, self).__init__( + (pre_processor, nn, _stack_note_channels) + ) class ADSRNoteTrackingProcessor(Processor): @@ -372,21 +409,34 @@ class ADSRNoteTrackingProcessor(Processor): [ 3.42, 43. , 0.74]]) """ - def __init__(self, onset_prob=0.8, note_prob=0.8, offset_prob=0.5, - attack_length=0.04, decay_length=0.04, release_length=0.02, - complete=True, onset_threshold=0.5, note_threshold=0.5, - fps=50, pitch_offset=21, **kwargs): - from .notes_hmm import (ADSRStateSpace, ADSRTransitionModel, - ADSRObservationModel) + def __init__( + self, + onset_prob=0.8, + note_prob=0.8, + offset_prob=0.5, + attack_length=0.04, + decay_length=0.04, + release_length=0.02, + complete=True, + onset_threshold=0.5, + note_threshold=0.5, + fps=50, + pitch_offset=21, + **kwargs, + ): + from .notes_hmm import ADSRStateSpace, ADSRTransitionModel, ADSRObservationModel from ..ml.hmm import HiddenMarkovModel + # state space - self.st = ADSRStateSpace(attack_length=int(attack_length * fps), - decay_length=int(decay_length * fps), - release_length=int(release_length * fps)) + self.st = ADSRStateSpace( + attack_length=int(attack_length * fps), + decay_length=int(decay_length * fps), + release_length=int(release_length * fps), + ) # transition model - self.tm = ADSRTransitionModel(self.st, onset_prob=onset_prob, - note_prob=note_prob, - offset_prob=offset_prob) + self.tm = ADSRTransitionModel( + self.st, onset_prob=onset_prob, note_prob=note_prob, offset_prob=offset_prob + ) # observation model self.om = ADSRObservationModel(self.st) # instantiate a HMM @@ -418,12 +468,11 @@ def process(self, activations, **kwargs): # process each pitch individually for pitch in range(activations.shape[1]): # decode activations for this pitch with HMM - with np.errstate(divide='ignore'): + with np.errstate(divide="ignore"): # ignore warnings when taking the log of 0 path, _ = self.hmm.viterbi(activations[:, pitch, :]) # extract HMM note segments - segments = np.logical_and(path > self.st.attack, - path < self.st.release) + segments = np.logical_and(path > self.st.attack, path < self.st.release) # extract start and end positions (transition points) idx = np.nonzero(np.diff(segments.astype(int)))[0] # add end if needed @@ -447,8 +496,13 @@ def process(self, activations, **kwargs): if onsets[onset:offset].max() < self.onset_threshold: continue # append segment as note - notes.append([onset / self.fps, pitch + self.pitch_offset, - (offset - onset) / self.fps]) + notes.append( + [ + onset / self.fps, + pitch + self.pitch_offset, + (offset - onset) / self.fps, + ] + ) # if no note notes are detected, return empty array if len(notes) == 0: return np.empty((0, 3)) diff --git a/madmom/ml/nn/__init__.py b/madmom/ml/nn/__init__.py index afbca0cae..dc2733e31 100644 --- a/madmom/ml/nn/__init__.py +++ b/madmom/ml/nn/__init__.py @@ -10,6 +10,8 @@ from __future__ import absolute_import, division, print_function +import os + import numpy as np from ...processors import Processor, ParallelProcessor, SequentialProcessor @@ -89,8 +91,23 @@ class NeuralNetwork(Processor): """ - def __init__(self, layers): + def __init__(self, layers=None, runtime=None): self.layers = layers + self._runtime = runtime + + @classmethod + def load(cls, *infile, **kwargs): + if infile: + nn_file = infile[0] + else: + nn_file = kwargs.pop("nn_file", kwargs.pop("infile", None)) + if nn_file is None: + raise TypeError("NeuralNetwork.load() missing required model file") + from .onnx_runtime import OnnxNeuralNetworkRuntime + + _validate_nn_model_file(nn_file) + runtime = OnnxNeuralNetworkRuntime(_runtime_model_source(nn_file), **kwargs) + return cls(runtime=runtime) def process(self, data, reset=True, **kwargs): """ @@ -109,6 +126,10 @@ def process(self, data, reset=True, **kwargs): Network predictions for this data. """ + if self._runtime is not None: + return self._runtime.process(data, reset=reset, **kwargs) + if self.layers is None: + raise ValueError("neural network must define `layers` or runtime") # make data at least 2d (required by NN-layers) if isinstance(data, np.ndarray) and data.ndim < 2: data = np.array(data, subok=True, copy=False, ndmin=2) @@ -117,21 +138,71 @@ def process(self, data, reset=True, **kwargs): # activate the layer and feed the output into the next one data = layer(data, reset=reset) # squeeze predictions to contain only true dimensions + if isinstance(data, np.ndarray): + if data.ndim > 1: + return data.squeeze() + return data try: - return data.squeeze() - except AttributeError: - # multi-task networks have multiple outputs and return lists return tuple([d.squeeze() for d in data]) + except (AttributeError, TypeError): + return data def reset(self): """ Reset the neural network to its initial state. """ + if self._runtime is not None: + self._runtime.reset() + return + if self.layers is None: + raise ValueError("neural network must define `layers` or runtime") for layer in self.layers: layer.reset() +def _model_extension(nn_file): + if hasattr(nn_file, "name") and isinstance(nn_file.name, str): + file_name = nn_file.name + else: + try: + file_name = os.fspath(nn_file) + except TypeError: + return "" + return os.fspath(os.path.splitext(file_name)[1].lower()) + + +def _unsupported_pickle_runtime_message(nn_file): + return ( + "loading pickled neural networks at runtime is unsupported: '%s'. " + "Convert legacy .pkl models with " + "`python tools/convert_models_to_onnx.py --convert` and load the " + "generated .onnx artifact instead." % nn_file + ) + + +def _unsupported_model_format_message(nn_file, model_ext): + return "unsupported neural network model format '%s' for '%s'; expected '.onnx'" % ( + model_ext, + nn_file, + ) + + +def _validate_nn_model_file(nn_file): + model_ext = _model_extension(nn_file) + if model_ext == ".pkl": + raise ValueError(_unsupported_pickle_runtime_message(nn_file)) + if model_ext not in ("", ".onnx"): + raise ValueError(_unsupported_model_format_message(nn_file, model_ext)) + return nn_file + + +def _runtime_model_source(model_file): + if hasattr(model_file, "read") and callable(model_file.read): + return model_file.read() + return model_file + + class NeuralNetworkEnsemble(SequentialProcessor): """ Neural Network ensemble class. @@ -166,22 +237,21 @@ class NeuralNetworkEnsemble(SequentialProcessor): """ - def __init__(self, networks, ensemble_fn=average_predictions, - num_threads=None, **kwargs): - networks_processor = ParallelProcessor(networks, - num_threads=num_threads) - super(NeuralNetworkEnsemble, self).__init__((networks_processor, - ensemble_fn)) + def __init__( + self, networks, ensemble_fn=average_predictions, num_threads=None, **kwargs + ): + networks_processor = ParallelProcessor(networks, num_threads=num_threads) + super(NeuralNetworkEnsemble, self).__init__((networks_processor, ensemble_fn)) @classmethod - def load(cls, nn_files, **kwargs): + def load(cls, *infile, **kwargs): """ Instantiate a new Neural Network ensemble from a list of files. Parameters ---------- - nn_files : list - List of neural network model file names. + infile : tuple + If the first argument is present, it is used as `nn_files`. kwargs : dict, optional Keyword arguments passed to NeuralNetworkEnsemble. @@ -191,7 +261,20 @@ def load(cls, nn_files, **kwargs): NeuralNetworkEnsemble instance. """ - networks = [NeuralNetwork.load(f) for f in nn_files] + if infile: + nn_files = infile[0] + else: + nn_files = kwargs.pop("nn_files", kwargs.pop("infile", None)) + + if nn_files is None: + raise TypeError( + "NeuralNetworkEnsemble.load() missing required model file list" + ) + intra_op_num_threads = kwargs.pop("intra_op_num_threads", 2) + networks = [ + NeuralNetwork.load(f, intra_op_num_threads=intra_op_num_threads) + for f in nn_files + ] return cls(networks, **kwargs) @staticmethod @@ -213,13 +296,25 @@ def add_arguments(parser, nn_files): """ # pylint: disable=signature-differs + import argparse from madmom.utils import OverrideDefaultListAction + + def _nn_file_argument(value): + try: + return _validate_nn_model_file(value) + except ValueError as exc: + raise argparse.ArgumentTypeError(str(exc)) + # add neural network options - g = parser.add_argument_group('neural network arguments') - g.add_argument('--nn_files', action=OverrideDefaultListAction, - type=str, default=nn_files, - help='average the predictions of these pre-trained ' - 'neural networks (multiple files can be given, ' - 'one file per argument)') + g = parser.add_argument_group("neural network arguments") + g.add_argument( + "--nn_files", + action=OverrideDefaultListAction, + type=_nn_file_argument, + default=nn_files, + help="average the predictions of these pre-trained " + "ONNX neural networks (multiple files can be given, " + "one file per argument)", + ) # return the argument group so it can be modified if needed return g diff --git a/madmom/ml/nn/onnx_runtime.py b/madmom/ml/nn/onnx_runtime.py new file mode 100644 index 000000000..326fba0e8 --- /dev/null +++ b/madmom/ml/nn/onnx_runtime.py @@ -0,0 +1,239 @@ +from __future__ import absolute_import, division, print_function + +import importlib + +import numpy as np + +from ...processors import Processor + + +CPU_EXECUTION_PROVIDER = "CPUExecutionProvider" + + +def _import_onnxruntime(): + try: + return importlib.import_module("onnxruntime") + except ImportError as error: + raise ImportError( + "onnxruntime is required for madmom.ml.nn runtime; install the " + "optional dependency and load ONNX model files" + ) from error + + +def _validate_providers(providers): + if providers is None: + return [CPU_EXECUTION_PROVIDER] + if not isinstance(providers, (list, tuple)): + raise ValueError("providers must be a list or tuple") + normalized = list(providers) + if any(provider != CPU_EXECUTION_PROVIDER for provider in normalized): + raise ValueError( + "only CPUExecutionProvider is supported for madmom.ml.nn runtime" + ) + return [CPU_EXECUTION_PROVIDER] + + +def _create_session_options(ort, intra_op_num_threads=2): + session_options = ort.SessionOptions() + session_options.execution_mode = ort.ExecutionMode.ORT_SEQUENTIAL + session_options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL + session_options.intra_op_num_threads = intra_op_num_threads + session_options.inter_op_num_threads = 1 + session_options.enable_mem_pattern = True + session_options.enable_cpu_mem_arena = True + return session_options + + +def _state_shape(input_meta): + shape = [] + for dim in input_meta.shape: + if isinstance(dim, np.integer): + dim = int(dim) + if not isinstance(dim, int) or dim <= 0: + raise ValueError( + "state input '%s' must have static positive dimensions; got %s" + % (input_meta.name, input_meta.shape) + ) + shape.append(dim) + if not shape: + raise ValueError("state input '%s' shape must not be empty" % input_meta.name) + return tuple(shape) + + +class OnnxNeuralNetworkRuntime(Processor): + def __init__(self, model_file, providers=None, intra_op_num_threads=2): + self._model_source = model_file + self._ort = _import_onnxruntime() + self.providers = _validate_providers(providers) + self._intra_op_num_threads = intra_op_num_threads + self._configure_session() + + def _configure_session(self): + self.session_options = _create_session_options( + self._ort, intra_op_num_threads=self._intra_op_num_threads + ) + self.session = self._ort.InferenceSession( + self._model_source, + sess_options=self.session_options, + providers=self.providers, + ) + self._inputs = list(self.session.get_inputs()) + self._overridable_inputs = list(self.session.get_overridable_initializers()) + self._outputs = list(self.session.get_outputs()) + if not self._inputs: + raise ValueError("ONNX model must expose at least one input") + if not self._outputs: + raise ValueError("ONNX model must expose at least one output") + self.data_input_name = self._inputs[0].name + + state_input_meta = list(self._inputs[1:]) + existing_state_names = {state_input.name for state_input in state_input_meta} + for state_input in self._overridable_inputs: + if state_input.name == self.data_input_name: + continue + if state_input.name in existing_state_names: + continue + state_input_meta.append(state_input) + existing_state_names.add(state_input.name) + + self.state_input_meta = state_input_meta + self.state_input_names = [ + state_input.name for state_input in self.state_input_meta + ] + self._state_shapes = { + state_input.name: _state_shape(state_input) + for state_input in self.state_input_meta + } + overridable_initializers = set( + initializer.name + for initializer in self.session.get_overridable_initializers() + ) + self._state_inputs_with_default_initializer = { + name for name in self.state_input_names if name in overridable_initializers + } + + state_count = len(self.state_input_names) + if state_count > len(self._outputs): + raise ValueError( + "ONNX model has %d state inputs but only %d outputs" + % (state_count, len(self._outputs)) + ) + split_index = len(self._outputs) - state_count + self.prediction_output_names = [ + output.name for output in self._outputs[:split_index] + ] + self.state_output_names = [ + output.name for output in self._outputs[split_index:] + ] + if not self.prediction_output_names: + raise ValueError("ONNX model must expose at least one prediction output") + + self._initial_state = { + state_input_name: np.zeros(state_shape, dtype=np.float32) + for state_input_name, state_shape in self._state_shapes.items() + if state_input_name not in self._state_inputs_with_default_initializer + } + self._state = {} + self.reset() + + def __getstate__(self): + state_values = { + state_input_name: ( + None + if self._state[state_input_name] is None + else np.array( + self._state[state_input_name], dtype=np.float32, copy=True + ) + ) + for state_input_name in self.state_input_names + } + return { + "model_source": self._model_source, + "providers": list(self.providers), + "intra_op_num_threads": self._intra_op_num_threads, + "state_values": state_values, + } + + def __setstate__(self, state): + self._model_source = state["model_source"] + self._ort = _import_onnxruntime() + self.providers = _validate_providers(state["providers"]) + self._intra_op_num_threads = state.get("intra_op_num_threads", 2) + self._configure_session() + + state_values = state.get("state_values", {}) + for state_input_name in self.state_input_names: + restored_state = state_values.get( + state_input_name, self._state[state_input_name] + ) + if restored_state is None: + self._state[state_input_name] = None + continue + restored_state = np.asarray(restored_state, dtype=np.float32) + expected_shape = self._state_shapes[state_input_name] + if restored_state.shape != expected_shape: + raise ValueError( + "restored state shape mismatch for '%s': expected %s, got %s" + % (state_input_name, expected_shape, restored_state.shape) + ) + self._state[state_input_name] = np.array( + restored_state, + dtype=np.float32, + copy=True, + ) + + def reset(self): + self._state = {} + for state_input_name in self.state_input_names: + if state_input_name in self._state_inputs_with_default_initializer: + self._state[state_input_name] = None + else: + self._state[state_input_name] = np.array( + self._initial_state[state_input_name], + dtype=np.float32, + copy=True, + ) + + def process(self, data, reset=True, **kwargs): + del kwargs + if reset: + self.reset() + + input_data = np.asarray(data, dtype=np.float32) + if input_data.ndim < 2: + input_data = np.array(input_data, subok=True, copy=False, ndmin=2) + + ort_inputs = {self.data_input_name: input_data} + for state_input_name in self.state_input_names: + state_input_value = self._state[state_input_name] + if state_input_value is not None: + ort_inputs[state_input_name] = state_input_value + + raw_outputs = self.session.run( + self.prediction_output_names + self.state_output_names, + ort_inputs, + ) + output_map = dict( + zip(self.prediction_output_names + self.state_output_names, raw_outputs) + ) + + for index, state_input_name in enumerate(self.state_input_names): + state_output_name = self.state_output_names[index] + state_output = np.asarray(output_map[state_output_name], dtype=np.float32) + expected_shape = self._state_shapes[state_input_name] + if state_output.shape != expected_shape: + raise ValueError( + "state output shape mismatch for '%s': expected %s, got %s" + % (state_input_name, expected_shape, state_output.shape) + ) + self._state[state_input_name] = state_output + + predictions = [ + np.asarray(output_map[name]) for name in self.prediction_output_names + ] + if len(predictions) == 1: + pred = predictions[0] + if pred.ndim > 1: + return pred.squeeze() + return pred + return tuple(p.squeeze() if p.ndim > 1 else p for p in predictions) diff --git a/madmom/models b/madmom/models index 7e3dc1b0c..1491c4270 160000 --- a/madmom/models +++ b/madmom/models @@ -1 +1 @@ -Subproject commit 7e3dc1b0cad499792767074d03c38b194b9b0a79 +Subproject commit 1491c42703e177b812a8eb49a78e49438861ae36 diff --git a/setup.py b/setup.py index bde7c7df6..0ed3f3fb0 100755 --- a/setup.py +++ b/setup.py @@ -7,6 +7,7 @@ """ import glob +import pathlib from distutils.extension import Extension import numpy as np @@ -14,94 +15,98 @@ from setuptools import setup, find_packages # define version -version = '0.17.dev0' +version = "0.17.dev0" # define which extensions to compile include_dirs = [np.get_include()] extensions = [ Extension( - 'madmom.audio.comb_filters', - ['madmom/audio/comb_filters.pyx'], + "madmom.audio.comb_filters", + ["madmom/audio/comb_filters.pyx"], include_dirs=include_dirs, ), Extension( - 'madmom.features.beats_crf', - ['madmom/features/beats_crf.pyx'], + "madmom.features.beats_crf", + ["madmom/features/beats_crf.pyx"], include_dirs=include_dirs, ), - Extension('madmom.ml.hmm', ['madmom/ml/hmm.pyx'], include_dirs=include_dirs), + Extension("madmom.ml.hmm", ["madmom/ml/hmm.pyx"], include_dirs=include_dirs), Extension( - 'madmom.ml.nn.layers', ['madmom/ml/nn/layers.py'], include_dirs=include_dirs + "madmom.ml.nn.layers", ["madmom/ml/nn/layers.py"], include_dirs=include_dirs ), ] # define scripts to be installed by the PyPI package -scripts = glob.glob('bin/*') +scripts = glob.glob("bin/*") -# define the models to be included in the PyPI package +_models_dir = pathlib.Path("madmom/models") package_data = [ - 'models/LICENSE', - 'models/README.rst', - 'models/beats/*/*', - 'models/chords/*/*', - 'models/chroma/*/*', - 'models/downbeats/*/*', - 'models/key/*/*', - 'models/notes/*/*', - 'models/onsets/*/*', - 'models/patterns/*/*', + "models/LICENSE", + "models/README.rst", ] +for _model_file in sorted(_models_dir.rglob("*")): + if not _model_file.is_file(): + continue + if _model_file.suffix == ".pkl" and _model_file.with_suffix(".onnx").exists(): + continue + _rel = str(_model_file.relative_to("madmom")) + if _rel not in package_data: + package_data.append(_rel) # some PyPI metadata classifiers = [ - 'Development Status :: 3 - Beta', - 'Programming Language :: Python :: 3.9', - 'Programming Language :: Python :: 3.10', - 'Programming Language :: Python :: 3.11', - 'Programming Language :: Python :: 3.12', - 'Environment :: Console', - 'License :: OSI Approved :: BSD License', - 'License :: Free for non-commercial use', - 'Topic :: Multimedia :: Sound/Audio :: Analysis', - 'Topic :: Scientific/Engineering :: Artificial Intelligence', + "Development Status :: 4 - Beta", + "Programming Language :: Python :: 3.9", + "Programming Language :: Python :: 3.10", + "Programming Language :: Python :: 3.11", + "Programming Language :: Python :: 3.12", + "Programming Language :: Python :: 3.13", + "Programming Language :: Python :: 3.14", + "Environment :: Console", + "License :: OSI Approved :: BSD License", + "License :: Free for non-commercial use", + "Topic :: Multimedia :: Sound/Audio :: Analysis", + "Topic :: Scientific/Engineering :: Artificial Intelligence", ] # requirements requirements = [ - 'numpy>=1.13.4', - 'scipy>=1.13', - 'mido>=1.2.6', + "numpy>=1.13.4", + "scipy>=1.13", + "mido>=1.2.6", + "onnxruntime>=1.19.2", ] # docs to be included try: - long_description = open('README.rst', encoding='utf-8').read() - long_description += '\n' + open('CHANGES.rst', encoding='utf-8').read() + long_description = open("README.rst", encoding="utf-8").read() + long_description += "\n" + open("CHANGES.rst", encoding="utf-8").read() except TypeError: - long_description = open('README.rst').read() - long_description += '\n' + open('CHANGES.rst').read() + long_description = open("README.rst").read() + long_description += "\n" + open("CHANGES.rst").read() # the actual setup routine setup( - name='madmom', + name="madmom-onnx", version=version, - description='Python audio signal processing library', + description="Python audio signal processing library", long_description=long_description, - author='Department of Computational Perception, Johannes Kepler ' - 'University, Linz, Austria and Austrian Research Institute for ' - 'Artificial Intelligence (OFAI), Vienna, Austria', - author_email='madmom-users@googlegroups.com', - url='https://github.com/CPJKU/madmom', - license='BSD, CC BY-NC-SA', - packages=find_packages(exclude=['tests', 'docs']), + author="Department of Computational Perception, Johannes Kepler " + "University, Linz, Austria and Austrian Research Institute for " + "Artificial Intelligence (OFAI), Vienna, Austria", + author_email="madmom-users@googlegroups.com", + url="https://github.com/alumkal/madmom-onnx", + license="BSD, CC BY-NC-SA", + packages=find_packages(exclude=["tests", "docs"]), ext_modules=cythonize(extensions), - package_data={'madmom': package_data}, - exclude_package_data={'': ['tests', 'docs']}, + package_data={"madmom": package_data}, + include_package_data=False, + exclude_package_data={"": ["tests", "docs"]}, scripts=scripts, install_requires=requirements, - cmdclass={'build_ext': build_ext}, - setup_requires=['pytest-runner'], - tests_require=['pytest'], + cmdclass={"build_ext": build_ext}, + setup_requires=["pytest-runner"], + tests_require=["pytest"], classifiers=classifiers, ) diff --git a/tests/test_ml_nn.py b/tests/test_ml_nn.py index da66dd4aa..768d1a158 100644 --- a/tests/test_ml_nn.py +++ b/tests/test_ml_nn.py @@ -1,5 +1,6 @@ # encoding: utf-8 # pylint: skip-file +# pyright: reportUninitializedInstanceVariable=false """ This file contains tests for the madmom.ml.nn module. @@ -15,137 +16,152 @@ from madmom.models import * -class TestRNNClass(unittest.TestCase): +ONSETS_RNN_INPUT_SIZE = 236 +BEATS_LSTM_INPUT_SIZE = 162 +ONSETS_BRNN_INPUT_SIZE = 266 +ONSETS_BRNN_PP_INPUT_SIZE = 1 +NOTES_BRNN_INPUT_SIZE = 482 +BEATS_BLSTM_INPUT_SIZE = 266 + +class TestRNNClass(unittest.TestCase): def setUp(self): # uni-directional RNN self.rnn = NeuralNetwork.load(ONSETS_RNN[0]) - self.data = np.zeros((4, self.rnn.layers[0].weights.shape[0])) - self.data[1] = 1. - self.result = [1.78801871e-04, 8.00144131e-01, - 3.30476369e-05, 1.36037513e-04] + self.data = np.zeros((4, ONSETS_RNN_INPUT_SIZE)) + self.data[1] = 1.0 + self.result = [1.78801871e-04, 8.00144131e-01, 3.30476369e-05, 1.36037513e-04] def test_process(self): # process the whole sequence at once result = self.rnn.process(self.data) - self.assertTrue(np.allclose(result, self.result)) + self.assertTrue(np.allclose(result, self.result, atol=1e-7)) # two runs must produce the same output result_1 = self.rnn.process(self.data) - self.assertTrue(np.allclose(result_1, self.result)) + self.assertTrue(np.allclose(result_1, self.result, atol=1e-7)) # after resetting the RNN, it must produce the same output as before self.rnn.reset() - result_2 = [self.rnn.process(np.atleast_2d(d), reset=False) - for d in self.data] - self.assertTrue(np.allclose(np.hstack(result_2), self.result)) + result_2 = [self.rnn.process(np.atleast_2d(d), reset=False) for d in self.data] + self.assertTrue(np.allclose(np.hstack(result_2), self.result, atol=1e-7)) # without resetting it produces different results - result_3 = [self.rnn.process(np.atleast_2d(d), reset=False) - for d in self.data] - self.assertTrue(np.allclose(np.hstack(result_3), - [9.15636891e-04, 9.74331021e-01, - 4.83996118e-05, 2.72355013e-04])) + result_3 = [self.rnn.process(np.atleast_2d(d), reset=False) for d in self.data] + self.assertTrue( + np.allclose( + np.hstack(result_3), + [9.15636891e-04, 9.74331021e-01, 4.83996118e-05, 2.72355013e-04], + atol=1e-7, + ) + ) class TestLSTMClass(unittest.TestCase): - def setUp(self): # uni-directional LSTM-RNN self.rnn = NeuralNetwork.load(BEATS_LSTM[0]) - self.data = np.zeros((4, self.rnn.layers[0].cell.weights.shape[0])) - self.data[1] = 1. + self.data = np.zeros((4, BEATS_LSTM_INPUT_SIZE)) + self.data[1] = 1.0 self.result = [0.00126955, 0.03134079, 0.01535073, 0.00207471] def test_process(self): # process the whole sequence at once result = self.rnn.process(self.data) - self.assertTrue(np.allclose(result, self.result)) + self.assertTrue(np.allclose(result, self.result, atol=1e-7)) # two runs must produce the same output result_1 = self.rnn.process(self.data) - self.assertTrue(np.allclose(result_1, self.result)) + self.assertTrue(np.allclose(result_1, self.result, atol=1e-7)) # after resetting the RNN, it must produce the same output self.rnn.reset() - result_2 = [self.rnn.process(np.atleast_2d(d), reset=False) - for d in self.data] - self.assertTrue(np.allclose(np.hstack(result_2), self.result)) + result_2 = [self.rnn.process(np.atleast_2d(d), reset=False) for d in self.data] + self.assertTrue(np.allclose(np.hstack(result_2), self.result, atol=1e-7)) # without resetting it produces different output - result_3 = [self.rnn.process(np.atleast_2d(d), reset=False) - for d in self.data] - self.assertTrue(np.allclose(np.hstack(result_3), - [0.00054101, 0.05323271, - 0.0548761, 0.00785541])) + result_3 = [self.rnn.process(np.atleast_2d(d), reset=False) for d in self.data] + self.assertTrue( + np.allclose( + np.hstack(result_3), + [0.00054101, 0.05323271, 0.0548761, 0.00785541], + atol=1e-7, + ) + ) # class for testing all other (offline-only) networks class TestNeuralNetworkClass(unittest.TestCase): - def test_brnn(self): rnn = NeuralNetwork.load(ONSETS_BRNN[0]) - input_size = rnn.layers[0].fwd_layer.weights.shape[0] - data = np.zeros((4, input_size)) - data[1] = 1. - result = rnn.process(data) - self.assertTrue(np.allclose(result, [0.00461393, 0.46032878, - 0.04824624, 0.00083493])) + data = np.zeros((4, ONSETS_BRNN_INPUT_SIZE)) + data[1] = 1.0 + result = np.asarray(rnn.process(data)) + self.assertTrue( + np.allclose(result, [0.00461393, 0.46032878, 0.04824624, 0.00083493]) + ) def test_brnn_pp(self): rnn = NeuralNetwork.load(ONSETS_BRNN_PP[0]) - input_size = rnn.layers[0].fwd_layer.weights.shape[0] - data = np.zeros((4, input_size)) - data[1] = 1. + data = np.zeros((4, ONSETS_BRNN_PP_INPUT_SIZE)) + data[1] = 1.0 result = rnn.process(data) - self.assertTrue(np.allclose(result, - [3.88076517e-03, 1.67354920e-03, - 1.14450835e-03, 5.01533471e-05])) + self.assertTrue( + np.allclose( + result, [3.88076517e-03, 1.67354920e-03, 1.14450835e-03, 5.01533471e-05] + ) + ) def test_brnn_regression(self): rnn = NeuralNetwork.load(NOTES_BRNN[0]) - input_size = rnn.layers[0].fwd_layer.weights.shape[0] - data = np.zeros((4, input_size)) - data[1] = 1. - result = rnn.process(data) + data = np.zeros((4, NOTES_BRNN_INPUT_SIZE)) + data[1] = 1.0 + result = np.asarray(rnn.process(data)) self.assertEqual(result.shape, (4, 88)) - self.assertTrue(np.allclose(result[:, :2], - [[6.50841586e-05, 4.06891153e-04], - [-9.74552809e-04, -3.86762259e-03], - [1.09878686e-04, 1.54044293e-04], - [-8.16427571e-04, 4.62550714e-04]])) + self.assertTrue( + np.allclose( + result[:, :2], + [ + [6.50841586e-05, 4.06891153e-04], + [-9.74552809e-04, -3.86762259e-03], + [1.09878686e-04, 1.54044293e-04], + [-8.16427571e-04, 4.62550714e-04], + ], + ) + ) def test_blstm(self): rnn = NeuralNetwork.load(BEATS_BLSTM[0]) - input_size = rnn.layers[0].fwd_layer.cell.weights.shape[0] - data = np.zeros((4, input_size)) - data[1] = 1. + data = np.zeros((4, BEATS_BLSTM_INPUT_SIZE)) + data[1] = 1.0 result = rnn.process(data) - self.assertTrue(np.allclose(result, [0.0815198, 0.24451593, - 0.08786312, 0.01776425])) + self.assertTrue( + np.allclose(result, [0.0815198, 0.24451593, 0.08786312, 0.01776425]) + ) def test_cnn(self): cnn = NeuralNetwork.load(ONSETS_CNN[0]) data = np.zeros((19, 80, 3)) - data[10] = 1. + data[10] = 1.0 result = cnn.process(data) - self.assertTrue(np.allclose(result, [0.0021432, 0.02647826, 0.92750794, - 0.84207922, 0.21631248])) + self.assertTrue( + np.allclose( + result, [0.0021432, 0.02647826, 0.92750794, 0.84207922, 0.21631248] + ) + ) class TestFeedForwardLayerClass(unittest.TestCase): - def setUp(self): - # borrow an FeedForwardLayer from an existing network - rnn = NeuralNetwork.load(ONSETS_RNN[0]) - self.layer = rnn.layers[-1] - input_size = self.layer.weights.shape[0] + weights = np.array([[0.8], [-0.4], [0.2]]) + bias = np.array([0.2]) + self.layer = FeedForwardLayer(weights, bias, activation_fn=sigmoid) + input_size = weights.shape[0] self.data = np.zeros((4, input_size)) - self.data[1] = 1. - self.result = np.array([[0.16283005], [0.14362903], - [0.16283005], [0.16283005]]) + self.data[1, 0] = 1.0 + self.result = np.array([[0.549834], [0.73105858], [0.549834], [0.549834]]) def test_types(self): self.assertTrue(isinstance(self.layer, FeedForwardLayer)) self.assertTrue(self.layer.activation_fn == sigmoid) def test_init(self): - self.assertFalse(hasattr(self.layer, 'init')) - self.assertFalse(hasattr(self.layer, '_prev')) + self.assertFalse(hasattr(self.layer, "init")) + self.assertFalse(hasattr(self.layer, "_prev")) def test_activate(self): # test result @@ -163,68 +179,109 @@ def test_activate(self): class TestRecurrentLayerClass(unittest.TestCase): - def setUp(self): - # borrow an RecurrentLayer from an existing network - rnn = NeuralNetwork.load(ONSETS_RNN[0]) - self.layer = rnn.layers[0] - input_size = self.layer.weights.shape[0] + weights = np.array([[0.5, -0.3], [0.1, 0.2], [-0.2, 0.4]]) + recurrent_weights = np.array([[0.6, 0.1], [-0.2, 0.5]]) + bias = np.array([0.05, -0.1]) + self.layer = RecurrentLayer( + weights, bias, recurrent_weights, activation_fn=tanh + ) + input_size = weights.shape[0] self.data = np.zeros((4, input_size)) - self.data[1] = 1. + self.data[1, 0] = 1.0 def test_types(self): self.assertTrue(isinstance(self.layer, RecurrentLayer)) self.assertTrue(self.layer.activation_fn == tanh) def test_init(self): - self.assertTrue(hasattr(self.layer, 'init')) - self.assertTrue(hasattr(self.layer, '_prev')) + self.assertTrue(hasattr(self.layer, "init")) + self.assertTrue(hasattr(self.layer, "_prev")) def test_activate(self): # test result result_1 = self.layer(self.data) - self.assertTrue(np.allclose(result_1[0, :2], - [-0.33919713, -0.02091585])) - self.assertTrue(np.allclose(result_1[-1, -2:], - [0.4419672, -0.261151])) + self.assertTrue( + np.allclose( + result_1, + [ + [0.04995837, -0.09966799], + [0.53698454, -0.41764674], + [0.42658954, -0.24973008], + [0.34159732, -0.18021615], + ], + ) + ) # two runs must yield identical results result_2 = self.layer(self.data) self.assertTrue(np.allclose(result_2, result_1)) # last step must be preserved self.assertTrue(np.allclose(self.layer._prev, result_2[-1])) # initialisation must not change - self.assertTrue(np.allclose(self.layer.init, np.zeros(25))) + self.assertTrue(np.allclose(self.layer.init, np.zeros(2))) # reset layer, activate framewise must yield the same result self.layer.reset() - result_3 = [self.layer.activate(np.atleast_2d(d), reset=False) - for d in self.data] + result_3 = [ + self.layer.activate(np.atleast_2d(d), reset=False) for d in self.data + ] result_3 = np.vstack(result_3) self.assertTrue(np.allclose(result_3, result_1)) self.assertTrue(np.allclose(self.layer._prev, result_3[-1])) # activate framewise without resetting must yield a different result - result_4 = [self.layer.activate(np.atleast_2d(d), reset=False) - for d in self.data] + result_4 = [ + self.layer.activate(np.atleast_2d(d), reset=False) for d in self.data + ] result_4 = np.vstack(result_4) self.assertFalse(np.allclose(result_4, result_1)) - self.assertTrue(np.allclose(result_4[0, :2], - [-3.14807342e-01, -2.22700375e-01])) - self.assertTrue(np.allclose(result_4[-1, -2:], - [4.40259654e-01, -2.59141315e-01])) + self.assertTrue( + np.allclose( + result_4, + [ + [0.28305644, -0.15469631], + [0.63560996, -0.42111163], + [0.47428788, -0.24209167], + [0.36530222, -0.17189339], + ], + ) + ) # last step must be preserved self.assertTrue(np.allclose(self.layer._prev, result_4[-1])) # initialisation must not change - self.assertTrue(np.allclose(self.layer.init, np.zeros(25))) + self.assertTrue(np.allclose(self.layer.init, np.zeros(2))) class TestLSTMLayerClass(unittest.TestCase): - def setUp(self): - # borrow an LSTMLayer from an existing network - rnn = NeuralNetwork.load(BEATS_BLSTM[0]) - self.layer = rnn.layers[0].fwd_layer - input_size = self.layer.cell.weights.shape[0] + input_gate = Gate( + np.array([[0.4, -0.3], [0.1, 0.2], [-0.2, 0.5]]), + np.array([0.05, -0.02]), + np.array([[0.1, 0.2], [-0.2, 0.1]]), + activation_fn=sigmoid, + ) + forget_gate = Gate( + np.array([[0.2, 0.1], [-0.4, 0.3], [0.3, -0.2]]), + np.array([0.1, 0.2]), + np.array([[0.05, -0.1], [0.15, 0.2]]), + activation_fn=sigmoid, + ) + cell = Cell( + np.array([[0.5, -0.4], [0.2, 0.1], [-0.3, 0.2]]), + np.array([0.0, 0.05]), + np.array([[0.2, 0.0], [-0.1, 0.3]]), + activation_fn=tanh, + ) + output_gate = Gate( + np.array([[0.3, 0.2], [-0.1, 0.4], [0.2, -0.3]]), + np.array([-0.03, 0.04]), + np.array([[0.1, -0.2], [0.2, 0.1]]), + activation_fn=sigmoid, + ) + self.layer = LSTMLayer( + input_gate, forget_gate, cell, output_gate, activation_fn=tanh + ) + input_size = input_gate.weights.shape[0] self.data = np.zeros((4, input_size)) - self.data[1] = 1. + self.data[1, 0] = 1.0 def test_types(self): self.assertTrue(isinstance(self.layer, LSTMLayer)) @@ -239,16 +296,23 @@ def test_types(self): self.assertTrue(self.layer.output_gate.activation_fn == sigmoid) def test_init(self): - self.assertTrue(hasattr(self.layer, 'init')) - self.assertTrue(hasattr(self.layer, '_prev')) + self.assertTrue(hasattr(self.layer, "init")) + self.assertTrue(hasattr(self.layer, "_prev")) def test_activate(self): # test result result_1 = self.layer(self.data) - self.assertTrue(np.allclose(result_1[0, :2], - [8.15829188e-02, 1.14209838e-02])) - self.assertTrue(np.allclose(result_1[-1, -2:], - [-1.81968838e-01, -2.32963227e-02])) + self.assertTrue( + np.allclose( + result_1, + [ + [0.0, 0.01260939], + [0.1556166, -0.07017737], + [0.08170114, -0.02691965], + [0.04800879, -0.00431454], + ], + ) + ) # two runs must yield identical results result_2 = self.layer(self.data) self.assertTrue(np.allclose(result_2, result_1)) @@ -256,136 +320,179 @@ def test_activate(self): self.assertTrue(np.allclose(self.layer._prev, result_2[-1])) # reset layer, activate framewise must yield the same result self.layer.reset() - result_3 = [self.layer.activate(np.atleast_2d(d), reset=False) - for d in self.data] + result_3 = [ + self.layer.activate(np.atleast_2d(d), reset=False) for d in self.data + ] self.assertTrue(np.allclose(np.vstack(result_3), result_1)) # activate framewise without resetting must yield a different result - result_4 = [self.layer.activate(np.atleast_2d(d), reset=False) - for d in self.data] + result_4 = [ + self.layer.activate(np.atleast_2d(d), reset=False) for d in self.data + ] result_4 = np.vstack(result_4) self.assertFalse(np.allclose(result_4, result_1)) - self.assertTrue(np.allclose(result_4[0, :2], - [2.34878659e-01, -1.26488507e-03])) - self.assertTrue(np.allclose(result_4[-1, -2:], - [-2.45573238e-01, -2.92062219e-02])) + self.assertTrue( + np.allclose( + result_4, + [ + [0.02783435, 0.00990413], + [0.17430915, -0.07205639], + [0.09170447, -0.02797117], + [0.05389891, -0.00498242], + ], + ) + ) # last step must be preserved self.assertTrue(np.allclose(self.layer._prev, result_4[-1])) # initialisation must not change - self.assertTrue(np.allclose(self.layer.init, np.zeros(25))) + self.assertTrue(np.allclose(self.layer.init, np.zeros(2))) class TestGRUClass(unittest.TestCase): - - W_xr = np.array([[-0.42948743, -1.29989187], - [0.77213901, 0.86070993], - [1.13791823, -0.87066225]]) - W_xu = np.array([[0.44875312, 0.07172084], - [-0.24292999, 1.318794], - [1.0270179, 0.16293946]]) - W_xhu = np.array([[0.8812559, 1.35859991], - [1.04311944, -0.25449358], - [-1.09539597, 1.19808424]]) - W_hr = np.array([[0.96696973, 0.1384294], - [-0.09561655, -1.23413809]]) - W_hu = np.array([[0.04664641, 0.59561686], - [1.00325841, -0.11574791]]) - W_hhu = np.array([[1.19742848, 1.07850016], - [0.35234964, -1.45348681]]) + W_xr = np.array( + [ + [-0.42948743, -1.29989187], + [0.77213901, 0.86070993], + [1.13791823, -0.87066225], + ] + ) + W_xu = np.array( + [[0.44875312, 0.07172084], [-0.24292999, 1.318794], [1.0270179, 0.16293946]] + ) + W_xhu = np.array( + [[0.8812559, 1.35859991], [1.04311944, -0.25449358], [-1.09539597, 1.19808424]] + ) + W_hr = np.array([[0.96696973, 0.1384294], [-0.09561655, -1.23413809]]) + W_hu = np.array([[0.04664641, 0.59561686], [1.00325841, -0.11574791]]) + W_hhu = np.array([[1.19742848, 1.07850016], [0.35234964, -1.45348681]]) b_r = np.array([1.41851288, -0.39743243]) b_u = np.array([-0.78729095, 0.83385797]) b_hu = np.array([1.25143065, -0.97715625]) - IN = np.array([[0.91298812, -1.47626202, -1.08667502], - [0.49814883, -0.0104938, 0.93869008], - [-1.12282135, 0.3780883, 1.42017503], - [0.62669439, 0.89438929, -0.69354132], - [0.16162221, -1.00166208, 0.23579985]]) - OUT = np.array([[0.22772433, -0.13181415], - [0.49479958, 0.51224858], - [0.08539771, -0.56119639], - [0.1946809, -0.50421363], - [0.17403202, -0.27258521]]) + IN = np.array( + [ + [0.91298812, -1.47626202, -1.08667502], + [0.49814883, -0.0104938, 0.93869008], + [-1.12282135, 0.3780883, 1.42017503], + [0.62669439, 0.89438929, -0.69354132], + [0.16162221, -1.00166208, 0.23579985], + ] + ) + OUT = np.array( + [ + [0.22772433, -0.13181415], + [0.49479958, 0.51224858], + [0.08539771, -0.56119639], + [0.1946809, -0.50421363], + [0.17403202, -0.27258521], + ] + ) H = np.array([0.02345737, 0.34454183]) def setUp(self): self.reset_gate = Gate( - TestGRUClass.W_xr, TestGRUClass.b_r, TestGRUClass.W_hr, - activation_fn=sigmoid) + TestGRUClass.W_xr, + TestGRUClass.b_r, + TestGRUClass.W_hr, + activation_fn=sigmoid, + ) self.update_gate = Gate( - TestGRUClass.W_xu, TestGRUClass.b_u, TestGRUClass.W_hu, - activation_fn=sigmoid) + TestGRUClass.W_xu, + TestGRUClass.b_u, + TestGRUClass.W_hu, + activation_fn=sigmoid, + ) self.gru_cell = GRUCell( - TestGRUClass.W_xhu, TestGRUClass.b_hu, TestGRUClass.W_hhu) - self.gru_1 = GRULayer(self.reset_gate, self.update_gate, - self.gru_cell) - self.gru_2 = GRULayer(self.reset_gate, self.update_gate, - self.gru_cell, init=TestGRUClass.H) + TestGRUClass.W_xhu, TestGRUClass.b_hu, TestGRUClass.W_hhu + ) + self.gru_1 = GRULayer(self.reset_gate, self.update_gate, self.gru_cell) + self.gru_2 = GRULayer( + self.reset_gate, self.update_gate, self.gru_cell, init=TestGRUClass.H + ) def test_activate(self): self.assertTrue( - np.allclose(self.reset_gate.activate(TestGRUClass.IN[0, :], - TestGRUClass.H), np.array([0.20419282, 0.08861294]))) + np.allclose( + self.reset_gate.activate(TestGRUClass.IN[0, :], TestGRUClass.H), + np.array([0.20419282, 0.08861294]), + ) + ) self.assertTrue( - np.allclose(self.update_gate.activate(TestGRUClass.IN[0, :], - TestGRUClass.H), np.array([0.31254834, 0.2226105]))) + np.allclose( + self.update_gate.activate(TestGRUClass.IN[0, :], TestGRUClass.H), + np.array([0.31254834, 0.2226105]), + ) + ) self.assertTrue( - np.allclose(self.gru_cell.activate(TestGRUClass.IN[0, :], - TestGRUClass.H, TestGRUClass.H), - np.array([0.9366396, -0.67876764]))) + np.allclose( + self.gru_cell.activate( + TestGRUClass.IN[0, :], TestGRUClass.H, TestGRUClass.H + ), + np.array([0.9366396, -0.67876764]), + ) + ) # activating the layer normally - self.assertTrue(np.allclose(self.gru_1.activate(TestGRUClass.IN), - TestGRUClass.OUT)) + self.assertTrue( + np.allclose(self.gru_1.activate(TestGRUClass.IN), TestGRUClass.OUT) + ) # activating the layer a second time must give the same results - self.assertTrue(np.allclose(self.gru_1.activate(TestGRUClass.IN), - TestGRUClass.OUT)) + self.assertTrue( + np.allclose(self.gru_1.activate(TestGRUClass.IN), TestGRUClass.OUT) + ) # activating the other layer self.assertTrue( - np.allclose(self.gru_2.activate(TestGRUClass.IN), - np.array([[0.30988133, 0.13258138], - [0.60639685, 0.55714613], - [0.21366976, -0.55568963], - [0.30860096, -0.43686554], - [0.28866628, -0.23025239]]))) + np.allclose( + self.gru_2.activate(TestGRUClass.IN), + np.array( + [ + [0.30988133, 0.13258138], + [0.60639685, 0.55714613], + [0.21366976, -0.55568963], + [0.30860096, -0.43686554], + [0.28866628, -0.23025239], + ] + ), + ) + ) # reset layer, activate framewise must yield the same result self.gru_1.reset() - result_1 = [self.gru_1.activate(np.atleast_2d(d), reset=False) - for d in self.IN] + result_1 = [self.gru_1.activate(np.atleast_2d(d), reset=False) for d in self.IN] self.assertTrue(np.allclose(np.vstack(result_1), self.OUT)) # the previous state must be the last output self.assertTrue(np.allclose(self.gru_1._prev, self.OUT[-1])) # activate with same data without resetting - result_2 = [self.gru_1.activate(np.atleast_2d(d), reset=False) - for d in self.IN] + result_2 = [self.gru_1.activate(np.atleast_2d(d), reset=False) for d in self.IN] # results must differ self.assertFalse(np.allclose(result_1, result_2)) - self.assertTrue(np.allclose(np.vstack(result_2), - [[0.3254016, -0.33195618], - [0.53048187, 0.51293057], - [0.12495148, -0.559039], - [0.22988503, -0.48455557], - [0.20934506, -0.25959042]])) + self.assertTrue( + np.allclose( + np.vstack(result_2), + [ + [0.3254016, -0.33195618], + [0.53048187, 0.51293057], + [0.12495148, -0.559039], + [0.22988503, -0.48455557], + [0.20934506, -0.25959042], + ], + ) + ) # the previous state must be the last output - self.assertTrue(np.allclose(self.gru_1._prev, - [0.20934506, -0.25959042])) + self.assertTrue(np.allclose(self.gru_1._prev, [0.20934506, -0.25959042])) # initialisation must not change self.assertTrue(np.allclose(self.gru_1.init, [0, 0])) class TestMaxPoolLayerClass(unittest.TestCase): - def test_time_pooling(self): layer = MaxPoolLayer(size=(2, 1)) out = layer(np.arange(20).reshape(5, 4)) self.assertEqual(out.shape, (2, 4)) - self.assertTrue(np.allclose(out, [[4, 5, 6, 7], - [12, 13, 14, 15]])) + self.assertTrue(np.allclose(out, [[4, 5, 6, 7], [12, 13, 14, 15]])) def test_freq_pooling(self): layer = MaxPoolLayer(size=(1, 2)) out = layer(np.arange(20).reshape(5, 4)) self.assertEqual(out.shape, (5, 2)) - self.assertTrue(np.allclose(out, [[1, 3], [5, 7], [9, 11], - [13, 15], [17, 19]])) + self.assertTrue(np.allclose(out, [[1, 3], [5, 7], [9, 11], [13, 15], [17, 19]])) def test_max_pooling(self): layer = MaxPoolLayer((3, 3)) @@ -398,54 +505,83 @@ def test_max_pooling(self): class TestBatchNormLayerClass(unittest.TestCase): - - IN = np.array([[[0.32400414, 0.31483042], - [0.38269293, 0.04822304], - [0.03791266, 0.34776369]], - [[0.87113619, 0.62172854], - [0.87353969, 0.92837042], - [0.70359915, 0.49917081]], - [[0.42643583, 0.74653631], - [0.08519834, 0.35423595], - [0.34863797, 0.44895086]]]) - - BN_PARAMS = [np.array([-0.00098404, 0.00185387]), - np.array([-0.0068268, 0.0068859]), - np.array([-0.00289366, 0.00742069]), - np.array([-0.00177374, -0.00444383])] - - OUT = np.array([[[-0.00098008, 0.00184446], - [-0.00097937, 0.00185262], - [-0.00098355, 0.00184345]], - [[-0.00097346, 0.00183507], - [-0.00097343, 0.00182569], - [-0.00097549, 0.00183882]], - [[-0.00097884, 0.00183125], - [-0.00098297, 0.00184326], - [-0.00097978, 0.00184036]]]) + IN = np.array( + [ + [ + [0.32400414, 0.31483042], + [0.38269293, 0.04822304], + [0.03791266, 0.34776369], + ], + [ + [0.87113619, 0.62172854], + [0.87353969, 0.92837042], + [0.70359915, 0.49917081], + ], + [ + [0.42643583, 0.74653631], + [0.08519834, 0.35423595], + [0.34863797, 0.44895086], + ], + ] + ) + + BN_PARAMS = [ + np.array([-0.00098404, 0.00185387]), + np.array([-0.0068268, 0.0068859]), + np.array([-0.00289366, 0.00742069]), + np.array([-0.00177374, -0.00444383]), + ] + + OUT = np.array( + [ + [ + [-0.00098008, 0.00184446], + [-0.00097937, 0.00185262], + [-0.00098355, 0.00184345], + ], + [ + [-0.00097346, 0.00183507], + [-0.00097343, 0.00182569], + [-0.00097549, 0.00183882], + ], + [ + [-0.00097884, 0.00183125], + [-0.00098297, 0.00184326], + [-0.00097978, 0.00184036], + ], + ] + ) def test_batch_norm(self): params = TestBatchNormLayerClass.BN_PARAMS x = TestBatchNormLayerClass.IN y_true = TestBatchNormLayerClass.OUT - bnl = BatchNormLayer( - params[0], params[1], params[2], params[3], linear) + bnl = BatchNormLayer(params[0], params[1], params[2], params[3], linear) y = bnl.activate(x) self.assertTrue(np.allclose(y, y_true)) class TestAverageLayerClass(unittest.TestCase): - - IN = np.array([[[0.32400414, 0.31483042], - [0.38269293, 0.04822304], - [0.03791266, 0.34776369]], - [[0.87113619, 0.62172854], - [0.87353969, 0.92837042], - [0.70359915, 0.49917081]], - [[0.42643583, 0.74653631], - [0.08519834, 0.35423595], - [0.34863797, 0.44895086]]]) + IN = np.array( + [ + [ + [0.32400414, 0.31483042], + [0.38269293, 0.04822304], + [0.03791266, 0.34776369], + ], + [ + [0.87113619, 0.62172854], + [0.87353969, 0.92837042], + [0.70359915, 0.49917081], + ], + [ + [0.42643583, 0.74653631], + [0.08519834, 0.35423595], + [0.34863797, 0.44895086], + ], + ] + ) OUT_AVG = 0.46460927444444444 OUT_02 = np.array([0.55077857, 0.44537673, 0.39767252]) @@ -472,7 +608,6 @@ def test_average_layer(self): class TestReshapeLayerClass(unittest.TestCase): - IN = np.random.random((2, 3, 4)) def test_reshape_layer(self): @@ -489,7 +624,6 @@ def test_reshape_layer(self): class TestTransposeLayerClass(unittest.TestCase): - IN = np.random.random((2, 3, 4, 5)) def test_transpose_layer(self): @@ -513,18 +647,17 @@ def test_transpose_layer(self): class TestPadLayerClass(unittest.TestCase): - def test_constant_padding(self): - pl = PadLayer(width=2, axes=(0, 1), value=10.) + pl = PadLayer(width=2, axes=(0, 1), value=10.0) data = np.arange(40).reshape(5, 4, 2).astype(float) out = pl(data) self.assertEqual(out.shape, (9, 8, 2)) self.assertTrue(np.allclose(out[2:-2, 2:-2, :], data)) - self.assertTrue(np.allclose(out[:2, :, :], 10.)) - self.assertTrue(np.allclose(out[-2:, :, :], 10.)) - self.assertTrue(np.allclose(out[:, :2, :], 10.)) - self.assertTrue(np.allclose(out[:, -2:, :], 10.)) + self.assertTrue(np.allclose(out[:2, :, :], 10.0)) + self.assertTrue(np.allclose(out[-2:, :, :], 10.0)) + self.assertTrue(np.allclose(out[:, :2, :], 10.0)) + self.assertTrue(np.allclose(out[:, -2:, :], 10.0)) pl = PadLayer(width=3, axes=(2,), value=2.2) out = pl(data) @@ -536,36 +669,44 @@ def test_constant_padding(self): class TestConvolutionalLayerClass(unittest.TestCase): - def setUp(self): # 1x1 - tf_weights = np.array([[[[1., 0.5]]]]) + tf_weights = np.array([[[[1.0, 0.5]]]]) weights = np.transpose(tf_weights, axes=(2, 3, 0, 1)) weights = np.flip(np.flip(weights, axis=2), axis=3) - bias = np.array([0., 2.]) + bias = np.array([0.0, 2.0]) self.layer1x1 = ConvolutionalLayer(weights, bias) # 2x2 - tf_weights = np.array([[[[1., 0.25]], [[0., 0.]]], - [[[1., 1.]], [[1., 0.75]]]]) + tf_weights = np.array( + [[[[1.0, 0.25]], [[0.0, 0.0]]], [[[1.0, 1.0]], [[1.0, 0.75]]]] + ) weights = np.transpose(tf_weights, axes=(2, 3, 0, 1)) weights = np.flip(np.flip(weights, axis=2), axis=3) bias = np.array([0.5, 0.1]) self.layer2x2 = ConvolutionalLayer(weights, bias) # 3x3 - tf_weights = np.array([[[[1.]], [[1.]], [[1.]]], - [[[0.]], [[1.]], [[1.]]], - [[[0.]], [[0.]], [[2.]]]]) + tf_weights = np.array( + [ + [[[1.0]], [[1.0]], [[1.0]]], + [[[0.0]], [[1.0]], [[1.0]]], + [[[0.0]], [[0.0]], [[2.0]]], + ] + ) weights = np.transpose(tf_weights, axes=(2, 3, 0, 1)) weights = np.flip(np.flip(weights, axis=2), axis=3) - bias = np.array([1.]) + bias = np.array([1.0]) self.layer3x3 = ConvolutionalLayer(weights, bias) # 3x3 multi-channel - tf_weights = np.array([[[[1.], [2]], [[1], [0]], [[1], [3]]], - [[[0.], [1]], [[1.], [-4]], [[1.], [-5]]], - [[[0.], [2]], [[0.], [-2]], [[2.], [0]]]]) + tf_weights = np.array( + [ + [[[1.0], [2]], [[1], [0]], [[1], [3]]], + [[[0.0], [1]], [[1.0], [-4]], [[1.0], [-5]]], + [[[0.0], [2]], [[0.0], [-2]], [[2.0], [0]]], + ] + ) weights = np.transpose(tf_weights, axes=(2, 3, 0, 1)) weights = np.flip(np.flip(weights, axis=2), axis=3) - bias = np.array([1.]) + bias = np.array([1.0]) self.layer3x3m = ConvolutionalLayer(weights, bias) # data self.x = np.arange(20, dtype=np.float32).reshape((5, 4, 1)) @@ -573,42 +714,47 @@ def setUp(self): def test_1x1(self): out = self.layer1x1.activate(self.x) self.assertEqual(out.shape, (5, 4, 2)) - self.assertTrue(np.allclose( - out[..., 0], np.arange(20).reshape((5, 4)))) - self.assertTrue(np.allclose( - out[..., 1], np.arange(20).reshape((5, 4)) / 2. + 2)) + self.assertTrue(np.allclose(out[..., 0], np.arange(20).reshape((5, 4)))) + self.assertTrue( + np.allclose(out[..., 1], np.arange(20).reshape((5, 4)) / 2.0 + 2) + ) def test_1x1_pad_same(self): - self.layer1x1.pad = 'same' + self.layer1x1.pad = "same" out = self.layer1x1.activate(self.x) self.assertEqual(out.shape, (5, 4, 2)) def test_2x2(self): out = self.layer2x2.activate(self.x) - correct = np.array([[[9.5, 7.85], [12.5, 9.85], [15.5, 11.85]], - [[21.5, 15.85], [24.5, 17.85], [27.5, 19.85]], - [[33.5, 23.85], [36.5, 25.85], [39.5, 27.85]], - [[45.5, 31.85], [48.5, 33.85], [51.5, 35.85]]]) + correct = np.array( + [ + [[9.5, 7.85], [12.5, 9.85], [15.5, 11.85]], + [[21.5, 15.85], [24.5, 17.85], [27.5, 19.85]], + [[33.5, 23.85], [36.5, 25.85], [39.5, 27.85]], + [[45.5, 31.85], [48.5, 33.85], [51.5, 35.85]], + ] + ) self.assertEqual(out.shape, (4, 3, 2)) self.assertTrue(np.allclose(out, correct)) - self.assertTrue(np.allclose( - out[..., 0], np.arange(9.5, 57.5, 3).reshape((4, 4))[:, :-1])) - self.assertTrue(np.allclose( - out[..., 1], np.arange(7.85, 39.85, 2).reshape((4, 4))[:, :-1])) + self.assertTrue( + np.allclose(out[..., 0], np.arange(9.5, 57.5, 3).reshape((4, 4))[:, :-1]) + ) + self.assertTrue( + np.allclose(out[..., 1], np.arange(7.85, 39.85, 2).reshape((4, 4))[:, :-1]) + ) def test_2x2_pad_same(self): - self.layer2x2.pad = 'same' + self.layer2x2.pad = "same" out = self.layer2x2.activate(self.x) - correct = np.array([[[9.5, 7.85], [12.5, 9.85], - [15.5, 11.85], [10.5, 7.85]], - [[21.5, 15.85], [24.5, 17.85], - [27.5, 19.85], [18.5, 12.85]], - [[33.5, 23.85], [36.5, 25.85], - [39.5, 27.85], [26.5, 17.85]], - [[45.5, 31.85], [48.5, 33.85], - [51.5, 35.85], [34.5, 22.85]], - [[16.5, 4.10], [17.5, 4.35], - [18.5, 4.60], [19.5, 4.85]]]) + correct = np.array( + [ + [[9.5, 7.85], [12.5, 9.85], [15.5, 11.85], [10.5, 7.85]], + [[21.5, 15.85], [24.5, 17.85], [27.5, 19.85], [18.5, 12.85]], + [[33.5, 23.85], [36.5, 25.85], [39.5, 27.85], [26.5, 17.85]], + [[45.5, 31.85], [48.5, 33.85], [51.5, 35.85], [34.5, 22.85]], + [[16.5, 4.10], [17.5, 4.35], [18.5, 4.60], [19.5, 4.85]], + ] + ) self.assertEqual(out.shape, (5, 4, 2)) self.assertTrue(np.allclose(out, correct)) @@ -619,13 +765,19 @@ def test_3x3(self): self.assertTrue(np.allclose(out, correct)) def test_3x3_pad_same(self): - self.layer3x3.pad = 'same' + self.layer3x3.pad = "same" out = self.layer3x3.activate(self.x) - correct = np.array([[[[12], [16], [20], [4]], - [[29], [35], [42], [13]], - [[53], [63], [70], [25]], - [[77], [91], [98], [37]], - [[59], [75], [80], [49]]]]) + correct = np.array( + [ + [ + [[12], [16], [20], [4]], + [[29], [35], [42], [13]], + [[53], [63], [70], [25]], + [[77], [91], [98], [37]], + [[59], [75], [80], [49]], + ] + ] + ) self.assertEqual(out.shape, (5, 4, 1)) self.assertTrue(np.allclose(out, correct)) @@ -637,12 +789,16 @@ def test_3x3_multi_channel(self): self.assertEqual(out.shape, (3, 2, 1)) self.assertTrue(np.allclose(out, correct)) # with padding 'same' - self.layer3x3m.pad = 'same' + self.layer3x3m.pad = "same" out = self.layer3x3m.activate(x) - correct = np.array([[[-14], [-9], [-17], [-20]], - [[-59], [-18], [-10], [-16]], - [[-75], [14], [22], [0]], - [[-91], [46], [54], [16]], - [[-109], [-6], [-2], [36]]]) + correct = np.array( + [ + [[-14], [-9], [-17], [-20]], + [[-59], [-18], [-10], [-16]], + [[-75], [14], [22], [0]], + [[-91], [46], [54], [16]], + [[-109], [-6], [-2], [36]], + ] + ) self.assertEqual(out.shape, (5, 4, 1)) self.assertTrue(np.allclose(out, correct)) diff --git a/tests/test_ml_nn_onnx_parity.py b/tests/test_ml_nn_onnx_parity.py new file mode 100644 index 000000000..ee8e81c03 --- /dev/null +++ b/tests/test_ml_nn_onnx_parity.py @@ -0,0 +1,158 @@ +from __future__ import absolute_import, division, print_function + +import os +import unittest + +import numpy as np + + +INTENTIONAL_PARITY_BREAK_ENV = "MADMOM_ONNX_PARITY_BREAK_HARNESS" + +PARITY_TOLERANCE_POLICY = { + "primitives": { + "default": {"atol": 1e-6, "rtol": 1e-6}, + }, + "conv": { + "default": {"atol": 1e-6, "rtol": 1e-6}, + }, + "structure": { + "default": {"atol": 1e-6, "rtol": 1e-6}, + }, + "recurrent": { + "default": {"atol": 1e-6, "rtol": 1e-6}, + "online": {"atol": 1e-6, "rtol": 1e-6}, + "state": {"atol": 1e-6, "rtol": 1e-6}, + }, + "tcn": { + "default": {"atol": 1e-6, "rtol": 1e-6}, + }, + "wrapped": { + "default": {"atol": 1e-6, "rtol": 1e-6}, + }, + "runtime": { + "default": {"atol": 1e-6, "rtol": 1e-6}, + "online": {"atol": 1e-6, "rtol": 1e-6}, + }, +} + + +def parity_tolerance(family, layer_group="default"): + if family not in PARITY_TOLERANCE_POLICY: + raise KeyError("unknown ONNX parity family: %s" % family) + family_policy = PARITY_TOLERANCE_POLICY[family] + if layer_group in family_policy: + return family_policy[layer_group] + return family_policy["default"] + + +def assert_parity_close( + observed, + expected, + *, + family, + layer_group="default", + context="", +): + tolerance = parity_tolerance(family, layer_group) + atol = float(tolerance["atol"]) + rtol = float(tolerance["rtol"]) + + observed_array = np.asarray(observed, dtype=np.float32) + expected_array = np.asarray(expected, dtype=np.float32) + + if observed_array.shape != expected_array.shape: + raise AssertionError( + "ONNX parity shape mismatch family=%s layer_group=%s context=%s " + "expected_shape=%s observed_shape=%s" + % ( + family, + layer_group, + context, + expected_array.shape, + observed_array.shape, + ) + ) + + if np.allclose(observed_array, expected_array, atol=atol, rtol=rtol): + return + + abs_diff = np.abs(observed_array - expected_array) + if abs_diff.size == 0: + max_index = () + max_abs_diff = 0.0 + max_rel_diff = 0.0 + observed_value = float(observed_array.reshape(-1)[0]) + expected_value = float(expected_array.reshape(-1)[0]) + else: + flat_index = int(np.argmax(abs_diff)) + max_index = np.unravel_index(flat_index, abs_diff.shape) + max_abs_diff = float(abs_diff[max_index]) + denominator = max(abs(float(expected_array[max_index])), 1e-12) + max_rel_diff = max_abs_diff / denominator + observed_value = float(observed_array[max_index]) + expected_value = float(expected_array[max_index]) + + raise AssertionError( + "ONNX parity mismatch family=%s layer_group=%s context=%s " + "atol=%g rtol=%g max_abs_diff=%.10g max_rel_diff=%.10g " + "max_index=%s observed=%r expected=%r" + % ( + family, + layer_group, + context, + atol, + rtol, + max_abs_diff, + max_rel_diff, + max_index, + observed_value, + expected_value, + ) + ) + + +class TestMlNnOnnxParity(unittest.TestCase): + def test_parity_tolerance_policy_covers_supported_families(self): + expected_families = { + "primitives", + "conv", + "structure", + "recurrent", + "tcn", + "wrapped", + "runtime", + } + self.assertEqual(set(PARITY_TOLERANCE_POLICY), expected_families) + + def test_parity_failure_diagnostics_include_family_layer_and_context(self): + with self.assertRaises(AssertionError) as ctx: + assert_parity_close( + observed=np.asarray([0.1, 0.2], dtype=np.float32), + expected=np.asarray([0.1, 1.2], dtype=np.float32), + family="recurrent", + layer_group="online", + context="layer=lstm chunk=tail", + ) + + message = str(ctx.exception) + self.assertIn("family=recurrent", message) + self.assertIn("layer_group=online", message) + self.assertIn("layer=lstm", message) + self.assertIn("chunk=tail", message) + + @unittest.skipUnless( + os.getenv(INTENTIONAL_PARITY_BREAK_ENV) == "1", + "intentional ONNX parity break harness is disabled", + ) + def test_intentional_parity_break_harness(self): + assert_parity_close( + observed=np.asarray([0.0, 0.0], dtype=np.float32), + expected=np.asarray([0.0, 0.5], dtype=np.float32), + family="primitives", + layer_group="default", + context="intentional-break-harness", + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_ml_nn_onnx_primitives.py b/tests/test_ml_nn_onnx_primitives.py new file mode 100644 index 000000000..582b1660b --- /dev/null +++ b/tests/test_ml_nn_onnx_primitives.py @@ -0,0 +1,927 @@ +from __future__ import absolute_import, annotations, division, print_function + +import hashlib +import importlib +import importlib.util +import json +import os +import pickle +import tempfile +import unittest +from types import ModuleType, SimpleNamespace +from typing import ( + Callable, + Literal, + Optional, + Protocol, + Tuple, + Type, + Union, + cast, + final, +) + +try: + from typing import TypeAlias +except ImportError: + from typing_extensions import TypeAlias +from unittest import mock + +import numpy as np + +from tests.test_ml_nn_onnx_parity import assert_parity_close + + +TESTS_DIR = os.path.dirname(os.path.abspath(__file__)) +REPO_ROOT = os.path.dirname(TESTS_DIR) +SCRIPT_PATH = os.path.join(REPO_ROOT, "tools", "convert_models_to_onnx.py") + +ActivationFn: TypeAlias = Callable[[np.ndarray, Optional[np.ndarray]], np.ndarray] +AxisType: TypeAlias = Union[int, Tuple[int, ...], None] +AverageDType: TypeAlias = Union[Type[np.float32], np.dtype[np.float32], None] + + +def linear(x: np.ndarray, out: np.ndarray | None = None) -> np.ndarray: + if out is None or x is out: + return x + out[:] = x + return out + + +def tanh(x: np.ndarray, out: np.ndarray | None = None) -> np.ndarray: + return np.tanh(x, out) + + +def sigmoid(x: np.ndarray, out: np.ndarray | None = None) -> np.ndarray: + if out is None: + out = np.asarray(0.5 * x) + else: + if out is not x: + out[:] = x + out *= 0.5 + _ = np.tanh(out, out=out) + out += 1.0 + out *= 0.5 + return out + + +def relu(x: np.ndarray, out: np.ndarray | None = None) -> np.ndarray: + return np.maximum(x, 0, out=out) + + +def elu(x: np.ndarray, out: np.ndarray | None = None) -> np.ndarray: + if out is None: + out = x.copy() + elif out is not x: + out[:] = x[:] + mask = x < 0 + out[mask] = np.exp(x[mask]) - 1 + return out + + +def softmax(x: np.ndarray, out: np.ndarray | None = None) -> np.ndarray: + shifted = cast(np.ndarray, x - np.amax(x, axis=1, keepdims=True)) + exponentiated = np.exp(shifted) + denominator = cast(np.ndarray, np.sum(exponentiated, axis=1, keepdims=True)) + probabilities = cast(np.ndarray, exponentiated / denominator) + if out is None: + return probabilities + out[:] = probabilities + return out + + +def unsupported_activation(x: np.ndarray, out: np.ndarray | None = None) -> np.ndarray: + if out is None: + return x + out[:] = x + return out + + +for _activation in (linear, tanh, sigmoid, relu, elu, softmax, unsupported_activation): + _activation.__module__ = __name__ + + +@final +class FeedForwardLayer(object): + weights: np.ndarray + bias: np.ndarray + activation_fn: ActivationFn | None + + def __init__( + self, + weights: np.ndarray, + bias: np.ndarray, + activation_fn: ActivationFn | None = None, + ): + self.weights = np.asarray(weights, dtype=np.float32) + self.bias = np.asarray(bias, dtype=np.float32).flatten() + self.activation_fn = activation_fn + + def __call__(self, data: np.ndarray, **kwargs: object) -> np.ndarray: + del kwargs + out = cast(np.ndarray, np.matmul(data, self.weights) + self.bias) + if self.activation_fn is not None: + _ = self.activation_fn(out, out) + return out + + +@final +class BatchNormLayer(object): + beta: np.ndarray + gamma: np.ndarray + mean: np.ndarray + inv_std: np.ndarray + activation_fn: ActivationFn | None + + def __init__( + self, + beta: np.ndarray, + gamma: np.ndarray, + mean: np.ndarray, + inv_std: np.ndarray, + activation_fn: ActivationFn | None = None, + ): + self.beta = np.asarray(beta, dtype=np.float32) + self.gamma = np.asarray(gamma, dtype=np.float32) + self.mean = np.asarray(mean, dtype=np.float32) + self.inv_std = np.asarray(inv_std, dtype=np.float32) + self.activation_fn = activation_fn + + def __call__(self, data: np.ndarray, **kwargs: object) -> np.ndarray: + del kwargs + out = (data - self.mean) * (self.gamma * self.inv_std) + self.beta + if self.activation_fn is not None: + _ = self.activation_fn(out, out) + return out + + +@final +class ReshapeLayer(object): + newshape: tuple[int, ...] + order: Literal["C", "F", "A"] + + def __init__(self, newshape: tuple[int, ...], order: Literal["C", "F", "A"] = "C"): + self.newshape = newshape + self.order = order + + def __call__(self, data: np.ndarray, **kwargs: object) -> np.ndarray: + del kwargs + return np.reshape(data, self.newshape, order=self.order) + + +@final +class TransposeLayer(object): + axes: tuple[int, ...] | None + + def __init__(self, axes: tuple[int, ...] | None = None): + self.axes = axes + + def __call__(self, data: np.ndarray, **kwargs: object) -> np.ndarray: + del kwargs + return np.transpose(data, self.axes) + + +@final +class PadLayer(object): + width: int + axes: tuple[int, ...] + value: float + + def __init__(self, width: int, axes: tuple[int, ...], value: float = 0.0): + self.width = width + self.axes = axes + self.value = value + + def __call__(self, data: np.ndarray, **kwargs: object) -> np.ndarray: + del kwargs + shape = list(data.shape) + data_idxs = [slice(None) for _ in range(len(shape))] + for axis in self.axes: + shape[axis] += self.width * 2 + data_idxs[axis] = slice(self.width, -self.width) + padded = np.full(tuple(shape), self.value, dtype=np.float32) + padded[tuple(data_idxs)] = data + return padded + + +@final +class AverageLayer(object): + axis: AxisType + dtype: AverageDType + keepdims: bool + + def __init__( + self, + axis: AxisType = None, + dtype: AverageDType = None, + keepdims: bool = False, + ): + self.axis = axis + self.dtype = dtype + self.keepdims = keepdims + + def __call__(self, data: np.ndarray, **kwargs: object) -> np.ndarray: + del kwargs + return cast( + np.ndarray, + np.mean(data, axis=self.axis, dtype=self.dtype, keepdims=self.keepdims), + ) + + +@final +class PrimitiveNetwork(object): + layers: list[object] + + def __init__(self, layers: list[object]): + self.layers = layers + + def process(self, data: np.ndarray) -> np.ndarray: + if data.ndim < 2: + data = np.array(data, subok=True, copy=False, ndmin=2) + output = data + for layer in self.layers: + output = cast(Callable[[np.ndarray], np.ndarray], layer)(output) + return np.asarray(output).squeeze() + + +@final +class UnsupportedLayer(object): + def __call__(self, data: np.ndarray, **kwargs: object) -> np.ndarray: + del kwargs + return data + + +class ConverterModule(Protocol): + DEFAULT_OPSET: int + + def convert_manifest_entries( + self, + manifest: dict[str, object], + models_dir: str, + converter: object = None, + run_onnx_checker: bool = True, + ) -> dict[str, object]: ... + + +converter_module: ConverterModule | None = None + + +def _load_converter_module() -> ConverterModule: + spec = importlib.util.spec_from_file_location("convert_models_to_onnx", SCRIPT_PATH) + if spec is None or spec.loader is None: + raise RuntimeError("failed to load converter module spec") + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + module_object = cast(object, module) + return cast(ConverterModule, module_object) + + +def _get_converter_module() -> ConverterModule: + global converter_module + if converter_module is None: + converter_module = _load_converter_module() + return converter_module + + +def _sha256_digest(path: str) -> str: + digest = hashlib.sha256() + with open(path, "rb") as infile: + while True: + chunk = infile.read(1024 * 1024) + if not chunk: + break + digest.update(chunk) + return "sha256:%s" % digest.hexdigest() + + +def _single_convertible_manifest( + module: ConverterModule, + source_rel: str, + target_rel: str, + source_path: str, +) -> dict[str, object]: + artifact: dict[str, object] = { + "source_pkl": source_rel, + "target_onnx": target_rel, + "hash": _sha256_digest(source_path), + "opset": module.DEFAULT_OPSET, + "status": "nn_convertible", + "convertible": True, + "conversion_mode": "direct_nn", + } + return {"schema_version": 1, "artifacts": [artifact]} + + +@final +class _FakeTensorProto(object): + FLOAT: int = 1 + INT64: int = 7 + + +@final +class _FakeNumpyHelper(object): + @staticmethod + def from_array(array: np.ndarray, name: str | None = None) -> dict[str, object]: + value = np.asarray(array) + return { + "name": name, + "array": value.tolist(), + "dtype": str(value.dtype), + } + + +@final +class _FakeHelper(object): + @staticmethod + def make_tensor_value_info( + name: str, elem_type: int, shape: object + ) -> dict[str, object]: + return {"name": name, "elem_type": elem_type, "shape": shape} + + @staticmethod + def make_node( + op_type: str, + inputs: list[str], + outputs: list[str], + name: str | None = None, + **kwargs: object, + ) -> dict[str, object]: + return { + "op_type": op_type, + "inputs": list(inputs), + "outputs": list(outputs), + "name": name, + "attrs": kwargs, + } + + @staticmethod + def make_graph( + nodes: list[dict[str, object]], + name: str, + inputs: list[dict[str, object]], + outputs: list[dict[str, object]], + initializer: list[dict[str, object]] | None = None, + ) -> dict[str, object]: + return { + "name": name, + "nodes": list(nodes), + "inputs": list(inputs), + "outputs": list(outputs), + "initializer": list(initializer or []), + } + + @staticmethod + def make_model( + graph: dict[str, object], + producer_name: str, + opset_imports: list[dict[str, object]], + ) -> object: + return _FakeModel( + { + "graph": graph, + "producer_name": producer_name, + "opset_imports": opset_imports, + } + ) + + @staticmethod + def make_opsetid(domain: str, version: int) -> dict[str, object]: + return {"domain": domain, "version": version} + + +@final +class _FakeModel(object): + payload: dict[str, object] + + def __init__(self, payload: dict[str, object]): + self.payload = payload + + def SerializeToString(self, deterministic: bool = False) -> bytes: + return json.dumps( + self.payload, + sort_keys=deterministic, + separators=(",", ":"), + ).encode("utf-8") + + +@final +class _FakeChecker(object): + @staticmethod + def check_model(model_path: str) -> None: + del model_path + + +def _fake_onnx_module() -> object: + return SimpleNamespace( + TensorProto=_FakeTensorProto, + helper=_FakeHelper(), + numpy_helper=_FakeNumpyHelper(), + checker=_FakeChecker(), + ) + + +def _model_payload(model_bytes: bytes) -> dict[str, object]: + return cast(dict[str, object], json.loads(model_bytes.decode("utf-8"))) + + +def _to_int_tuple(value: np.ndarray) -> tuple[int, ...]: + flattened = np.asarray(value, dtype=np.int64).reshape(-1) + return tuple(int(item) for item in flattened) + + +def _evaluate_fake_onnx( + payload: dict[str, object], input_array: np.ndarray +) -> np.ndarray: + graph = cast(dict[str, object], payload["graph"]) + initializer_items = cast(list[dict[str, object]], graph["initializer"]) + values: dict[str, np.ndarray] = {"input": np.asarray(input_array, dtype=np.float32)} + + for item in initializer_items: + dtype = np.dtype(cast(str, item["dtype"])) + values[cast(str, item["name"])] = np.asarray(item["array"], dtype=dtype) + + nodes = cast(list[dict[str, object]], graph["nodes"]) + for node in nodes: + op_type = cast(str, node["op_type"]) + inputs = [values[name] for name in cast(list[str], node["inputs"])] + attrs = cast(dict[str, object], node.get("attrs", {})) + output_name = cast(list[str], node["outputs"])[0] + + if op_type == "MatMul": + values[output_name] = np.matmul(inputs[0], inputs[1]) + elif op_type == "Add": + values[output_name] = inputs[0] + inputs[1] + elif op_type == "Mul": + values[output_name] = inputs[0] * inputs[1] + elif op_type == "Tanh": + values[output_name] = np.tanh(inputs[0]) + elif op_type == "Sigmoid": + values[output_name] = 1.0 / (1.0 + np.exp(-inputs[0])) + elif op_type == "Relu": + values[output_name] = np.maximum(inputs[0], 0) + elif op_type == "Elu": + alpha = float(cast(float, attrs.get("alpha", 1.0))) + values[output_name] = np.where( + inputs[0] > 0, + inputs[0], + alpha * (np.exp(inputs[0]) - 1.0), + ) + elif op_type == "Softmax": + axis = int(cast(int, attrs.get("axis", 1))) + shifted = cast( + np.ndarray, + inputs[0] - np.max(inputs[0], axis=axis, keepdims=True), + ) + exponentiated = np.exp(shifted) + denominator = cast( + np.ndarray, np.sum(exponentiated, axis=axis, keepdims=True) + ) + values[output_name] = cast(np.ndarray, exponentiated / denominator) + elif op_type == "Reshape": + target_shape = _to_int_tuple(inputs[1]) + values[output_name] = np.reshape(inputs[0], target_shape) + elif op_type == "Transpose": + perm = attrs.get("perm") + if perm is None: + values[output_name] = np.transpose(inputs[0]) + else: + values[output_name] = np.transpose( + inputs[0], axes=tuple(cast(list[int], perm)) + ) + elif op_type == "Pad": + pads = list(_to_int_tuple(inputs[1])) + rank = len(inputs[0].shape) + begins = pads[:rank] + ends = pads[rank:] + pad_width = list(zip(begins, ends)) + constant_array = np.asarray(inputs[2], dtype=np.float32).reshape(()) + constant_value = float(constant_array.item()) + values[output_name] = np.pad( + inputs[0], + pad_width=pad_width, + mode="constant", + constant_values=constant_value, + ) + elif op_type == "ReduceMean": + keepdims = bool(int(cast(int, attrs.get("keepdims", 1)))) + if len(inputs) == 1: + values[output_name] = np.mean(inputs[0], keepdims=keepdims) + else: + axes_values = _to_int_tuple(inputs[1]) + values[output_name] = np.mean( + inputs[0], + axis=axes_values, + keepdims=keepdims, + ) + elif op_type == "Cast": + values[output_name] = inputs[0].astype(np.float32) + elif op_type == "Squeeze": + if len(inputs) == 1: + values[output_name] = np.squeeze(inputs[0]) + else: + axes_values = _to_int_tuple(inputs[1]) + values[output_name] = np.squeeze(inputs[0], axis=axes_values) + else: + raise AssertionError("unsupported fake op: %s" % op_type) + + output_name = cast(list[dict[str, object]], graph["outputs"])[0]["name"] + return values[cast(str, output_name)] + + +class TestMlNnOnnxPrimitives(unittest.TestCase): + def _patched_import(self): + fake_onnx = _fake_onnx_module() + original_import = importlib.import_module + + def side_effect(name: str, package: str | None = None) -> ModuleType | object: + del package + if name == "onnx": + return fake_onnx + return original_import(name) + + return mock.patch("importlib.import_module", side_effect=side_effect) + + def test_primitive_fixture_converts_deterministically_and_matches_reference(self): + module = _get_converter_module() + + model = PrimitiveNetwork( + layers=[ + FeedForwardLayer( + weights=np.array( + [ + [0.1, -0.4, 0.8, 1.0], + [0.2, 0.5, -0.2, 0.1], + [-0.6, 0.7, 0.3, -0.9], + ], + dtype=np.float32, + ), + bias=np.array([0.4, -0.2, 0.3, 0.1], dtype=np.float32), + activation_fn=relu, + ), + BatchNormLayer( + beta=np.array([0.01, -0.02, 0.03, 0.04], dtype=np.float32), + gamma=np.array([1.1, 0.9, 1.2, 0.7], dtype=np.float32), + mean=np.array([0.3, -0.1, 0.5, 0.2], dtype=np.float32), + inv_std=np.array([0.8, 1.1, 0.6, 0.9], dtype=np.float32), + activation_fn=tanh, + ), + ReshapeLayer(newshape=(2, 2, 2), order="C"), + TransposeLayer(axes=(1, 0, 2)), + PadLayer(width=1, axes=(2,), value=0.25), + AverageLayer(axis=2, keepdims=False), + FeedForwardLayer( + weights=np.array( + [[1.2, -0.4, 0.5], [0.3, 0.7, -1.1]], + dtype=np.float32, + ), + bias=np.array([0.2, -0.3, 0.1], dtype=np.float32), + activation_fn=softmax, + ), + ] + ) + input_data = np.array( + [[1.0, -0.5, 0.3], [0.2, 0.4, -1.0]], + dtype=np.float32, + ) + expected = model.process(input_data) + + with tempfile.TemporaryDirectory() as temp_dir: + source_rel = "fixtures/primitives.pkl" + target_rel = "fixtures/primitives.onnx" + source_path = os.path.join(temp_dir, "fixtures", "primitives.pkl") + target_path = os.path.join(temp_dir, "fixtures", "primitives.onnx") + os.makedirs(os.path.dirname(source_path), exist_ok=True) + with open(source_path, "wb") as outfile: + pickle.dump(model, outfile) + + manifest = _single_convertible_manifest( + module=module, + source_rel=source_rel, + target_rel=target_rel, + source_path=source_path, + ) + + with self._patched_import(): + first = module.convert_manifest_entries( + manifest=manifest, + models_dir=temp_dir, + run_onnx_checker=False, + ) + second = module.convert_manifest_entries( + manifest=manifest, + models_dir=temp_dir, + run_onnx_checker=False, + ) + + first_result = cast(list[dict[str, object]], first["results"])[0] + second_result = cast(list[dict[str, object]], second["results"])[0] + self.assertEqual(first_result["status"], "converted") + self.assertEqual(second_result["status"], "converted") + self.assertEqual(first_result["output_hash"], second_result["output_hash"]) + + with open(target_path, "rb") as infile: + first_bytes = infile.read() + with open(target_path, "rb") as infile: + second_bytes = infile.read() + self.assertEqual(first_bytes, second_bytes) + + model_payload = _model_payload(first_bytes) + graph_payload = cast(dict[str, object], model_payload["graph"]) + graph_inputs = cast(list[dict[str, object]], graph_payload["inputs"]) + graph_outputs = cast(list[dict[str, object]], graph_payload["outputs"]) + self.assertEqual( + graph_inputs[0]["shape"], + ["input_dim_0", "input_dim_1"], + ) + self.assertEqual(graph_outputs[0]["shape"], ["output_0_dim_0"]) + + observed = _evaluate_fake_onnx(model_payload, input_data) + assert_parity_close( + observed, + expected, + family="primitives", + layer_group="default", + context="fixture=primitives", + ) + + op_types = [ + cast(str, node["op_type"]) + for node in cast( + list[dict[str, object]], + cast(dict[str, object], model_payload["graph"])["nodes"], + ) + ] + self.assertIn("Reshape", op_types) + self.assertIn("Transpose", op_types) + self.assertIn("Pad", op_types) + self.assertIn("ReduceMean", op_types) + self.assertIn("Softmax", op_types) + self.assertIn("Squeeze", op_types) + + def test_supported_activation_mappings_emit_expected_ops(self): + module = _get_converter_module() + input_data = np.array( + [[0.2, -0.1], [1.0, 0.5]], + dtype=np.float32, + ) + cases: list[tuple[str, ActivationFn, str | None]] = [ + ("linear", linear, None), + ("tanh", tanh, "Tanh"), + ("sigmoid", sigmoid, "Sigmoid"), + ("relu", relu, "Relu"), + ("elu", elu, "Elu"), + ("softmax", softmax, "Softmax"), + ] + + for activation_name, activation_fn, expected_op in cases: + with self.subTest(activation=activation_name): + model = PrimitiveNetwork( + layers=[ + FeedForwardLayer( + weights=np.array( + [[1.0, -0.5], [0.3, 0.7]], + dtype=np.float32, + ), + bias=np.array([0.1, -0.2], dtype=np.float32), + activation_fn=activation_fn, + ) + ] + ) + expected = model.process(input_data) + + with tempfile.TemporaryDirectory() as temp_dir: + source_rel = "fixtures/%s.pkl" % activation_name + target_rel = "fixtures/%s.onnx" % activation_name + source_path = os.path.join( + temp_dir, "fixtures", "%s.pkl" % activation_name + ) + target_path = os.path.join( + temp_dir, "fixtures", "%s.onnx" % activation_name + ) + os.makedirs(os.path.dirname(source_path), exist_ok=True) + with open(source_path, "wb") as outfile: + pickle.dump(model, outfile) + + manifest = _single_convertible_manifest( + module=module, + source_rel=source_rel, + target_rel=target_rel, + source_path=source_path, + ) + + with self._patched_import(): + report = module.convert_manifest_entries( + manifest=manifest, + models_dir=temp_dir, + run_onnx_checker=False, + ) + + result = cast(list[dict[str, object]], report["results"])[0] + self.assertEqual(result["status"], "converted") + with open(target_path, "rb") as infile: + model_bytes = infile.read() + + model_payload = _model_payload(model_bytes) + observed = _evaluate_fake_onnx(model_payload, input_data) + assert_parity_close( + observed, + expected, + family="primitives", + layer_group="default", + context="activation=%s" % activation_name, + ) + + op_types = [ + cast(str, node["op_type"]) + for node in cast( + list[dict[str, object]], + cast(dict[str, object], model_payload["graph"])["nodes"], + ) + ] + activation_ops = {"Tanh", "Sigmoid", "Relu", "Elu", "Softmax"} + if expected_op is None: + self.assertFalse(activation_ops.intersection(op_types)) + else: + self.assertIn(expected_op, op_types) + + def test_output_squeeze_parity_matches_scalar_reference(self): + module = _get_converter_module() + model = PrimitiveNetwork( + layers=[ + AverageLayer(axis=None, keepdims=False), + ] + ) + input_data = np.array( + [[1.0, 3.0, -2.0], [2.0, 0.0, 4.0]], + dtype=np.float32, + ) + expected = model.process(input_data) + + with tempfile.TemporaryDirectory() as temp_dir: + source_rel = "fixtures/squeeze_scalar.pkl" + target_rel = "fixtures/squeeze_scalar.onnx" + source_path = os.path.join(temp_dir, "fixtures", "squeeze_scalar.pkl") + target_path = os.path.join(temp_dir, "fixtures", "squeeze_scalar.onnx") + os.makedirs(os.path.dirname(source_path), exist_ok=True) + with open(source_path, "wb") as outfile: + pickle.dump(model, outfile) + + manifest = _single_convertible_manifest( + module=module, + source_rel=source_rel, + target_rel=target_rel, + source_path=source_path, + ) + + with self._patched_import(): + report = module.convert_manifest_entries( + manifest=manifest, + models_dir=temp_dir, + run_onnx_checker=False, + ) + + result = cast(list[dict[str, object]], report["results"])[0] + self.assertEqual(result["status"], "converted") + with open(target_path, "rb") as infile: + model_bytes = infile.read() + + model_payload = _model_payload(model_bytes) + observed = _evaluate_fake_onnx(model_payload, input_data) + self.assertEqual(np.asarray(observed).ndim, 0) + self.assertEqual(np.asarray(expected).ndim, 0) + assert_parity_close( + observed, + expected, + family="primitives", + layer_group="default", + context="fixture=squeeze_scalar", + ) + + op_types = [ + cast(str, node["op_type"]) + for node in cast( + list[dict[str, object]], + cast(dict[str, object], model_payload["graph"])["nodes"], + ) + ] + self.assertIn("ReduceMean", op_types) + self.assertIn("Squeeze", op_types) + + def test_unknown_activation_returns_actionable_conversion_error(self): + module = _get_converter_module() + + model = PrimitiveNetwork( + layers=[ + FeedForwardLayer( + weights=np.array([[1.0], [2.0]], dtype=np.float32), + bias=np.array([0.5], dtype=np.float32), + activation_fn=unsupported_activation, + ) + ] + ) + + with tempfile.TemporaryDirectory() as temp_dir: + source_rel = "fixtures/unsupported_activation.pkl" + target_rel = "fixtures/unsupported_activation.onnx" + source_path = os.path.join( + temp_dir, "fixtures", "unsupported_activation.pkl" + ) + os.makedirs(os.path.dirname(source_path), exist_ok=True) + with open(source_path, "wb") as outfile: + pickle.dump(model, outfile) + + manifest = _single_convertible_manifest( + module=module, + source_rel=source_rel, + target_rel=target_rel, + source_path=source_path, + ) + + with self._patched_import(): + report = module.convert_manifest_entries( + manifest=manifest, + models_dir=temp_dir, + run_onnx_checker=False, + ) + + result = cast(list[dict[str, object]], report["results"])[0] + self.assertEqual(result["status"], "conversion_error") + self.assertIn("UnsupportedActivationError", cast(str, result["error"])) + self.assertIn("unsupported_activation", cast(str, result["error"])) + self.assertIn("supported activations", cast(str, result["error"])) + + def test_unsupported_reshape_layout_returns_actionable_conversion_error(self): + module = _get_converter_module() + + model = PrimitiveNetwork( + layers=[ + ReshapeLayer(newshape=(4, 1), order="F"), + ] + ) + + with tempfile.TemporaryDirectory() as temp_dir: + source_rel = "fixtures/unsupported_layout.pkl" + target_rel = "fixtures/unsupported_layout.onnx" + source_path = os.path.join(temp_dir, "fixtures", "unsupported_layout.pkl") + os.makedirs(os.path.dirname(source_path), exist_ok=True) + with open(source_path, "wb") as outfile: + pickle.dump(model, outfile) + + manifest = _single_convertible_manifest( + module=module, + source_rel=source_rel, + target_rel=target_rel, + source_path=source_path, + ) + + with self._patched_import(): + report = module.convert_manifest_entries( + manifest=manifest, + models_dir=temp_dir, + run_onnx_checker=False, + ) + + result = cast(list[dict[str, object]], report["results"])[0] + self.assertEqual(result["status"], "conversion_error") + self.assertIn("UnsupportedLayoutError", cast(str, result["error"])) + self.assertIn("order='C'", cast(str, result["error"])) + + def test_unsupported_layer_returns_actionable_converter_unavailable(self): + module = _get_converter_module() + + model = PrimitiveNetwork( + layers=[ + UnsupportedLayer(), + ] + ) + + with tempfile.TemporaryDirectory() as temp_dir: + source_rel = "fixtures/unsupported_layer.pkl" + target_rel = "fixtures/unsupported_layer.onnx" + source_path = os.path.join(temp_dir, "fixtures", "unsupported_layer.pkl") + os.makedirs(os.path.dirname(source_path), exist_ok=True) + with open(source_path, "wb") as outfile: + pickle.dump(model, outfile) + + manifest = _single_convertible_manifest( + module=module, + source_rel=source_rel, + target_rel=target_rel, + source_path=source_path, + ) + + with self._patched_import(): + report = module.convert_manifest_entries( + manifest=manifest, + models_dir=temp_dir, + run_onnx_checker=False, + ) + + result = cast(list[dict[str, object]], report["results"])[0] + self.assertEqual(result["status"], "converter_unavailable") + self.assertIn("NotImplementedError", cast(str, result["error"])) + self.assertIn("UnsupportedLayer", cast(str, result["error"])) + + +if __name__ == "__main__": + _ = unittest.main() diff --git a/tests/test_ml_nn_onnx_tcn.py b/tests/test_ml_nn_onnx_tcn.py new file mode 100644 index 000000000..9dee041c3 --- /dev/null +++ b/tests/test_ml_nn_onnx_tcn.py @@ -0,0 +1,537 @@ +from __future__ import absolute_import, annotations, division, print_function + +import importlib +import unittest +from types import ModuleType +from typing import Any, cast, final +from unittest import mock + +import numpy as np + +from tests.test_ml_nn_onnx_parity import assert_parity_close +from tests.test_ml_nn_onnx_primitives import _fake_onnx_module, _get_converter_module + + +def linear(x: np.ndarray, out: np.ndarray | None = None) -> np.ndarray: + if out is None or out is x: + return x + out[:] = x + return out + + +def relu(x: np.ndarray, out: np.ndarray | None = None) -> np.ndarray: + res = np.maximum(x, 0) + if out is not None: + out[:] = res + return out + return res + + +@final +class ConvolutionalLayer(object): + def __init__(self, weights, bias, stride=None, pad="valid", activation_fn=None): + self.weights = np.asarray(weights, dtype=np.float32) + self.bias = np.asarray(bias, dtype=np.float32) + self.stride = stride + self.pad = pad + self.activation_fn = activation_fn + + +@final +class FeedForwardLayer(object): + def __init__(self, weights, bias, activation_fn=None): + self.weights = np.asarray(weights, dtype=np.float32) + self.bias = np.asarray(bias, dtype=np.float32) + self.activation_fn = activation_fn + + +@final +class TCNBlock(object): + def __init__( + self, + dilated_conv, + dilation_rate, + activation_fn=None, + skip_conv=None, + residual_conv=None, + ): + self.dilated_conv = dilated_conv + self.dilation_rate = dilation_rate + self.activation_fn = activation_fn + self.skip_conv = skip_conv + self.residual_conv = residual_conv + + +@final +class TCNLayer(object): + def __init__(self, tcn_blocks, activation_fn=None, skip_connections=False): + self.tcn_blocks = tcn_blocks + self.skip_connections = skip_connections + self.activation_fn = activation_fn + + +@final +class LayerContainer(object): + def __init__(self, layers): + self.layers = list(layers) + + +def _to_int_tuple(value: np.ndarray) -> tuple[int, ...]: + flattened = np.asarray(value, dtype=np.int64).reshape(-1) + return tuple(int(item) for item in flattened) + + +def _evaluate_conv_with_padding_and_dilation( + x: np.ndarray, + weights: np.ndarray, + bias: np.ndarray, + *, + strides: tuple[int, int], + pads: tuple[int, int, int, int], + dilations: tuple[int, int], +) -> np.ndarray: + n, c_in, height, width = x.shape + c_out, c_weights, kernel_height, kernel_width = weights.shape + if c_weights != c_in: + raise AssertionError("conv channel mismatch") + + stride_h, stride_w = strides + dilation_h, dilation_w = dilations + pad_top, pad_left, pad_bottom, pad_right = pads + effective_kernel_h = (kernel_height - 1) * dilation_h + 1 + effective_kernel_w = (kernel_width - 1) * dilation_w + 1 + + padded = np.pad( + x, + ((0, 0), (0, 0), (pad_top, pad_bottom), (pad_left, pad_right)), + mode="constant", + ) + out_height = (padded.shape[2] - effective_kernel_h) // stride_h + 1 + out_width = (padded.shape[3] - effective_kernel_w) // stride_w + 1 + out = np.zeros((n, c_out, out_height, out_width), dtype=np.float32) + + for batch_idx in range(n): + for out_channel in range(c_out): + for out_row in range(out_height): + start_row = out_row * stride_h + for out_col in range(out_width): + start_col = out_col * stride_w + total = float(bias[out_channel]) + for in_channel in range(c_in): + for kernel_row in range(kernel_height): + source_row = start_row + kernel_row * dilation_h + for kernel_col in range(kernel_width): + source_col = start_col + kernel_col * dilation_w + total += ( + padded[ + batch_idx, in_channel, source_row, source_col + ] + * weights[ + out_channel, in_channel, kernel_row, kernel_col + ] + ) + out[batch_idx, out_channel, out_row, out_col] = total + return out + + +def _evaluate_fake_tcn_outputs( + payload: dict[str, object], input_array: np.ndarray +) -> dict[str, np.ndarray]: + graph = cast(dict[str, object], payload["graph"]) + initializers = cast(list[dict[str, object]], graph["initializer"]) + values: dict[str, np.ndarray] = {"input": np.asarray(input_array, dtype=np.float32)} + + for item in initializers: + dtype = np.dtype(cast(str, item["dtype"])) + values[cast(str, item["name"])] = np.asarray(item["array"], dtype=dtype) + + nodes = cast(list[dict[str, object]], graph["nodes"]) + for node in nodes: + op_type = cast(str, node["op_type"]) + inputs = [values[name] for name in cast(list[str], node["inputs"])] + attrs = cast(dict[str, object], node.get("attrs", {})) + output_name = cast(list[str], node["outputs"])[0] + + if op_type == "Unsqueeze": + axes = list(_to_int_tuple(inputs[1])) + output = inputs[0] + for axis in sorted(axes): + normalized_axis = axis + if normalized_axis < 0: + normalized_axis += output.ndim + 1 + output = np.expand_dims(output, axis=normalized_axis) + values[output_name] = output + elif op_type == "Squeeze": + if len(inputs) == 1: + values[output_name] = np.squeeze(inputs[0]) + else: + axes = _to_int_tuple(inputs[1]) + values[output_name] = np.squeeze(inputs[0], axis=axes) + elif op_type == "Transpose": + perm = attrs.get("perm") + if perm is None: + values[output_name] = np.transpose(inputs[0]) + else: + values[output_name] = np.transpose( + inputs[0], axes=tuple(cast(list[int], perm)) + ) + elif op_type == "Conv": + stride_attr = attrs.get("strides", [1, 1]) + pad_attr = attrs.get("pads", [0, 0, 0, 0]) + dilation_attr = attrs.get("dilations", [1, 1]) + values[output_name] = _evaluate_conv_with_padding_and_dilation( + np.asarray(inputs[0], dtype=np.float32), + np.asarray(inputs[1], dtype=np.float32), + np.asarray(inputs[2], dtype=np.float32).reshape(-1), + strides=cast(tuple[int, int], tuple(cast(list[int], stride_attr))), + pads=cast(tuple[int, int, int, int], tuple(cast(list[int], pad_attr))), + dilations=cast(tuple[int, int], tuple(cast(list[int], dilation_attr))), + ) + elif op_type == "MatMul": + values[output_name] = np.matmul(inputs[0], inputs[1]) + elif op_type == "Concat": + axis = int(cast(int, attrs["axis"])) + values[output_name] = np.concatenate(inputs, axis=axis) + elif op_type == "Relu": + values[output_name] = np.maximum(inputs[0], 0) + elif op_type == "Add": + values[output_name] = inputs[0] + inputs[1] + else: + raise AssertionError("unsupported fake op in tcn test: %s" % op_type) + + outputs: dict[str, np.ndarray] = {} + output_infos = cast(list[dict[str, object]], graph["outputs"]) + for output_info in output_infos: + name = cast(str, output_info["name"]) + outputs[name] = values[name] + return outputs + + +def _evaluate_tcn_block_reference( + input_data: np.ndarray, block: TCNBlock +) -> tuple[np.ndarray, np.ndarray]: + conv = cast(ConvolutionalLayer, block.dilated_conv) + weights = np.asarray(conv.weights, dtype=np.float32) + bias = np.asarray(conv.bias, dtype=np.float32).reshape(-1) + rate = int(block.dilation_rate) + + in_data = np.asarray(input_data, dtype=np.float32) + time_steps, freq_bins, in_channels = in_data.shape + if freq_bins != 1: + raise AssertionError("reference supports only freq_bins == 1") + out_channels = weights.shape[1] + out = np.zeros((time_steps, 1, out_channels), dtype=np.float32) + + for time_index in range(time_steps): + for out_channel in range(out_channels): + total = float(bias[out_channel]) + for in_channel in range(in_channels): + for kernel_index in range(weights.shape[3]): + source_index = time_index - kernel_index * rate + if source_index < 0: + continue + total += ( + in_data[source_index, 0, in_channel] + * weights[in_channel, out_channel, 0, kernel_index] + ) + out[time_index, 0, out_channel] = total + + if block.activation_fn is not None: + out = np.maximum(out, 0) + skip = out + residual = in_data + skip + return residual, skip + + +class TestTCNConversion(unittest.TestCase): + def _patched_import(self): + fake_onnx = _fake_onnx_module() + original_import = importlib.import_module + + def side_effect(name: str, package: str | None = None) -> ModuleType | object: + del package + if name == "onnx": + return fake_onnx + return original_import(name) + + return mock.patch("importlib.import_module", side_effect=side_effect) + + def test_tcn_block_single_conv(self): + conv = ConvolutionalLayer( + weights=np.ones((1, 4, 1, 3)), bias=np.zeros(4), pad="valid" + ) + block = TCNBlock(dilated_conv=conv, dilation_rate=2, activation_fn=relu) + + converter: Any = _get_converter_module() + with self._patched_import(): + onnx_model: Any = converter._build_onnx_model_from_primitive_layers( + [block], opset=18 + ) + self.assertIsNotNone(onnx_model) + # Check that we have a Conv node with dilations + graph = cast(dict[str, Any], onnx_model.payload["graph"]) + nodes = cast(list[dict[str, Any]], graph["nodes"]) + conv_nodes = [n for n in nodes if n["op_type"] == "Conv"] + self.assertEqual(len(conv_nodes), 1) + self.assertEqual(conv_nodes[0]["attrs"]["dilations"], [2, 1]) # dilations + + def test_tcn_layer(self): + conv = ConvolutionalLayer( + weights=np.ones((1, 4, 1, 3)), bias=np.zeros(4), pad="valid" + ) + block = TCNBlock(dilated_conv=conv, dilation_rate=1, activation_fn=relu) + layer = TCNLayer( + tcn_blocks=[block, block], activation_fn=relu, skip_connections=True + ) + + converter: Any = _get_converter_module() + with self._patched_import(): + onnx_model: Any = converter._build_onnx_model_from_primitive_layers( + [layer], opset=18 + ) + self.assertIsNotNone(onnx_model) + # 2 blocks * 1 conv = 2 Conv nodes + graph = cast(dict[str, Any], onnx_model.payload["graph"]) + nodes = cast(list[dict[str, Any]], graph["nodes"]) + conv_nodes = [n for n in nodes if n["op_type"] == "Conv"] + self.assertEqual(len(conv_nodes), 2) + + def test_tcn_block_multi_conv(self): + conv1 = ConvolutionalLayer( + weights=np.ones((1, 2, 1, 3)), bias=np.zeros(2), pad="valid" + ) + conv2 = ConvolutionalLayer( + weights=np.ones((1, 2, 1, 3)), bias=np.zeros(2), pad="valid" + ) + block = TCNBlock( + dilated_conv=[conv1, conv2], dilation_rate=[1, 2], activation_fn=relu + ) + + converter: Any = _get_converter_module() + with self._patched_import(): + onnx_model: Any = converter._build_onnx_model_from_primitive_layers( + [block], opset=18 + ) + self.assertIsNotNone(onnx_model) + # 2 convs = 2 Conv nodes + 1 Concat node + graph = cast(dict[str, Any], onnx_model.payload["graph"]) + nodes = cast(list[dict[str, Any]], graph["nodes"]) + conv_nodes = [n for n in nodes if n["op_type"] == "Conv"] + self.assertEqual(len(conv_nodes), 2) + concat_nodes = [n for n in nodes if n["op_type"] == "Concat"] + self.assertTrue(len(concat_nodes) >= 1) + + def test_tcn_block_residual_path(self): + conv = ConvolutionalLayer( + weights=np.ones((1, 4, 1, 3)), bias=np.zeros(4), pad="valid" + ) + res_conv = ConvolutionalLayer( + weights=np.ones((1, 4, 1, 1)), bias=np.zeros(4), pad="valid" + ) + block = TCNBlock( + dilated_conv=conv, + dilation_rate=1, + activation_fn=relu, + residual_conv=res_conv, + ) + + converter: Any = _get_converter_module() + with self._patched_import(): + onnx_model: Any = converter._build_onnx_model_from_primitive_layers( + [block], opset=18 + ) + self.assertIsNotNone(onnx_model) + graph = cast(dict[str, Any], onnx_model.payload["graph"]) + nodes = cast(list[dict[str, Any]], graph["nodes"]) + # Should have 2 Conv nodes (one for dilated_conv, one for residual_conv) + # plus one for skip_conv (none here) + # and one Add node for residual summation + conv_nodes = [n for n in nodes if n["op_type"] == "Conv"] + self.assertEqual(len(conv_nodes), 2) + add_nodes = [n for n in nodes if n["op_type"] == "Add"] + self.assertTrue(len(add_nodes) >= 1) + + def test_tcn_block_skip_conv(self): + conv = ConvolutionalLayer( + weights=np.ones((1, 4, 1, 3)), bias=np.zeros(4), pad="valid" + ) + skip_conv = ConvolutionalLayer( + weights=np.ones((4, 8, 1, 1)), bias=np.zeros(8), pad="valid" + ) + block = TCNBlock( + dilated_conv=conv, dilation_rate=1, activation_fn=relu, skip_conv=skip_conv + ) + + converter: Any = _get_converter_module() + with self._patched_import(): + onnx_model: Any = converter._build_onnx_model_from_primitive_layers( + [block], opset=18 + ) + self.assertIsNotNone(onnx_model) + graph = cast(dict[str, Any], onnx_model.payload["graph"]) + nodes = cast(list[dict[str, Any]], graph["nodes"]) + # 1 dilated_conv + 1 skip_conv = 2 Conv nodes + conv_nodes = [n for n in nodes if n["op_type"] == "Conv"] + self.assertEqual(len(conv_nodes), 2) + + def test_tcn_block_parity_residual_and_skip_outputs(self): + conv = ConvolutionalLayer( + weights=np.array( + [ + [[[0.2, -0.1, 0.4], [0.0, 0.0, 0.0]]], + [[[0.1, 0.3, -0.2], [0.0, 0.0, 0.0]]], + ], + dtype=np.float32, + ).transpose(0, 2, 1, 3), + bias=np.array([0.05, -0.02], dtype=np.float32), + pad="valid", + ) + block = TCNBlock(dilated_conv=conv, dilation_rate=2, activation_fn=relu) + model = LayerContainer([block]) + input_data = np.array( + [ + [[0.2, -0.1]], + [[0.0, 0.3]], + [[-0.4, 0.5]], + [[0.7, -0.2]], + [[0.1, 0.4]], + [[-0.3, 0.6]], + ], + dtype=np.float32, + ) + expected_residual, expected_skip = _evaluate_tcn_block_reference( + input_data, block + ) + + converter: Any = _get_converter_module() + with self._patched_import(): + onnx_model: Any = converter._build_onnx_model_from_primitive_layers( + model.layers, opset=18 + ) + + payload = cast(dict[str, object], onnx_model.payload) + graph = cast(dict[str, object], payload["graph"]) + output_infos = cast(list[dict[str, object]], graph["outputs"]) + observed_outputs = _evaluate_fake_tcn_outputs(payload, input_data) + observed_ordered = [ + observed_outputs[cast(str, output_info["name"])] + for output_info in output_infos + ] + self.assertEqual(len(observed_ordered), 2) + + assert_parity_close( + observed_ordered[0], + expected_residual.squeeze(), + family="tcn", + layer_group="default", + context="output=residual", + ) + assert_parity_close( + observed_ordered[1], + expected_skip.squeeze(), + family="tcn", + layer_group="default", + context="output=skip", + ) + + def test_tcn_block_invalid_weights_rank(self): + # Weight rank 3 instead of 4 + conv = ConvolutionalLayer( + weights=np.ones((1, 4, 3)), bias=np.zeros(4), pad="valid" + ) + block = TCNBlock(dilated_conv=conv, dilation_rate=1, activation_fn=relu) + + converter: Any = _get_converter_module() + with self._patched_import(): + with self.assertRaisesRegex(Exception, "weights must be rank-4"): + converter._build_onnx_model_from_primitive_layers([block], opset=18) + + def test_tcn_block_feedforward_skip_and_residual_paths(self): + conv = ConvolutionalLayer( + weights=np.array( + [ + [ + [[0.2, -0.1, 0.05]], + [[-0.05, 0.12, 0.08]], + ], + [ + [[-0.3, 0.25, 0.1]], + [[0.2, -0.15, 0.07]], + ], + ], + dtype=np.float32, + ), + bias=np.array([0.02, -0.03], dtype=np.float32), + pad="valid", + ) + skip_ff = FeedForwardLayer( + weights=np.array([[0.4, -0.2], [0.1, 0.3]], dtype=np.float32), + bias=np.array([0.05, -0.04], dtype=np.float32), + activation_fn=linear, + ) + residual_ff = FeedForwardLayer( + weights=np.array([[0.3, 0.2], [-0.1, 0.5]], dtype=np.float32), + bias=np.array([0.01, -0.02], dtype=np.float32), + activation_fn=linear, + ) + block = TCNBlock( + dilated_conv=conv, + dilation_rate=1, + activation_fn=relu, + skip_conv=skip_ff, + residual_conv=residual_ff, + ) + model = LayerContainer([block]) + input_data = np.array( + [ + [[0.2, -0.1]], + [[0.0, 0.3]], + [[-0.4, 0.5]], + [[0.7, -0.2]], + ], + dtype=np.float32, + ) + + _, skip_source = _evaluate_tcn_block_reference(input_data, block) + expected_skip = np.matmul(skip_source, skip_ff.weights) + skip_ff.bias + expected_residual_path = ( + np.matmul(input_data, residual_ff.weights) + residual_ff.bias + ) + expected_residual = expected_residual_path + expected_skip + + converter: Any = _get_converter_module() + with self._patched_import(): + onnx_model: Any = converter._build_onnx_model_from_primitive_layers( + model.layers, opset=18 + ) + + payload = cast(dict[str, object], onnx_model.payload) + graph = cast(dict[str, object], payload["graph"]) + nodes = cast(list[dict[str, object]], graph["nodes"]) + conv_nodes = [n for n in nodes if n["op_type"] == "Conv"] + matmul_nodes = [n for n in nodes if n["op_type"] == "MatMul"] + self.assertEqual(len(conv_nodes), 1) + self.assertEqual(len(matmul_nodes), 2) + + output_infos = cast(list[dict[str, object]], graph["outputs"]) + observed_outputs = _evaluate_fake_tcn_outputs(payload, input_data) + observed_ordered = [ + observed_outputs[cast(str, output_info["name"])] + for output_info in output_infos + ] + self.assertEqual(len(observed_ordered), 2) + + assert_parity_close( + observed_ordered[0], + expected_residual.squeeze(), + family="tcn", + layer_group="default", + context="output=residual feedforward-branches", + ) + assert_parity_close( + observed_ordered[1], + expected_skip.squeeze(), + family="tcn", + layer_group="default", + context="output=skip feedforward-branches", + ) diff --git a/tools/convert_models_to_onnx.py b/tools/convert_models_to_onnx.py new file mode 100644 index 000000000..6bf0680f3 --- /dev/null +++ b/tools/convert_models_to_onnx.py @@ -0,0 +1,3090 @@ +#!/usr/bin/env python +from __future__ import absolute_import, annotations, division, print_function + +import argparse +import glob +import hashlib +import importlib +import importlib.machinery +import json +import os +import pickle +import re +import sys +import types +from contextlib import contextmanager +from collections.abc import Callable +from collections.abc import Iterator +from collections.abc import Sequence +from typing import Optional, Protocol, cast + +import numpy as np + + +ManifestEntry = dict[str, object] +ManifestData = dict[str, object] + +SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__)) +REPO_ROOT = os.path.dirname(SCRIPT_DIR) +DEFAULT_MODELS_DIR = os.path.join(REPO_ROOT, "madmom", "models") +DEFAULT_SCHEMA_PATH = os.path.join(DEFAULT_MODELS_DIR, "manifest.schema.json") +DEFAULT_OPSET = 18 +MODE_DRY_RUN_INVENTORY = "dry-run-inventory" +MODE_CONVERT = "convert" +MODE_CHECK = "check" + +RESULT_CONVERTED = "converted" +RESULT_SKIPPED_OUT_OF_SCOPE = "skipped_out_of_scope" +RESULT_CORRUPT_SOURCE = "corrupt_source" +RESULT_CONVERTER_UNAVAILABLE = "converter_unavailable" +RESULT_CONVERSION_ERROR = "conversion_error" +RESULT_CHECK_FAILED = "check_failed" + + +class CliArgs(argparse.Namespace): + dry_run_inventory: bool = False + convert: bool = False + check: bool = False + models_dir: str = DEFAULT_MODELS_DIR + schema: str = DEFAULT_SCHEMA_PATH + manifest: str | None = None + opset: int = DEFAULT_OPSET + + +ConverterFn = Callable[[str, ManifestEntry, int], object] + + +class _SerializableModel(Protocol): + def SerializeToString(self, deterministic: bool = False) -> object: ... + + +class _OnnxChecker(Protocol): + def check_model(self, model_path: str) -> None: ... + + +class _OnnxModule(Protocol): + checker: _OnnxChecker + + +class _OnnxTensorProto(Protocol): + FLOAT: int + INT64: int + + +class _OnnxNumpyHelper(Protocol): + def from_array(self, array: np.ndarray, name: str | None = None) -> object: ... + + +class _OnnxGraphHelper(Protocol): + def make_tensor_value_info( + self, name: str, elem_type: int, shape: object + ) -> object: ... + + def make_node( + self, + op_type: str, + inputs: Sequence[str], + outputs: Sequence[str], + name: str | None = None, + **kwargs: object, + ) -> object: ... + + def make_graph( + self, + nodes: Sequence[object], + name: str, + inputs: Sequence[object], + outputs: Sequence[object], + initializer: Sequence[object] | None = None, + ) -> object: ... + + def make_model( + self, + graph: object, + producer_name: str, + opset_imports: Sequence[object], + ) -> object: ... + + +class _OnnxBuilderModule(_OnnxModule, Protocol): + TensorProto: _OnnxTensorProto + helper: _OnnxGraphHelper + numpy_helper: _OnnxNumpyHelper + + +class UnsupportedActivationError(ValueError): + pass + + +class UnsupportedLayoutError(ValueError): + pass + + +class SequentialLayer(object): + def __init__(self, layers: Sequence[object]): + self.layers = list(layers) + + +class ParallelLayer(object): + def __init__(self, layers: Sequence[object]): + self.layers = list(layers) + + +def _infer_layer_input_rank(layer: object) -> int | None: + layer_type = layer.__class__.__name__ + + if layer_type in {"SequentialLayer", "MultiTaskLayer"}: + sub_layers = getattr(layer, "layers", None) + if isinstance(sub_layers, list): + for sub_layer in sub_layers: + inferred = _infer_layer_input_rank(sub_layer) + if inferred is not None: + return inferred + return None + + if layer_type == "ParallelLayer": + sub_layers = getattr(layer, "layers", None) + if isinstance(sub_layers, list): + inferred_values = [ + inferred + for inferred in ( + _infer_layer_input_rank(sub_layer) for sub_layer in sub_layers + ) + if inferred is not None + ] + if inferred_values: + return max(inferred_values) + return None + + if layer_type == "BidirectionalLayer": + fwd_layer = getattr(layer, "fwd_layer", None) + if fwd_layer is not None: + return _infer_layer_input_rank(fwd_layer) + return None + + if layer_type == "TCNLayer": + blocks = getattr(layer, "tcn_blocks", None) + if isinstance(blocks, list) and blocks: + return _infer_layer_input_rank(blocks[0]) + return 2 + + if layer_type == "TCNBlock": + dilated_conv = getattr(layer, "dilated_conv", None) + if isinstance(dilated_conv, list) and dilated_conv: + first_conv = dilated_conv[0] + else: + first_conv = dilated_conv + if first_conv is None: + return 2 + return _infer_layer_input_rank(first_conv) + + if layer_type == "ConvolutionalLayer": + weights = np.asarray(getattr(layer, "weights", []), dtype=np.float32) + if weights.ndim == 4 and int(weights.shape[0]) > 1: + return 3 + return 2 + + if layer_type in { + "FeedForwardLayer", + "RecurrentLayer", + "LSTMLayer", + "GRULayer", + }: + return 2 + + return None + + +def _infer_graph_input_rank(layers: Sequence[object]) -> int: + if not layers: + return 2 + + for layer in layers: + inferred = _infer_layer_input_rank(layer) + if inferred is not None: + return inferred + return 2 + + +def _normalized_sys_path_entry(path: str) -> str: + return os.path.normcase(os.path.abspath(path if path else os.getcwd())) + + +@contextmanager +def _temporary_repo_root_on_syspath() -> Iterator[None]: + repo_root_normalized = _normalized_sys_path_entry(REPO_ROOT) + has_repo_root = any( + _normalized_sys_path_entry(path) == repo_root_normalized for path in sys.path + ) + added_repo_root = False + if not has_repo_root: + sys.path.insert(0, REPO_ROOT) + added_repo_root = True + + try: + yield + finally: + if added_repo_root: + for index, path in enumerate(sys.path): + if _normalized_sys_path_entry(path) == repo_root_normalized: + del sys.path[index] + break + + +@contextmanager +def _temporary_local_madmom_package() -> Iterator[None]: + package_name = "madmom" + if package_name in sys.modules: + yield + return + + package_dir = os.path.join(REPO_ROOT, package_name) + module = types.ModuleType(package_name) + module.__package__ = package_name + module.__path__ = [package_dir] # type: ignore[attr-defined] + module.__file__ = os.path.join(package_dir, "__init__.py") + module.__spec__ = importlib.machinery.ModuleSpec( + name=package_name, + loader=None, + is_package=True, + ) + sys.modules[package_name] = module + + try: + yield + finally: + current = sys.modules.get(package_name) + if current is module: + del sys.modules[package_name] + + +@contextmanager +def _temporary_numpy_shape_base_alias() -> Iterator[None]: + legacy_module_name = "numpy.lib.shape_base" + if legacy_module_name in sys.modules: + yield + return + + try: + replacement_module = importlib.import_module("numpy.lib._shape_base_impl") + except ImportError: + yield + return + + sys.modules[legacy_module_name] = replacement_module + try: + yield + finally: + current = sys.modules.get(legacy_module_name) + if current is replacement_module: + del sys.modules[legacy_module_name] + + +def _as_dict(value: object) -> dict[str, object] | None: + if isinstance(value, dict): + return cast(dict[str, object], value) + return None + + +def _as_list(value: object) -> list[object] | None: + if isinstance(value, list): + return cast(list[object], value) + return None + + +def _relative_posix_path(path: str, root: str) -> str: + return os.path.relpath(path, root).replace(os.sep, "/") + + +def _iter_model_pickles(models_dir: str) -> list[str]: + pattern = os.path.join(models_dir, "**", "*.pkl") + return sorted(glob.glob(pattern, recursive=True)) + + +def _sha256_digest(path: str) -> str: + digest = hashlib.sha256() + with open(path, "rb") as infile: + while True: + data = infile.read(1024 * 1024) + if not data: + break + digest.update(data) + return "sha256:%s" % digest.hexdigest() + + +def _classify_by_path( + source_pkl: str, model_hash: str, target_onnx: str, opset: int +) -> ManifestEntry: + basename = os.path.basename(source_pkl).lower() + + if source_pkl.startswith("patterns/"): + return { + "source_pkl": source_pkl, + "target_onnx": None, + "hash": model_hash, + "opset": None, + "model_type": "gmm_pattern_model", + "status": "out_of_scope", + "rationale": "GMM pattern artifact is outside the NN conversion scope.", + "conversion_mode": "none", + "convertible": False, + } + + if source_pkl.startswith("chords/") and "crf" in basename: + return { + "source_pkl": source_pkl, + "target_onnx": None, + "hash": model_hash, + "opset": None, + "model_type": "conditional_random_field", + "status": "out_of_scope", + "rationale": "CRF artifact is not an nn-convertible model.", + "conversion_mode": "none", + "convertible": False, + } + + if source_pkl.startswith("notes/") and "cnn" in basename: + return { + "source_pkl": source_pkl, + "target_onnx": target_onnx, + "hash": model_hash, + "opset": opset, + "model_type": "wrapped_nn_processor", + "status": "nn_convertible", + "rationale": "Notes CNN artifact is a SequentialProcessor wrapper; convert only the wrapped NN core.", + "conversion_mode": "wrapped_nn_core", + "convertible": True, + } + + nn_family_prefixes = ( + "beats/", + "chroma/", + "chords/", + "downbeats/", + "key/", + "notes/", + "onsets/", + ) + if source_pkl.startswith(nn_family_prefixes): + return { + "source_pkl": source_pkl, + "target_onnx": target_onnx, + "hash": model_hash, + "opset": opset, + "model_type": "neural_network", + "status": "nn_convertible", + "rationale": "Artifact contains NN layers and is in conversion scope.", + "conversion_mode": "direct_nn", + "convertible": True, + } + + return { + "source_pkl": source_pkl, + "target_onnx": None, + "hash": model_hash, + "opset": None, + "model_type": "unsupported_artifact", + "status": "out_of_scope", + "rationale": "Artifact does not expose an NN model supported by this migration.", + "conversion_mode": "none", + "convertible": False, + } + + +def classify_artifact( + source_file: str, models_dir: str, opset: int = DEFAULT_OPSET +) -> ManifestEntry: + source_pkl = _relative_posix_path(source_file, models_dir) + target_onnx = source_pkl[:-4] + ".onnx" + model_hash = _sha256_digest(source_file) + return _classify_by_path(source_pkl, model_hash, target_onnx, opset) + + +def build_supported_artifact_matrix( + models_dir: str = DEFAULT_MODELS_DIR, opset: int = DEFAULT_OPSET +) -> ManifestData: + artifacts: list[ManifestEntry] = [] + for source_file in _iter_model_pickles(models_dir): + artifacts.append(classify_artifact(source_file, models_dir, opset=opset)) + return {"schema_version": 1, "artifacts": artifacts} + + +def _json_type_matches(value: object, json_type: str) -> bool: + if json_type == "object": + return isinstance(value, dict) + if json_type == "array": + return isinstance(value, list) + if json_type == "string": + return isinstance(value, str) + if json_type == "integer": + return isinstance(value, int) and not isinstance(value, bool) + if json_type == "boolean": + return isinstance(value, bool) + if json_type == "null": + return value is None + return True + + +def validate_manifest_shape(manifest: object, schema: object) -> list[str]: + errors: list[str] = [] + manifest_dict = _as_dict(manifest) + schema_dict = _as_dict(schema) + if manifest_dict is None: + return ["manifest must be a JSON object"] + if schema_dict is None: + return ["schema must be a JSON object"] + + required_keys = _as_list(schema_dict.get("required")) + if required_keys is not None: + for required_key in required_keys: + if isinstance(required_key, str) and required_key not in manifest_dict: + errors.append("missing top-level key: %s" % required_key) + + artifacts = _as_list(manifest_dict.get("artifacts")) + if artifacts is None: + errors.append("top-level artifacts must be a list") + return errors + + properties = _as_dict(schema_dict.get("properties")) + artifacts_schema = _as_dict(properties.get("artifacts") if properties else None) + item_schema = _as_dict(artifacts_schema.get("items") if artifacts_schema else None) + + required = _as_list(item_schema.get("required") if item_schema else None) + property_map = _as_dict(item_schema.get("properties") if item_schema else None) + allow_extra_value = item_schema.get("additionalProperties") if item_schema else True + allow_extra = bool(allow_extra_value) + + for index, artifact in enumerate(artifacts): + prefix = "artifacts[%d]" % index + artifact_dict = _as_dict(artifact) + if artifact_dict is None: + errors.append("%s must be an object" % prefix) + continue + + if required is not None: + for required_key in required: + if isinstance(required_key, str) and required_key not in artifact_dict: + errors.append("%s missing key: %s" % (prefix, required_key)) + + if not allow_extra and property_map is not None: + for key in artifact_dict: + if key not in property_map: + errors.append("%s has unexpected key: %s" % (prefix, key)) + + if property_map is None: + continue + + for key, prop_value in property_map.items(): + prop = _as_dict(prop_value) + if prop is None or key not in artifact_dict: + continue + value = artifact_dict[key] + + expected_type = prop.get("type") + if isinstance(expected_type, list): + valid = False + for type_value in cast(list[object], expected_type): + if isinstance(type_value, str) and _json_type_matches( + value, type_value + ): + valid = True + if not valid: + errors.append("%s.%s has invalid type" % (prefix, key)) + elif isinstance(expected_type, str): + if not _json_type_matches(value, expected_type): + errors.append("%s.%s has invalid type" % (prefix, key)) + + enum_values = _as_list(prop.get("enum")) + if enum_values is not None and value not in enum_values: + errors.append("%s.%s has unsupported value" % (prefix, key)) + + pattern_value = prop.get("pattern") + if isinstance(pattern_value, str) and isinstance(value, str): + if re.match(pattern_value, value) is None: + errors.append( + "%s.%s does not match expected pattern" % (prefix, key) + ) + + status = artifact_dict.get("status") + convertible = artifact_dict.get("convertible") + if status == "nn_convertible" and convertible is not True: + errors.append( + "%s.convertible must be true for nn_convertible status" % prefix + ) + if status == "out_of_scope" and convertible is not False: + errors.append( + "%s.convertible must be false for out_of_scope status" % prefix + ) + + return errors + + +def validate_conversion_report_shape(report: object) -> list[str]: + errors: list[str] = [] + report_dict = _as_dict(report) + if report_dict is None: + return ["report must be a JSON object"] + + results = _as_list(report_dict.get("results")) + if results is None: + return ["top-level results must be a list"] + + mode_value = report_dict.get("mode") + if mode_value is not None and mode_value != MODE_CONVERT: + errors.append("mode must be 'convert' when provided") + + for index, result in enumerate(results): + prefix = "results[%d]" % index + result_dict = _as_dict(result) + if result_dict is None: + errors.append("%s must be an object" % prefix) + continue + + required_keys = [ + "source_pkl", + "target_onnx", + "source_hash", + "status", + "output_hash", + "checker_invoked", + "error", + ] + for required_key in required_keys: + if required_key not in result_dict: + errors.append("%s missing key: %s" % (prefix, required_key)) + + source_pkl = result_dict.get("source_pkl") + if not isinstance(source_pkl, str): + errors.append("%s.source_pkl has invalid type" % prefix) + + target_onnx = result_dict.get("target_onnx") + if target_onnx is not None and not isinstance(target_onnx, str): + errors.append("%s.target_onnx has invalid type" % prefix) + + source_hash = result_dict.get("source_hash") + if not isinstance(source_hash, str): + errors.append("%s.source_hash has invalid type" % prefix) + + status = result_dict.get("status") + if not isinstance(status, str): + errors.append("%s.status has invalid type" % prefix) + + output_hash = result_dict.get("output_hash") + if output_hash is not None and not isinstance(output_hash, str): + errors.append("%s.output_hash has invalid type" % prefix) + + checker_invoked = result_dict.get("checker_invoked") + if not isinstance(checker_invoked, bool): + errors.append("%s.checker_invoked has invalid type" % prefix) + + error_value = result_dict.get("error") + if error_value is not None and not isinstance(error_value, str): + errors.append("%s.error has invalid type" % prefix) + + return errors + + +def validate_check_input_shape(manifest: object, schema: object) -> list[str]: + manifest_dict = _as_dict(manifest) + if manifest_dict is None: + return ["manifest must be a JSON object"] + + artifacts = _as_list(manifest_dict.get("artifacts")) + if artifacts is not None: + return validate_manifest_shape(manifest, schema) + + results = _as_list(manifest_dict.get("results")) + if results is not None: + return validate_conversion_report_shape(manifest) + + return ["manifest must contain an artifacts or results list"] + + +def _load_json(path: str) -> object: + with open(path, "r") as infile: + return cast(object, json.load(infile)) + + +def _write_json(data: object, path: str) -> None: + with open(path, "w") as outfile: + json.dump(data, outfile, indent=2, sort_keys=True) + _ = outfile.write("\n") + + +def determine_mode(cli_args: CliArgs) -> str: + selected: list[str] = [] + if cli_args.dry_run_inventory: + selected.append(MODE_DRY_RUN_INVENTORY) + if cli_args.convert: + selected.append(MODE_CONVERT) + if cli_args.check: + selected.append(MODE_CHECK) + + if len(selected) > 1: + raise SystemExit( + "use exactly one of --dry-run-inventory, --convert, or --check" + ) + if not selected: + return MODE_DRY_RUN_INVENTORY + return selected[0] + + +def _serialize_onnx_model_bytes(model: object) -> bytes: + if isinstance(model, bytes): + return model + if isinstance(model, bytearray): + return bytes(model) + + serializable = cast(_SerializableModel, model) + try: + serialize = serializable.SerializeToString + except AttributeError as error: + raise TypeError( + "converter must return ONNX bytes or a serializable model" + ) from error + if not callable(serialize): + raise TypeError("converter must return ONNX bytes or a serializable model") + + serialized: object + try: + serialized = serialize(deterministic=True) + except TypeError: + serialized = serialize() + + if isinstance(serialized, bytes): + return serialized + if isinstance(serialized, bytearray): + return bytes(serialized) + raise TypeError("SerializeToString() must return bytes") + + +def _artifact_source_sort_key(artifact: object) -> str: + artifact_dict = _as_dict(artifact) + if artifact_dict is None: + return "" + source_pkl = artifact_dict.get("source_pkl") + if isinstance(source_pkl, str): + return source_pkl + return "" + + +def write_deterministic_onnx(model: object, output_path: str) -> str: + onnx_bytes = _serialize_onnx_model_bytes(model) + output_dir = os.path.dirname(output_path) + if output_dir: + os.makedirs(output_dir, exist_ok=True) + with open(output_path, "wb") as outfile: + _ = outfile.write(onnx_bytes) + return "sha256:%s" % hashlib.sha256(onnx_bytes).hexdigest() + + +def run_onnx_checker_if_available(model_path: str) -> bool: + try: + imported_module = importlib.import_module("onnx") + onnx_module = cast(_OnnxModule, cast(object, imported_module)) + except ImportError: + return False + + onnx_module.checker.check_model(model_path) + return True + + +def _import_onnx_for_conversion() -> _OnnxBuilderModule: + try: + imported_module = importlib.import_module("onnx") + except ImportError as error: + raise NotImplementedError( + "primitive NN conversion requires optional dependency 'onnx'" + ) from error + return cast(_OnnxBuilderModule, cast(object, imported_module)) + + +class _PrimitiveGraphBuilder: + def __init__(self, onnx_module: _OnnxBuilderModule, opset: int): + self.onnx_module = onnx_module + self.opset = opset + self.nodes: list[object] = [] + self.initializers: list[object] = [] + self.current_tensor = "input" + self.current_rank: int | None = 3 + self._counter = 0 + + def _name(self, stem: str) -> str: + name = "%s_%d" % (stem, self._counter) + self._counter += 1 + return name + + def add_initializer(self, stem: str, array: np.ndarray) -> str: + name = self._name(stem) + initializer = self.onnx_module.numpy_helper.from_array(array, name=name) + self.initializers.append(initializer) + return name + + def add_node( + self, + op_type: str, + inputs: Sequence[str], + stem: str, + **attrs: object, + ) -> str: + output = self._name(stem) + node = self.onnx_module.helper.make_node( + op_type, + list(inputs), + [output], + name=output, + **attrs, + ) + self.nodes.append(node) + self.current_tensor = output + return output + + +def _layer_label(index: int, layer: object) -> str: + return "%s[%d]" % (layer.__class__.__name__, index) + + +def _require_attr(layer: object, name: str, label: str) -> object: + if not hasattr(layer, name): + raise UnsupportedLayoutError( + "%s is missing required attribute '%s'" % (label, name) + ) + return getattr(layer, name) + + +def _as_float32_array(value: object, what: str) -> np.ndarray: + try: + array = np.asarray(value, dtype=np.float32) + except (TypeError, ValueError) as error: + raise UnsupportedLayoutError("%s must be numeric" % what) from error + return array + + +def _as_int(value: object, what: str) -> int: + if isinstance(value, bool) or not isinstance(value, int): + raise UnsupportedLayoutError("%s must be an int" % what) + return value + + +def _as_positive_int(value: object, what: str) -> int: + int_value = _as_int(value, what) + if int_value <= 0: + raise UnsupportedLayoutError("%s must be > 0" % what) + return int_value + + +def _as_int_pair(value: object, what: str) -> tuple[int, int]: + if isinstance(value, bool): + raise UnsupportedLayoutError("%s must be int or length-2 tuple/list" % what) + if isinstance(value, int): + pair = (value, value) + elif isinstance(value, (list, tuple)): + if len(value) != 2: + raise UnsupportedLayoutError("%s must contain exactly 2 ints" % what) + first = _as_positive_int(value[0], "%s[0]" % what) + second = _as_positive_int(value[1], "%s[1]" % what) + pair = (first, second) + else: + raise UnsupportedLayoutError("%s must be int or length-2 tuple/list" % what) + + if pair[0] <= 0 or pair[1] <= 0: + raise UnsupportedLayoutError("%s values must be > 0" % what) + return pair + + +def _add_axes_initializer( + builder: _PrimitiveGraphBuilder, stem: str, axes: Sequence[int] +) -> str: + return builder.add_initializer(stem, np.asarray(list(axes), dtype=np.int64)) + + +def _append_unsqueeze( + builder: _PrimitiveGraphBuilder, + axes: Sequence[int], + stem: str, +) -> None: + axes_name = _add_axes_initializer(builder, "%s_axes" % stem, axes) + _ = builder.add_node("Unsqueeze", [builder.current_tensor, axes_name], stem) + if builder.current_rank is not None: + builder.current_rank += len(axes) + + +def _append_squeeze( + builder: _PrimitiveGraphBuilder, + axes: Sequence[int], + stem: str, +) -> None: + axes_name = _add_axes_initializer(builder, "%s_axes" % stem, axes) + _ = builder.add_node("Squeeze", [builder.current_tensor, axes_name], stem) + if builder.current_rank is not None: + builder.current_rank -= len(axes) + + +def _prepare_spatial_input_for_nchw( + builder: _PrimitiveGraphBuilder, + label: str, + what: str, +) -> bool: + if builder.current_rank == 2: + _append_unsqueeze(builder, [2], "%s_expand_channel" % what) + if builder.current_rank != 3: + raise UnsupportedLayoutError( + "%s internal error while preparing channels" % label + ) + channel_added = True + elif builder.current_rank == 3: + channel_added = False + else: + raise UnsupportedLayoutError( + "%s requires rank-2 or rank-3 input in [time, freq, channel] layout" % label + ) + + _append_unsqueeze(builder, [0], "%s_add_batch" % what) + _ = builder.add_node( + "Transpose", + [builder.current_tensor], + "%s_to_nchw" % what, + perm=[0, 3, 1, 2], + ) + builder.current_rank = 4 + return channel_added + + +def _restore_spatial_output_from_nchw( + builder: _PrimitiveGraphBuilder, + *, + channel_added: bool, + what: str, + remove_added_channel: bool = True, +) -> None: + _ = builder.add_node( + "Transpose", + [builder.current_tensor], + "%s_to_nhwc" % what, + perm=[0, 2, 3, 1], + ) + builder.current_rank = 4 + _append_squeeze(builder, [0], "%s_remove_batch" % what) + builder.current_rank = 3 + if channel_added and remove_added_channel: + _append_squeeze(builder, [2], "%s_remove_channel" % what) + builder.current_rank = 2 + + +def _normalize_axes(axis_values: Sequence[int], rank: int, label: str) -> list[int]: + normalized: list[int] = [] + for axis in axis_values: + normalized_axis = axis + if axis < 0: + normalized_axis = rank + axis + if normalized_axis < 0 or normalized_axis >= rank: + raise UnsupportedLayoutError( + "%s axis %d is out of bounds for rank %d" % (label, axis, rank) + ) + normalized.append(normalized_axis) + if len(set(normalized)) != len(normalized): + raise UnsupportedLayoutError("%s axes must be unique" % label) + return normalized + + +def _activation_name(activation_fn: object, label: str) -> str | None: + if activation_fn is None: + return None + activation_name = getattr(activation_fn, "__name__", None) + if not isinstance(activation_name, str): + raise UnsupportedActivationError( + "%s activation must expose a __name__ attribute" % label + ) + if activation_name in {"linear", "tanh", "sigmoid", "relu", "elu", "softmax"}: + return activation_name + raise UnsupportedActivationError( + "%s activation '%s' is unsupported; supported activations: " + "linear, tanh, sigmoid, relu, elu, softmax" % (label, activation_name) + ) + + +def _append_activation( + builder: _PrimitiveGraphBuilder, + layer: object, + label: str, +) -> None: + activation_name = _activation_name( + _require_attr(layer, "activation_fn", label), label + ) + if activation_name is None or activation_name == "linear": + return + + if activation_name == "softmax": + if builder.current_rank is not None and builder.current_rank < 2: + raise UnsupportedLayoutError("%s softmax requires rank >= 2" % label) + _ = builder.add_node("Softmax", [builder.current_tensor], "softmax", axis=-1) + return + + op_map = { + "tanh": "Tanh", + "sigmoid": "Sigmoid", + "relu": "Relu", + "elu": "Elu", + } + op_type = op_map[activation_name] + if op_type == "Elu": + _ = builder.add_node(op_type, [builder.current_tensor], "elu", alpha=1.0) + else: + _ = builder.add_node(op_type, [builder.current_tensor], activation_name) + + +def _as_1d_float32_array(value: object, what: str) -> np.ndarray: + array = _as_float32_array(value, what) + if array.ndim != 1: + raise UnsupportedLayoutError( + "%s must be rank-1; got shape %s" % (what, tuple(array.shape)) + ) + return array + + +def _onnx_recurrent_activation_name( + activation_fn: object, + label: str, + allowed: set[str], +) -> str: + activation_name = _activation_name(activation_fn, label) + if activation_name is None: + raise UnsupportedActivationError( + "%s recurrent activation must not be None" % label + ) + + if activation_name not in allowed: + supported = ", ".join(sorted(allowed)) + raise UnsupportedActivationError( + "%s recurrent activation '%s' is unsupported; supported activations: %s" + % (label, activation_name, supported) + ) + + activation_map = { + "tanh": "Tanh", + "sigmoid": "Sigmoid", + "relu": "Relu", + } + return activation_map[activation_name] + + +def _temporarily_convert_tensor( + builder: _PrimitiveGraphBuilder, + tensor_name: str, + tensor_rank: int | None, + axes: Sequence[int], + stem: str, + op_type: str, +) -> str: + previous_tensor = builder.current_tensor + previous_rank = builder.current_rank + builder.current_tensor = tensor_name + builder.current_rank = tensor_rank + if op_type == "Unsqueeze": + _append_unsqueeze(builder, axes, stem) + else: + _append_squeeze(builder, axes, stem) + output = builder.current_tensor + builder.current_tensor = previous_tensor + builder.current_rank = previous_rank + return output + + +def _unsqueeze_tensor( + builder: _PrimitiveGraphBuilder, + tensor_name: str, + tensor_rank: int | None, + axes: Sequence[int], + stem: str, +) -> str: + return _temporarily_convert_tensor( + builder=builder, + tensor_name=tensor_name, + tensor_rank=tensor_rank, + axes=axes, + stem=stem, + op_type="Unsqueeze", + ) + + +def _squeeze_tensor( + builder: _PrimitiveGraphBuilder, + tensor_name: str, + tensor_rank: int | None, + axes: Sequence[int], + stem: str, +) -> str: + return _temporarily_convert_tensor( + builder=builder, + tensor_name=tensor_name, + tensor_rank=tensor_rank, + axes=axes, + stem=stem, + op_type="Squeeze", + ) + + +def _append_multi_output_node( + builder: _PrimitiveGraphBuilder, + op_type: str, + inputs: Sequence[str], + outputs: Sequence[str], + stem: str, + **attrs: object, +) -> None: + node_name = builder._name(stem) + node = builder.onnx_module.helper.make_node( + op_type, + list(inputs), + list(outputs), + name=node_name, + **attrs, + ) + builder.nodes.append(node) + + +TensorReference = tuple[str, Optional[int]] +TensorReferences = list[TensorReference] + + +def _validate_state_size( + state: np.ndarray, + expected_size: int, + what: str, +) -> None: + if state.size != expected_size: + raise UnsupportedLayoutError( + "%s size mismatch: expected %d, got %d" % (what, expected_size, state.size) + ) + + +def _validate_gate_shapes( + gate: object, + *, + gate_name: str, + label: str, + input_size: int, + hidden_size: int, + allow_peephole: bool, +) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray | None, object]: + gate_label = "%s.%s" % (label, gate_name) + weights = _as_float32_array( + _require_attr(gate, "weights", gate_label), "%s.weights" % gate_label + ) + recurrent_weights = _as_float32_array( + _require_attr(gate, "recurrent_weights", gate_label), + "%s.recurrent_weights" % gate_label, + ) + bias = _as_1d_float32_array( + _require_attr(gate, "bias", gate_label), "%s.bias" % gate_label + ) + + if weights.shape != (input_size, hidden_size): + raise UnsupportedLayoutError( + "%s.weights must have shape (%d, %d); got %s" + % (gate_label, input_size, hidden_size, tuple(weights.shape)) + ) + if recurrent_weights.shape != (hidden_size, hidden_size): + raise UnsupportedLayoutError( + "%s.recurrent_weights must have shape (%d, %d); got %s" + % (gate_label, hidden_size, hidden_size, tuple(recurrent_weights.shape)) + ) + if bias.size != hidden_size: + raise UnsupportedLayoutError( + "%s.bias size mismatch: expected %d, got %d" + % (gate_label, hidden_size, bias.size) + ) + + peephole_value = getattr(gate, "peephole_weights", None) + peephole_weights: np.ndarray | None + if peephole_value is None: + peephole_weights = None + else: + if not allow_peephole: + raise UnsupportedLayoutError( + "%s peephole weights are unsupported in this context" % gate_label + ) + peephole_weights = _as_1d_float32_array( + peephole_value, "%s.peephole_weights" % gate_label + ) + if peephole_weights.size != hidden_size: + raise UnsupportedLayoutError( + "%s.peephole_weights size mismatch: expected %d, got %d" + % (gate_label, hidden_size, peephole_weights.size) + ) + + activation_fn = _require_attr(gate, "activation_fn", gate_label) + return weights, recurrent_weights, bias, peephole_weights, activation_fn + + +def _convert_feed_forward_layer( + builder: _PrimitiveGraphBuilder, + layer: object, + label: str, +) -> None: + weights = _as_float32_array( + _require_attr(layer, "weights", label), "%s.weights" % label + ) + if weights.ndim not in (1, 2): + raise UnsupportedLayoutError( + "%s weights must be rank-1 or rank-2; got rank %d" % (label, weights.ndim) + ) + + if builder.current_rank is not None and builder.current_rank < 1: + raise UnsupportedLayoutError("%s input rank must be >= 1 for MatMul" % label) + + bias = _as_float32_array(_require_attr(layer, "bias", label), "%s.bias" % label) + bias = bias.reshape(-1) + + weights_name = builder.add_initializer("weights", weights) + bias_name = builder.add_initializer("bias", bias) + matmul_output = builder.add_node( + "MatMul", + [builder.current_tensor, weights_name], + "matmul", + ) + _ = builder.add_node("Add", [matmul_output, bias_name], "bias_add") + + if builder.current_rank is not None: + if weights.ndim == 1: + builder.current_rank = builder.current_rank - 1 + else: + builder.current_rank = builder.current_rank + + _append_activation(builder, layer, label) + + +def _convert_batch_norm_layer( + builder: _PrimitiveGraphBuilder, + layer: object, + label: str, +) -> None: + beta = _as_float32_array(_require_attr(layer, "beta", label), "%s.beta" % label) + gamma = _as_float32_array(_require_attr(layer, "gamma", label), "%s.gamma" % label) + mean = _as_float32_array(_require_attr(layer, "mean", label), "%s.mean" % label) + inv_std = _as_float32_array( + _require_attr(layer, "inv_std", label), "%s.inv_std" % label + ) + + scale = gamma * inv_std + offset = beta - (mean * scale) + + scale_name = builder.add_initializer("batchnorm_scale", scale) + offset_name = builder.add_initializer("batchnorm_offset", offset) + scaled = builder.add_node( + "Mul", + [builder.current_tensor, scale_name], + "batchnorm_mul", + ) + _ = builder.add_node("Add", [scaled, offset_name], "batchnorm_add") + _append_activation(builder, layer, label) + + +def _convert_reshape_layer( + builder: _PrimitiveGraphBuilder, + layer: object, + label: str, +) -> None: + order = _require_attr(layer, "order", label) + if order != "C": + raise UnsupportedLayoutError( + "%s supports only order='C'; got order=%r" % (label, order) + ) + + newshape_value = _require_attr(layer, "newshape", label) + if isinstance(newshape_value, int): + newshape = [newshape_value] + elif isinstance(newshape_value, (list, tuple)): + newshape = list(newshape_value) + else: + raise UnsupportedLayoutError("%s newshape must be int, list, or tuple" % label) + + normalized_shape: list[int] = [] + for dim in newshape: + if isinstance(dim, np.integer): + normalized_dim: object = int(dim) + else: + normalized_dim = dim + + if isinstance(normalized_dim, bool) or not isinstance(normalized_dim, int): + raise UnsupportedLayoutError("%s newshape values must be ints" % label) + if normalized_dim == 0: + raise UnsupportedLayoutError( + "%s newshape with 0 is unsupported for ONNX parity" % label + ) + normalized_shape.append(normalized_dim) + + shape_name = builder.add_initializer( + "reshape_shape", np.asarray(normalized_shape, dtype=np.int64) + ) + _ = builder.add_node("Reshape", [builder.current_tensor, shape_name], "reshape") + builder.current_rank = len(normalized_shape) + + +def _convert_transpose_layer( + builder: _PrimitiveGraphBuilder, + layer: object, + label: str, +) -> None: + axes = _require_attr(layer, "axes", label) + if axes is None: + if builder.current_rank is None: + raise UnsupportedLayoutError( + "%s axes=None requires known rank for reverse permutation" % label + ) + _ = builder.add_node("Transpose", [builder.current_tensor], "transpose") + return + + if not isinstance(axes, (list, tuple)): + raise UnsupportedLayoutError("%s axes must be list or tuple" % label) + + permutation: list[int] = [] + for axis in axes: + if not isinstance(axis, int): + raise UnsupportedLayoutError("%s axes must contain ints" % label) + permutation.append(axis) + + if builder.current_rank is not None: + permutation = _normalize_axes(permutation, builder.current_rank, label) + if len(set(permutation)) != len(permutation): + raise UnsupportedLayoutError("%s axes must be unique" % label) + if sorted(permutation) != list(range(len(permutation))): + raise UnsupportedLayoutError( + "%s axes must be a full permutation of [0, ..., rank-1]" % label + ) + + _ = builder.add_node( + "Transpose", + [builder.current_tensor], + "transpose", + perm=permutation, + ) + builder.current_rank = len(permutation) + + +def _convert_pad_layer( + builder: _PrimitiveGraphBuilder, + layer: object, + label: str, +) -> None: + if builder.current_rank is None: + raise UnsupportedLayoutError( + "%s requires known input rank for static pads" % label + ) + + width_value = _require_attr(layer, "width", label) + if not isinstance(width_value, int): + raise UnsupportedLayoutError("%s width must be an int" % label) + if width_value < 0: + raise UnsupportedLayoutError("%s width must be non-negative" % label) + + axes_value = _require_attr(layer, "axes", label) + if not isinstance(axes_value, (list, tuple)): + raise UnsupportedLayoutError("%s axes must be list or tuple" % label) + + axes: list[int] = [] + for axis in axes_value: + if not isinstance(axis, int): + raise UnsupportedLayoutError("%s axes must contain ints" % label) + axes.append(axis) + normalized_axes = _normalize_axes(axes, builder.current_rank, label) + + pads = np.zeros(builder.current_rank * 2, dtype=np.int64) + for axis in normalized_axes: + pads[axis] = width_value + pads[builder.current_rank + axis] = width_value + + value_array = _as_float32_array( + _require_attr(layer, "value", label), "%s.value" % label + ) + if value_array.size != 1: + raise UnsupportedLayoutError("%s value must be scalar" % label) + + pads_name = builder.add_initializer("pad_width", pads) + value_name = builder.add_initializer( + "pad_value", np.asarray(value_array.item(), dtype=np.float32) + ) + _ = builder.add_node( + "Pad", + [builder.current_tensor, pads_name, value_name], + "pad", + mode="constant", + ) + + +def _convert_average_layer( + builder: _PrimitiveGraphBuilder, + layer: object, + label: str, +) -> None: + dtype_value = _require_attr(layer, "dtype", label) + if dtype_value is not None: + dtype_name: str | None = None + if dtype_value is np.float32: + dtype_name = "float32" + elif isinstance(dtype_value, np.dtype): + dtype_name = dtype_value.name + elif isinstance(dtype_value, str): + dtype_name = dtype_value + + if dtype_name not in {"float32", "single"}: + raise UnsupportedLayoutError( + "%s supports only dtype=None or dtype=np.float32" % label + ) + _ = builder.add_node( + "Cast", + [builder.current_tensor], + "average_cast", + to=builder.onnx_module.TensorProto.FLOAT, + ) + + axis_value = _require_attr(layer, "axis", label) + keepdims_value = _require_attr(layer, "keepdims", label) + keepdims = bool(keepdims_value) + + axes: list[int] | None + if axis_value is None: + axes = None + elif isinstance(axis_value, int): + axes = [axis_value] + elif isinstance(axis_value, (list, tuple)): + axes = [] + for axis in axis_value: + if not isinstance(axis, int): + raise UnsupportedLayoutError("%s axis values must be ints" % label) + axes.append(axis) + else: + raise UnsupportedLayoutError( + "%s axis must be None, int, list, or tuple" % label + ) + + normalized_axes: list[int] | None = axes + if axes is not None and builder.current_rank is not None: + normalized_axes = _normalize_axes(axes, builder.current_rank, label) + + reduce_inputs = [builder.current_tensor] + if normalized_axes is not None: + axes_name = builder.add_initializer( + "reduce_axes", np.asarray(normalized_axes, dtype=np.int64) + ) + reduce_inputs.append(axes_name) + + _ = builder.add_node( + "ReduceMean", + reduce_inputs, + "reduce_mean", + keepdims=1 if keepdims else 0, + ) + + if builder.current_rank is not None: + if normalized_axes is None: + builder.current_rank = builder.current_rank if keepdims else 0 + else: + reduced_rank = builder.current_rank - len(normalized_axes) + builder.current_rank = builder.current_rank if keepdims else reduced_rank + + +def _convert_convolutional_layer( + builder: _PrimitiveGraphBuilder, + layer: object, + label: str, +) -> None: + weights = _as_float32_array( + _require_attr(layer, "weights", label), "%s.weights" % label + ) + if weights.ndim != 4: + raise UnsupportedLayoutError( + "%s weights must have rank 4 [channel, feature, time, freq]" % label + ) + + num_channels = weights.shape[0] + bias = _as_float32_array(_require_attr(layer, "bias", label), "%s.bias" % label) + bias = bias.reshape(-1) + output_features = weights.shape[1] + if bias.size == 1 and output_features > 1: + bias = np.repeat(bias, output_features) + elif bias.size != output_features: + raise UnsupportedLayoutError( + "%s bias size (%d) must be 1 or match number of feature maps (%d)" + % (label, bias.size, output_features) + ) + + stride_value = _require_attr(layer, "stride", label) + if stride_value in (None, 1, (1, 1)): + strides = (1, 1) + else: + strides = _as_int_pair(stride_value, "%s.stride" % label) + + pad_value = _require_attr(layer, "pad", label) + if not isinstance(pad_value, str): + raise UnsupportedLayoutError("%s pad must be a string" % label) + if pad_value not in {"valid", "same"}: + raise UnsupportedLayoutError( + "%s supports only pad='valid' or pad='same'; got %r" % (label, pad_value) + ) + + if builder.current_rank == 2 and num_channels != 1: + raise UnsupportedLayoutError( + "%s rank-2 input implies one channel, but weights expect %d channels" + % (label, num_channels) + ) + if builder.current_rank not in (2, 3): + raise UnsupportedLayoutError( + "%s requires rank-2 or rank-3 input in [time, freq, channel] layout" % label + ) + + channel_added = _prepare_spatial_input_for_nchw(builder, label, "conv") + if channel_added and num_channels != 1: + raise UnsupportedLayoutError( + "%s rank-2 input can only be used with single-channel weights" % label + ) + + transformed_weights = np.flip(weights, axis=(2, 3)).transpose(1, 0, 2, 3) + weights_name = builder.add_initializer("conv_weights", transformed_weights) + bias_name = builder.add_initializer("conv_bias", bias) + + conv_attrs: dict[str, object] = { + "strides": [strides[0], strides[1]], + } + if pad_value == "same": + conv_attrs["auto_pad"] = "SAME_UPPER" + + _ = builder.add_node( + "Conv", + [builder.current_tensor, weights_name, bias_name], + "conv", + **conv_attrs, + ) + _restore_spatial_output_from_nchw( + builder, + channel_added=channel_added, + what="conv", + remove_added_channel=False, + ) + builder.current_rank = 3 + _append_activation(builder, layer, label) + + +def _convert_max_pool_layer( + builder: _PrimitiveGraphBuilder, + layer: object, + label: str, +) -> None: + axis_value = _require_attr(layer, "axis", label) + stride_value = _require_attr(layer, "stride", label) + + def normalize_pool_pair_value(value: object) -> object: + if isinstance(value, (bool, int, list, tuple)): + return value + + value_array = np.asarray(value) + if value_array.ndim == 0: + return value_array.item() + if value_array.ndim == 1: + return value_array.tolist() + return value + + if axis_value is not None: + if stride_value is not None: + raise UnsupportedLayoutError( + "%s does not support axis pooling when stride is set" % label + ) + + if isinstance(axis_value, int): + axes = [axis_value] + elif isinstance(axis_value, (list, tuple)): + axes = [_as_int(axis, "%s.axis" % label) for axis in axis_value] + else: + raise UnsupportedLayoutError( + "%s axis must be an int, tuple, list, or None" % label + ) + + if builder.current_rank is not None: + axes = _normalize_axes(axes, builder.current_rank, label) + axes_name = _add_axes_initializer(builder, "pool_axis", axes) + _ = builder.add_node( + "ReduceMax", + [builder.current_tensor, axes_name], + "pool_axis_reduce", + keepdims=0, + ) + if builder.current_rank is not None: + builder.current_rank -= len(axes) + return + + size_value = _require_attr(layer, "size", label) + size_value = normalize_pool_pair_value(size_value) + size = _as_int_pair(size_value, "%s.size" % label) + + if stride_value is None: + stride = size + else: + stride_value = normalize_pool_pair_value(stride_value) + stride = _as_int_pair(stride_value, "%s.stride" % label) + + channel_added = _prepare_spatial_input_for_nchw(builder, label, "pool") + _ = builder.add_node( + "MaxPool", + [builder.current_tensor], + "max_pool", + kernel_shape=[size[0], size[1]], + strides=[stride[0], stride[1]], + ) + _restore_spatial_output_from_nchw( + builder, + channel_added=channel_added, + what="pool", + remove_added_channel=True, + ) + + +def _convert_tcn_block( + builder: _PrimitiveGraphBuilder, + layer: object, + label: str, + layer_index: int, + squeeze_outputs: bool = True, +) -> TensorReferences: + dilated_conv = _require_attr(layer, "dilated_conv", label) + dilation_rate = _require_attr(layer, "dilation_rate", label) + + if isinstance(dilated_conv, list): + convs = dilated_conv + if isinstance(dilation_rate, (list, tuple)): + rates = list(dilation_rate) + else: + rates = [dilation_rate] * len(convs) + else: + convs = [dilated_conv] + rates = [dilation_rate] + + if len(rates) != len(convs): + raise UnsupportedLayoutError( + "%s dilation_rate length (%d) must match dilated_conv length (%d)" + % (label, len(rates), len(convs)) + ) + + original_input_tensor = builder.current_tensor + original_input_rank = builder.current_rank + + if builder.current_rank == 2: + _append_unsqueeze(builder, [1], "tcn_block_input_expand") + builder.current_rank = 3 + channel_added = True + elif builder.current_rank == 3: + channel_added = False + else: + raise UnsupportedLayoutError( + "%s requires rank-2 or rank-3 input in [time, freq, channel] layout" % label + ) + + res_path_base = builder.current_tensor + res_rank_base = builder.current_rank + + _append_unsqueeze(builder, [0], "tcn_block_add_batch") + _ = builder.add_node( + "Transpose", + [builder.current_tensor], + "tcn_block_to_nchw", + perm=[0, 3, 1, 2], + ) + builder.current_rank = 4 + nchw_input_tensor = builder.current_tensor + + conv_outputs = [] + for i, (conv, rate) in enumerate(zip(convs, rates)): + rate_int = _as_positive_int(rate, "%s.dilation_rate[%d]" % (label, i)) + weights = _as_float32_array( + _require_attr(conv, "weights", "%s.conv[%d]" % (label, i)), + "%s.weights" % label, + ) + if weights.ndim != 4: + raise UnsupportedLayoutError( + "%s.conv[%d] weights must be rank-4" % (label, i) + ) + + # Weights in madmom TCN are (C_in, C_out, 1, K) + # Transpose to (C_out, C_in, K, 1) for Conv along H (time) in NCHW + transformed_weights = weights.transpose(1, 0, 3, 2) + # flip along K axis (axis 2) to match scipy.ndimage.convolve + transformed_weights = np.flip(transformed_weights, axis=2) + + bias = _as_float32_array( + _require_attr(conv, "bias", "%s.conv[%d]" % (label, i)), "%s.bias" % label + ) + bias = bias.reshape(-1) + + weights_name = builder.add_initializer( + "tcn_conv_weights_%d_%d" % (layer_index, i), transformed_weights + ) + bias_name = builder.add_initializer( + "tcn_conv_bias_%d_%d" % (layer_index, i), bias + ) + + kernel_size = transformed_weights.shape[2] + left_pad = 2 * rate_int + right_pad = max((kernel_size - 3) * rate_int, 0) + + conv_output = builder.add_node( + "Conv", + [nchw_input_tensor, weights_name, bias_name], + "tcn_conv_%d_%d" % (layer_index, i), + dilations=[rate_int, 1], + pads=[left_pad, 0, right_pad, 0], + ) + _append_activation(builder, conv, label + ".conv[%d]" % i) + conv_outputs.append(builder.current_tensor) + + if len(conv_outputs) > 1: + out = builder.add_node( + "Concat", conv_outputs, "tcn_conv_concat_%d" % layer_index, axis=1 + ) + else: + out = conv_outputs[0] + + builder.current_tensor = out + _restore_spatial_output_from_nchw( + builder, + channel_added=channel_added, + what="tcn_block_out", + remove_added_channel=False, + ) + builder.current_rank = 3 + + _append_activation(builder, layer, label) + activated_out = builder.current_tensor + + def convert_tcn_branch_projection(branch_layer: object, branch_label: str) -> None: + branch_weights = _as_float32_array( + _require_attr(branch_layer, "weights", branch_label), + "%s.weights" % branch_label, + ) + if branch_weights.ndim == 4: + _convert_convolutional_layer(builder, branch_layer, branch_label) + return + if branch_weights.ndim in (1, 2): + _convert_feed_forward_layer(builder, branch_layer, branch_label) + return + raise UnsupportedLayoutError( + "%s weights must have rank 4 [channel, feature, time, freq] or rank-1/rank-2 feed-forward projection" + % branch_label + ) + + skip_conv = getattr(layer, "skip_conv", None) + if skip_conv is not None: + convert_tcn_branch_projection(skip_conv, label + ".skip_conv") + + final_skip_out = builder.current_tensor + final_skip_rank = builder.current_rank + + residual_conv = getattr(layer, "residual_conv", None) + if residual_conv is not None: + builder.current_tensor = res_path_base + builder.current_rank = res_rank_base + convert_tcn_branch_projection(residual_conv, label + ".residual_conv") + res_path = builder.current_tensor + else: + res_path = res_path_base + + final_out_res = builder.add_node( + "Add", [res_path, final_skip_out], "tcn_residual_add_%d" % layer_index + ) + + def _squeeze_to_rank2( + tensor_name: str, tensor_rank: int | None, stem: str + ) -> TensorReference: + if tensor_rank == 3: + return _squeeze_tensor(builder, tensor_name, 3, [1], stem), 2 + return tensor_name, tensor_rank + + if squeeze_outputs: + res_ref = _squeeze_to_rank2(final_out_res, 3, "tcn_block_res_squeeze") + skip_ref = _squeeze_to_rank2( + final_skip_out, final_skip_rank, "tcn_block_skip_squeeze" + ) + else: + res_ref = (final_out_res, 3) + skip_ref = (final_skip_out, final_skip_rank) + + return [res_ref, skip_ref] + + +def _convert_tcn_layer( + builder: _PrimitiveGraphBuilder, + layer: object, + label: str, + input_values: TensorReferences, + recurrent_state_interfaces: list[tuple[str, str, np.ndarray]], + layer_counter: list[int], +) -> TensorReferences: + tcn_blocks = _require_layer_list(layer, label, "tcn_blocks") + skip_connections = bool(getattr(layer, "skip_connections", False)) + + current_data = input_values[0] + all_skips = [] + + for i, block in enumerate(tcn_blocks): + _set_builder_input(builder, current_data) + block_outputs = _convert_tcn_block( + builder, + block, + label + ".block[%d]" % i, + layer_counter[0], + squeeze_outputs=False, + ) + layer_counter[0] += 1 + current_data = block_outputs[0] + all_skips.append(block_outputs[1]) + + if len(all_skips) > 1: + current_skip = all_skips[0][0] + for i in range(1, len(all_skips)): + current_skip = builder.add_node( + "Add", + [current_skip, all_skips[i][0]], + "tcn_skip_sum_%d_%d" % (layer_counter[0], i), + ) + skip_total = (current_skip, all_skips[0][1]) + else: + skip_total = all_skips[0] + + _set_builder_input(builder, current_data) + _append_activation(builder, layer, label) + current_data = (builder.current_tensor, builder.current_rank) + + if skip_connections: + return [current_data, skip_total] + else: + return [current_data] + + +def _convert_stride_layer( + builder: _PrimitiveGraphBuilder, + layer: object, + label: str, +) -> None: + block_size_value = _require_attr(layer, "block_size", label) + if not isinstance(block_size_value, (bool, int)): + block_size_array = np.asarray(block_size_value) + if block_size_array.ndim == 0: + block_size_value = block_size_array.item() + block_size = _as_positive_int(block_size_value, "%s.block_size" % label) + if builder.current_rank is not None and builder.current_rank < 1: + raise UnsupportedLayoutError("%s requires rank >= 1 input" % label) + + sliced_segments: list[str] = [] + source_tensor = builder.current_tensor + max_end = np.asarray(np.iinfo(np.int64).max, dtype=np.int64) + axis_zero = builder.add_initializer("stride_axis", np.asarray([0], dtype=np.int64)) + unit_step = builder.add_initializer("stride_step", np.asarray([1], dtype=np.int64)) + unsqueeze_axis = _add_axes_initializer(builder, "stride_unsqueeze_axis", [1]) + + for offset in range(block_size): + starts_name = builder.add_initializer( + "stride_starts", np.asarray([offset], dtype=np.int64) + ) + if offset == block_size - 1: + ends_array = np.asarray([max_end.item()], dtype=np.int64) + else: + ends_array = np.asarray([-(block_size - offset - 1)], dtype=np.int64) + ends_name = builder.add_initializer("stride_ends", ends_array) + sliced = builder.add_node( + "Slice", + [source_tensor, starts_name, ends_name, axis_zero, unit_step], + "stride_slice", + ) + segment = builder.add_node( + "Unsqueeze", + [sliced, unsqueeze_axis], + "stride_segment", + ) + sliced_segments.append(segment) + + stacked = builder.add_node("Concat", sliced_segments, "stride_concat", axis=1) + shape = builder.add_node("Shape", [stacked], "stride_shape") + first_dim_index = builder.add_initializer( + "stride_first_dim", np.asarray([0], dtype=np.int64) + ) + first_dim = builder.add_node( + "Gather", + [shape, first_dim_index], + "stride_first_dim_value", + axis=0, + ) + minus_one = builder.add_initializer( + "stride_minus_one", np.asarray([-1], dtype=np.int64) + ) + target_shape = builder.add_node( + "Concat", + [first_dim, minus_one], + "stride_target_shape", + axis=0, + ) + _ = builder.add_node("Reshape", [stacked, target_shape], "stride_flatten") + builder.current_rank = 2 + + +def _convert_recurrent_layer( + builder: _PrimitiveGraphBuilder, + layer: object, + label: str, + layer_index: int, + recurrent_state_interfaces: list[tuple[str, str, np.ndarray]], +) -> None: + weights = _as_float32_array( + _require_attr(layer, "weights", label), "%s.weights" % label + ) + recurrent_weights = _as_float32_array( + _require_attr(layer, "recurrent_weights", label), "%s.recurrent_weights" % label + ) + bias = _as_1d_float32_array(_require_attr(layer, "bias", label), "%s.bias" % label) + init = _as_1d_float32_array(_require_attr(layer, "init", label), "%s.init" % label) + + if weights.ndim != 2: + raise UnsupportedLayoutError( + "%s.weights must be rank-2; got rank %d" % (label, weights.ndim) + ) + input_size, hidden_size = weights.shape + if recurrent_weights.shape != (hidden_size, hidden_size): + raise UnsupportedLayoutError( + "%s.recurrent_weights must have shape (%d, %d); got %s" + % (label, hidden_size, hidden_size, tuple(recurrent_weights.shape)) + ) + if bias.size != hidden_size: + raise UnsupportedLayoutError( + "%s.bias size mismatch: expected %d, got %d" + % (label, hidden_size, bias.size) + ) + _validate_state_size(init, hidden_size, "%s.init" % label) + + activation_name = _onnx_recurrent_activation_name( + _require_attr(layer, "activation_fn", label), + label, + allowed={"tanh", "sigmoid", "relu"}, + ) + + x_name = _unsqueeze_tensor( + builder, + builder.current_tensor, + builder.current_rank, + [1], + "recurrent_input_batch_%d" % layer_index, + ) + state_input_name = "state_%d_hidden_in" % layer_index + initial_h_name = _unsqueeze_tensor( + builder, + state_input_name, + 1, + [0, 1], + "recurrent_state_batch_%d" % layer_index, + ) + + w_name = builder.add_initializer( + "recurrent_w_%d" % layer_index, + np.expand_dims(weights.T, axis=0), + ) + r_name = builder.add_initializer( + "recurrent_r_%d" % layer_index, + np.expand_dims(recurrent_weights.T, axis=0), + ) + b_name = builder.add_initializer( + "recurrent_b_%d" % layer_index, + np.expand_dims( + np.concatenate( + [ + bias, + np.zeros(hidden_size, dtype=np.float32), + ] + ), + axis=0, + ), + ) + + sequence_output_name = builder._name("recurrent_sequence_%d" % layer_index) + hidden_output_name = builder._name("recurrent_hidden_%d" % layer_index) + _append_multi_output_node( + builder, + "RNN", + [x_name, w_name, r_name, b_name, "", initial_h_name], + [sequence_output_name, hidden_output_name], + "recurrent_node_%d" % layer_index, + activations=[activation_name], + hidden_size=hidden_size, + ) + + builder.current_tensor = _squeeze_tensor( + builder, + sequence_output_name, + 4, + [1, 2], + "recurrent_sequence_squeeze_%d" % layer_index, + ) + builder.current_rank = 2 + + state_output_name = _squeeze_tensor( + builder, + hidden_output_name, + 3, + [0, 1], + "recurrent_state_squeeze_%d" % layer_index, + ) + recurrent_state_interfaces.append( + (state_input_name, state_output_name, np.asarray(init, dtype=np.float32)) + ) + + +def _convert_lstm_layer( + builder: _PrimitiveGraphBuilder, + layer: object, + label: str, + layer_index: int, + recurrent_state_interfaces: list[tuple[str, str, np.ndarray]], +) -> None: + input_gate = _require_attr(layer, "input_gate", label) + forget_gate = _require_attr(layer, "forget_gate", label) + output_gate = _require_attr(layer, "output_gate", label) + cell = _require_attr(layer, "cell", label) + + input_weights = _as_float32_array( + _require_attr(input_gate, "weights", "%s.input_gate" % label), + "%s.input_gate.weights" % label, + ) + if input_weights.ndim != 2: + raise UnsupportedLayoutError( + "%s.input_gate.weights must be rank-2; got rank %d" + % (label, input_weights.ndim) + ) + input_size, hidden_size = input_weights.shape + + input_params = _validate_gate_shapes( + input_gate, + gate_name="input_gate", + label=label, + input_size=input_size, + hidden_size=hidden_size, + allow_peephole=True, + ) + forget_params = _validate_gate_shapes( + forget_gate, + gate_name="forget_gate", + label=label, + input_size=input_size, + hidden_size=hidden_size, + allow_peephole=True, + ) + cell_params = _validate_gate_shapes( + cell, + gate_name="cell", + label=label, + input_size=input_size, + hidden_size=hidden_size, + allow_peephole=False, + ) + output_params = _validate_gate_shapes( + output_gate, + gate_name="output_gate", + label=label, + input_size=input_size, + hidden_size=hidden_size, + allow_peephole=True, + ) + + input_activation = _onnx_recurrent_activation_name( + input_params[4], + "%s.input_gate" % label, + allowed={"sigmoid", "tanh", "relu"}, + ) + forget_activation = _onnx_recurrent_activation_name( + forget_params[4], + "%s.forget_gate" % label, + allowed={"sigmoid", "tanh", "relu"}, + ) + output_activation = _onnx_recurrent_activation_name( + output_params[4], + "%s.output_gate" % label, + allowed={"sigmoid", "tanh", "relu"}, + ) + if input_activation != forget_activation or input_activation != output_activation: + raise UnsupportedActivationError( + "%s requires identical input/forget/output gate activations; got %s, %s, %s" + % (label, input_activation, forget_activation, output_activation) + ) + + cell_activation = _onnx_recurrent_activation_name( + cell_params[4], + "%s.cell" % label, + allowed={"tanh", "sigmoid", "relu"}, + ) + state_activation = _onnx_recurrent_activation_name( + _require_attr(layer, "activation_fn", label), + label, + allowed={"tanh", "sigmoid", "relu"}, + ) + + init = _as_1d_float32_array(_require_attr(layer, "init", label), "%s.init" % label) + cell_init = _as_1d_float32_array( + _require_attr(layer, "cell_init", label), "%s.cell_init" % label + ) + _validate_state_size(init, hidden_size, "%s.init" % label) + _validate_state_size(cell_init, hidden_size, "%s.cell_init" % label) + + x_name = _unsqueeze_tensor( + builder, + builder.current_tensor, + builder.current_rank, + [1], + "lstm_input_batch_%d" % layer_index, + ) + hidden_state_input_name = "state_%d_hidden_in" % layer_index + cell_state_input_name = "state_%d_cell_in" % layer_index + initial_h_name = _unsqueeze_tensor( + builder, + hidden_state_input_name, + 1, + [0, 1], + "lstm_hidden_batch_%d" % layer_index, + ) + initial_c_name = _unsqueeze_tensor( + builder, + cell_state_input_name, + 1, + [0, 1], + "lstm_cell_batch_%d" % layer_index, + ) + + gate_weights = [ + input_params[0].T, + output_params[0].T, + forget_params[0].T, + cell_params[0].T, + ] + recurrent_gate_weights = [ + input_params[1].T, + output_params[1].T, + forget_params[1].T, + cell_params[1].T, + ] + gate_biases = [ + input_params[2], + output_params[2], + forget_params[2], + cell_params[2], + ] + + w_name = builder.add_initializer( + "lstm_w_%d" % layer_index, + np.expand_dims(np.concatenate(gate_weights, axis=0), axis=0), + ) + r_name = builder.add_initializer( + "lstm_r_%d" % layer_index, + np.expand_dims(np.concatenate(recurrent_gate_weights, axis=0), axis=0), + ) + b_name = builder.add_initializer( + "lstm_b_%d" % layer_index, + np.expand_dims( + np.concatenate( + [ + np.concatenate(gate_biases, axis=0), + np.zeros(hidden_size * 4, dtype=np.float32), + ] + ), + axis=0, + ), + ) + + peepholes = [input_params[3], output_params[3], forget_params[3]] + has_peepholes = any(item is not None for item in peepholes) + recurrent_inputs = [ + x_name, + w_name, + r_name, + b_name, + "", + initial_h_name, + initial_c_name, + ] + if has_peepholes: + p_values: list[np.ndarray] = [] + for item in peepholes: + if item is None: + p_values.append(np.zeros(hidden_size, dtype=np.float32)) + else: + p_values.append(item) + p_name = builder.add_initializer( + "lstm_p_%d" % layer_index, + np.expand_dims(np.concatenate(p_values, axis=0), axis=0), + ) + recurrent_inputs.append(p_name) + + sequence_output_name = builder._name("lstm_sequence_%d" % layer_index) + hidden_output_name = builder._name("lstm_hidden_%d" % layer_index) + cell_output_name = builder._name("lstm_cell_%d" % layer_index) + _append_multi_output_node( + builder, + "LSTM", + recurrent_inputs, + [sequence_output_name, hidden_output_name, cell_output_name], + "lstm_node_%d" % layer_index, + activations=[input_activation, cell_activation, state_activation], + hidden_size=hidden_size, + ) + + builder.current_tensor = _squeeze_tensor( + builder, + sequence_output_name, + 4, + [1, 2], + "lstm_sequence_squeeze_%d" % layer_index, + ) + builder.current_rank = 2 + + hidden_state_output = _squeeze_tensor( + builder, + hidden_output_name, + 3, + [0, 1], + "lstm_hidden_squeeze_%d" % layer_index, + ) + cell_state_output = _squeeze_tensor( + builder, + cell_output_name, + 3, + [0, 1], + "lstm_cell_squeeze_%d" % layer_index, + ) + + recurrent_state_interfaces.append( + ( + hidden_state_input_name, + hidden_state_output, + np.asarray(init, dtype=np.float32), + ) + ) + recurrent_state_interfaces.append( + ( + cell_state_input_name, + cell_state_output, + np.asarray(cell_init, dtype=np.float32), + ) + ) + + +def _convert_gru_layer( + builder: _PrimitiveGraphBuilder, + layer: object, + label: str, + layer_index: int, + recurrent_state_interfaces: list[tuple[str, str, np.ndarray]], +) -> None: + reset_gate = _require_attr(layer, "reset_gate", label) + update_gate = _require_attr(layer, "update_gate", label) + cell = _require_attr(layer, "cell", label) + + reset_weights = _as_float32_array( + _require_attr(reset_gate, "weights", "%s.reset_gate" % label), + "%s.reset_gate.weights" % label, + ) + if reset_weights.ndim != 2: + raise UnsupportedLayoutError( + "%s.reset_gate.weights must be rank-2; got rank %d" + % (label, reset_weights.ndim) + ) + input_size, hidden_size = reset_weights.shape + + reset_params = _validate_gate_shapes( + reset_gate, + gate_name="reset_gate", + label=label, + input_size=input_size, + hidden_size=hidden_size, + allow_peephole=False, + ) + update_params = _validate_gate_shapes( + update_gate, + gate_name="update_gate", + label=label, + input_size=input_size, + hidden_size=hidden_size, + allow_peephole=False, + ) + cell_params = _validate_gate_shapes( + cell, + gate_name="cell", + label=label, + input_size=input_size, + hidden_size=hidden_size, + allow_peephole=False, + ) + + gate_activation_reset = _onnx_recurrent_activation_name( + reset_params[4], + "%s.reset_gate" % label, + allowed={"sigmoid", "tanh", "relu"}, + ) + gate_activation_update = _onnx_recurrent_activation_name( + update_params[4], + "%s.update_gate" % label, + allowed={"sigmoid", "tanh", "relu"}, + ) + if gate_activation_reset != "Sigmoid" or gate_activation_update != "Sigmoid": + raise UnsupportedActivationError( + "%s requires Sigmoid reset/update gate activations; got %s and %s" + % (label, gate_activation_reset, gate_activation_update) + ) + + candidate_activation = _onnx_recurrent_activation_name( + cell_params[4], + "%s.cell" % label, + allowed={"tanh", "sigmoid", "relu"}, + ) + + init = _as_1d_float32_array(_require_attr(layer, "init", label), "%s.init" % label) + _validate_state_size(init, hidden_size, "%s.init" % label) + + x_name = _unsqueeze_tensor( + builder, + builder.current_tensor, + builder.current_rank, + [1], + "gru_input_batch_%d" % layer_index, + ) + state_input_name = "state_%d_hidden_in" % layer_index + initial_h_name = _unsqueeze_tensor( + builder, + state_input_name, + 1, + [0, 1], + "gru_state_batch_%d" % layer_index, + ) + + z_input_weights = -update_params[0].T + z_recurrent_weights = -update_params[1].T + z_bias = -update_params[2] + + w_name = builder.add_initializer( + "gru_w_%d" % layer_index, + np.expand_dims( + np.concatenate( + [ + z_input_weights, + reset_params[0].T, + cell_params[0].T, + ], + axis=0, + ), + axis=0, + ), + ) + r_name = builder.add_initializer( + "gru_r_%d" % layer_index, + np.expand_dims( + np.concatenate( + [ + z_recurrent_weights, + reset_params[1].T, + cell_params[1].T, + ], + axis=0, + ), + axis=0, + ), + ) + b_name = builder.add_initializer( + "gru_b_%d" % layer_index, + np.expand_dims( + np.concatenate( + [ + np.concatenate( + [ + z_bias, + reset_params[2], + cell_params[2], + ], + axis=0, + ), + np.zeros(hidden_size * 3, dtype=np.float32), + ] + ), + axis=0, + ), + ) + + sequence_output_name = builder._name("gru_sequence_%d" % layer_index) + hidden_output_name = builder._name("gru_hidden_%d" % layer_index) + _append_multi_output_node( + builder, + "GRU", + [x_name, w_name, r_name, b_name, "", initial_h_name], + [sequence_output_name, hidden_output_name], + "gru_node_%d" % layer_index, + activations=[gate_activation_reset, candidate_activation], + hidden_size=hidden_size, + linear_before_reset=1, + ) + + builder.current_tensor = _squeeze_tensor( + builder, + sequence_output_name, + 4, + [1, 2], + "gru_sequence_squeeze_%d" % layer_index, + ) + builder.current_rank = 2 + + state_output_name = _squeeze_tensor( + builder, + hidden_output_name, + 3, + [0, 1], + "gru_state_squeeze_%d" % layer_index, + ) + recurrent_state_interfaces.append( + (state_input_name, state_output_name, np.asarray(init, dtype=np.float32)) + ) + + +def _append_output_squeeze( + builder: _PrimitiveGraphBuilder, + tensor_name: str, + tensor_rank: int | None, + stem: str, +) -> TensorReference: + previous_tensor = builder.current_tensor + previous_rank = builder.current_rank + builder.current_tensor = tensor_name + builder.current_rank = tensor_rank + output_name = builder.add_node("Squeeze", [builder.current_tensor], stem) + builder.current_rank = None + builder.current_tensor = previous_tensor + builder.current_rank = previous_rank + return output_name, None + + +def _symbolic_shape(rank: int, prefix: str) -> list[str]: + if rank < 0: + raise ValueError("rank must be >= 0") + return ["%s_dim_%d" % (prefix, index) for index in range(rank)] + + +def _checker_compatible_shape(rank: int | None, prefix: str) -> list[str]: + if rank is None: + return ["%s_dim_0" % prefix] + return _symbolic_shape(rank, prefix) + + +def _require_single_tensor( + values: TensorReferences, + label: str, + what: str, +) -> TensorReference: + if len(values) != 1: + raise UnsupportedLayoutError( + "%s requires a single input tensor for %s; got %d tensors" + % (label, what, len(values)) + ) + return values[0] + + +def _set_builder_input( + builder: _PrimitiveGraphBuilder, + tensor: TensorReference, +) -> None: + builder.current_tensor = tensor[0] + builder.current_rank = tensor[1] + + +def _require_layer_list(layer: object, label: str, attr_name: str) -> list[object]: + value = _require_attr(layer, attr_name, label) + if not isinstance(value, (list, tuple)): + raise UnsupportedLayoutError( + "%s.%s must be a list or tuple of layers" % (label, attr_name) + ) + return list(value) + + +def _reverse_tensor_along_axis( + builder: _PrimitiveGraphBuilder, + tensor_name: str, + tensor_rank: int | None, + axis: int, + stem: str, +) -> TensorReference: + if tensor_rank is not None and (axis < 0 or axis >= tensor_rank): + raise UnsupportedLayoutError( + "cannot reverse axis %d for rank %d tensor" % (axis, tensor_rank) + ) + + starts_name = builder.add_initializer( + "%s_starts" % stem, + np.asarray([-1], dtype=np.int64), + ) + ends_name = builder.add_initializer( + "%s_ends" % stem, + np.asarray([np.iinfo(np.int64).min], dtype=np.int64), + ) + axes_name = builder.add_initializer( + "%s_axes" % stem, + np.asarray([axis], dtype=np.int64), + ) + steps_name = builder.add_initializer( + "%s_steps" % stem, + np.asarray([-1], dtype=np.int64), + ) + + previous_tensor = builder.current_tensor + previous_rank = builder.current_rank + builder.current_tensor = tensor_name + builder.current_rank = tensor_rank + output_name = builder.add_node( + "Slice", + [tensor_name, starts_name, ends_name, axes_name, steps_name], + stem, + ) + builder.current_tensor = previous_tensor + builder.current_rank = previous_rank + return output_name, tensor_rank + + +def _convert_layer_graph( + builder: _PrimitiveGraphBuilder, + layer: object, + input_values: TensorReferences, + recurrent_state_interfaces: list[tuple[str, str, np.ndarray]], + layer_counter: list[int], +) -> TensorReferences: + layer_index = layer_counter[0] + layer_counter[0] += 1 + label = _layer_label(layer_index, layer) + layer_type = layer.__class__.__name__ + + if layer_type == "SequentialLayer": + layers = _require_layer_list(layer, label, "layers") + current_values = list(input_values) + for sub_layer in layers: + current_values = _convert_layer_graph( + builder, + sub_layer, + current_values, + recurrent_state_interfaces, + layer_counter, + ) + return current_values + + if layer_type == "ParallelLayer": + layers = _require_layer_list(layer, label, "layers") + outputs: TensorReferences = [] + for branch_index, sub_layer in enumerate(layers): + branch_outputs = _convert_layer_graph( + builder, + sub_layer, + list(input_values), + recurrent_state_interfaces, + layer_counter, + ) + if len(branch_outputs) != 1: + raise UnsupportedLayoutError( + "%s branch %d must produce exactly one tensor; got %d" + % (label, branch_index, len(branch_outputs)) + ) + outputs.append(branch_outputs[0]) + return outputs + + if layer_type == "MultiTaskLayer": + layers = _require_layer_list(layer, label, "layers") + mapping_value = _require_attr(layer, "mapping", label) + if mapping_value is None: + mapping = None + elif isinstance(mapping_value, dict): + mapping = mapping_value + else: + raise UnsupportedLayoutError("%s.mapping must be a dict or None" % label) + + outputs = [] + for task_index, sub_layer in enumerate(layers): + if mapping is None: + input_index = task_index + else: + try: + mapping_index = mapping[task_index] + except KeyError as error: + raise UnsupportedLayoutError( + "%s.mapping is missing task index %d" % (label, task_index) + ) from error + input_index = _as_int( + mapping_index, + "%s.mapping[%d]" % (label, task_index), + ) + + if input_index < 0 or input_index >= len(input_values): + raise UnsupportedLayoutError( + "%s task %d selects input index %d, but only %d inputs are available" + % (label, task_index, input_index, len(input_values)) + ) + + task_outputs = _convert_layer_graph( + builder, + sub_layer, + [input_values[input_index]], + recurrent_state_interfaces, + layer_counter, + ) + if len(task_outputs) != 1: + raise UnsupportedLayoutError( + "%s task %d must produce exactly one tensor; got %d" + % (label, task_index, len(task_outputs)) + ) + outputs.append(task_outputs[0]) + return outputs + + if layer_type == "BidirectionalLayer": + source_tensor_name, source_rank = _require_single_tensor( + input_values, + label, + "bidirectional input", + ) + + fwd_layer = _require_attr(layer, "fwd_layer", label) + bwd_layer = _require_attr(layer, "bwd_layer", label) + + fwd_outputs = _convert_layer_graph( + builder, + fwd_layer, + [(source_tensor_name, source_rank)], + recurrent_state_interfaces, + layer_counter, + ) + if len(fwd_outputs) != 1: + raise UnsupportedLayoutError( + "%s.fwd_layer must produce exactly one tensor; got %d" + % (label, len(fwd_outputs)) + ) + + reversed_input_name, reversed_input_rank = _reverse_tensor_along_axis( + builder, + source_tensor_name, + source_rank, + axis=0, + stem="bidirectional_input_reverse_%d" % layer_index, + ) + bwd_outputs = _convert_layer_graph( + builder, + bwd_layer, + [(reversed_input_name, reversed_input_rank)], + recurrent_state_interfaces, + layer_counter, + ) + if len(bwd_outputs) != 1: + raise UnsupportedLayoutError( + "%s.bwd_layer must produce exactly one tensor; got %d" + % (label, len(bwd_outputs)) + ) + + bwd_name, bwd_rank = bwd_outputs[0] + fwd_name, fwd_rank = fwd_outputs[0] + if fwd_rank is not None and fwd_rank < 2: + raise UnsupportedLayoutError( + "%s output rank must be >= 2 for bidirectional concatenation" % label + ) + if fwd_rank is not None and bwd_rank is not None and fwd_rank != bwd_rank: + raise UnsupportedLayoutError( + "%s forward/backward rank mismatch: %d vs %d" + % (label, fwd_rank, bwd_rank) + ) + + reversed_bwd_name, _ = _reverse_tensor_along_axis( + builder, + bwd_name, + bwd_rank, + axis=0, + stem="bidirectional_output_reverse_%d" % layer_index, + ) + + previous_tensor = builder.current_tensor + previous_rank = builder.current_rank + builder.current_tensor = fwd_name + builder.current_rank = fwd_rank + output_name = builder.add_node( + "Concat", + [fwd_name, reversed_bwd_name], + "bidirectional_concat_%d" % layer_index, + axis=1, + ) + output_rank = builder.current_rank + builder.current_tensor = previous_tensor + builder.current_rank = previous_rank + return [(output_name, output_rank)] + + tensor_input = _require_single_tensor(input_values, label, layer_type) + _set_builder_input(builder, tensor_input) + + if layer_type == "FeedForwardLayer": + _convert_feed_forward_layer(builder, layer, label) + return [(builder.current_tensor, builder.current_rank)] + if layer_type == "BatchNormLayer": + _convert_batch_norm_layer(builder, layer, label) + return [(builder.current_tensor, builder.current_rank)] + if layer_type == "ReshapeLayer": + _convert_reshape_layer(builder, layer, label) + return [(builder.current_tensor, builder.current_rank)] + if layer_type == "TransposeLayer": + _convert_transpose_layer(builder, layer, label) + return [(builder.current_tensor, builder.current_rank)] + if layer_type == "PadLayer": + _convert_pad_layer(builder, layer, label) + return [(builder.current_tensor, builder.current_rank)] + if layer_type == "AverageLayer": + _convert_average_layer(builder, layer, label) + return [(builder.current_tensor, builder.current_rank)] + if layer_type == "ConvolutionalLayer": + _convert_convolutional_layer(builder, layer, label) + return [(builder.current_tensor, builder.current_rank)] + if layer_type == "MaxPoolLayer": + _convert_max_pool_layer(builder, layer, label) + return [(builder.current_tensor, builder.current_rank)] + if layer_type == "StrideLayer": + _convert_stride_layer(builder, layer, label) + return [(builder.current_tensor, builder.current_rank)] + if layer_type == "RecurrentLayer": + _convert_recurrent_layer( + builder, + layer, + label, + layer_index, + recurrent_state_interfaces, + ) + return [(builder.current_tensor, builder.current_rank)] + if layer_type == "LSTMLayer": + _convert_lstm_layer( + builder, + layer, + label, + layer_index, + recurrent_state_interfaces, + ) + return [(builder.current_tensor, builder.current_rank)] + if layer_type == "GRULayer": + _convert_gru_layer( + builder, + layer, + label, + layer_index, + recurrent_state_interfaces, + ) + return [(builder.current_tensor, builder.current_rank)] + if layer_type == "TCNBlock": + return _convert_tcn_block(builder, layer, label, layer_index) + if layer_type == "TCNLayer": + return _convert_tcn_layer( + builder, + layer, + label, + input_values, + recurrent_state_interfaces, + layer_counter, + ) + + raise NotImplementedError( + "layer %s is not implemented; converter supports structural layers " + "SequentialLayer, ParallelLayer, MultiTaskLayer, BidirectionalLayer " + "and primitive layers FeedForwardLayer, BatchNormLayer, ReshapeLayer, " + "TransposeLayer, PadLayer, AverageLayer, ConvolutionalLayer, " + "MaxPoolLayer, StrideLayer, RecurrentLayer, LSTMLayer, GRULayer, " + "TCNBlock, and TCNLayer" % layer_type + ) + + +def _build_onnx_model_from_primitive_layers( + layers: Sequence[object], + opset: int, +) -> object: + if opset < 13: + raise NotImplementedError("primitive conversion requires opset >= 13") + + onnx_module = _import_onnx_for_conversion() + builder = _PrimitiveGraphBuilder(onnx_module=onnx_module, opset=opset) + input_rank = _infer_graph_input_rank(layers) + builder.current_rank = input_rank + recurrent_state_interfaces: list[tuple[str, str, np.ndarray]] = [] + + layer_counter = [0] + outputs: TensorReferences = [("input", input_rank)] + for layer in layers: + outputs = _convert_layer_graph( + builder, + layer, + outputs, + recurrent_state_interfaces, + layer_counter, + ) + + if not outputs: + raise UnsupportedLayoutError("network produced no outputs") + + input_info = onnx_module.helper.make_tensor_value_info( + "input", + onnx_module.TensorProto.FLOAT, + _symbolic_shape(input_rank, "input"), + ) + + graph_inputs = [input_info] + graph_outputs = [] + for output_index, (output_name, output_rank) in enumerate(outputs): + final_name, final_rank = _append_output_squeeze( + builder, output_name, output_rank, "output_%d_squeeze" % output_index + ) + graph_outputs.append( + onnx_module.helper.make_tensor_value_info( + final_name, + onnx_module.TensorProto.FLOAT, + _checker_compatible_shape(final_rank, "output_%d" % output_index), + ) + ) + for state_input_name, state_output_name, state_init in recurrent_state_interfaces: + if state_init.ndim != 1: + raise UnsupportedLayoutError( + "state initializer for '%s' must be rank-1; got %s" + % (state_input_name, tuple(state_init.shape)) + ) + hidden_size = int(state_init.size) + graph_inputs.append( + onnx_module.helper.make_tensor_value_info( + state_input_name, + onnx_module.TensorProto.FLOAT, + [hidden_size], + ) + ) + graph_outputs.append( + onnx_module.helper.make_tensor_value_info( + state_output_name, + onnx_module.TensorProto.FLOAT, + [hidden_size], + ) + ) + + graph = onnx_module.helper.make_graph( + builder.nodes, + "madmom_primitive_network", + graph_inputs, + graph_outputs, + initializer=builder.initializers, + ) + + make_opsetid = getattr(onnx_module.helper, "make_opsetid", None) + if make_opsetid is None: + make_opsetid = getattr(onnx_module.helper, "make_operatorsetid") + opset_import = make_opsetid("", opset) + model = onnx_module.helper.make_model( + graph, + producer_name="madmom-model-converter", + opset_imports=[opset_import], + ) + model.ir_version = 8 + return model + + +def _extract_convertible_layers( + model_object: object, entry: ManifestEntry +) -> list[object]: + conversion_mode = entry.get("conversion_mode") + if conversion_mode in (None, "direct_nn"): + layers = getattr(model_object, "layers", None) + if isinstance(layers, list): + return list(layers) + raise NotImplementedError( + "direct_nn conversion requires an object with a list-like 'layers' attribute" + ) + + if conversion_mode == "wrapped_nn_core": + + def _looks_like_layer(obj: object) -> bool: + layer_name = obj.__class__.__name__ + return layer_name.endswith("Layer") or layer_name in { + "TCNBlock", + "TCNLayer", + } + + def _extract_wrapped_node(obj: object) -> object | None: + layers = getattr(obj, "layers", None) + if isinstance(layers, list): + return SequentialLayer(layers) + + processors = getattr(obj, "processors", None) + if isinstance(processors, (list, tuple)): + converted_children = [] + for processor in processors: + converted = _extract_wrapped_node(processor) + if converted is not None: + converted_children.append(converted) + if not converted_children: + return None + + if obj.__class__.__name__ == "ParallelProcessor": + return ParallelLayer(converted_children) + + return SequentialLayer(converted_children) + + if _looks_like_layer(obj): + return obj + + return None + + extracted = _extract_wrapped_node(model_object) + if extracted is not None: + if extracted.__class__.__name__ == "SequentialLayer": + return list(cast(Sequence[object], getattr(extracted, "layers"))) + return [extracted] + + raise NotImplementedError( + "wrapped_nn_core conversion could not find a nested object with 'layers' attribute" + ) + + raise NotImplementedError("unsupported conversion mode: %r" % conversion_mode) + + +def _convert_pickled_model_to_onnx( + model_object: object, + entry: ManifestEntry, + opset: int, +) -> object: + layers = _extract_convertible_layers(model_object, entry) + return _build_onnx_model_from_primitive_layers(layers, opset=opset) + + +def _default_converter(source_file: str, entry: ManifestEntry, opset: int) -> object: + with _temporary_repo_root_on_syspath(): + with _temporary_local_madmom_package(): + with _temporary_numpy_shape_base_alias(): + with open(source_file, "rb") as infile: + try: + model_object = pickle.load(infile, encoding="latin1") + except TypeError: + model_object = pickle.load(infile) + return _convert_pickled_model_to_onnx( + model_object=model_object, + entry=entry, + opset=opset, + ) + + +def _error_result( + *, + entry: ManifestEntry, + status: str, + error: Exception, +) -> ManifestEntry: + return { + "source_pkl": entry["source_pkl"], + "target_onnx": entry["target_onnx"], + "source_hash": entry["hash"], + "status": status, + "output_hash": None, + "checker_invoked": False, + "error": "%s: %s" % (error.__class__.__name__, error), + } + + +def convert_manifest_entries( + manifest: ManifestData, + models_dir: str = DEFAULT_MODELS_DIR, + converter: ConverterFn | None = None, + run_onnx_checker: bool = True, +) -> ManifestData: + artifacts = _as_list(manifest.get("artifacts")) + if artifacts is None: + raise ValueError("manifest must contain an artifacts list") + + converter_fn = converter if converter is not None else _default_converter + results: list[ManifestEntry] = [] + sorted_artifacts = sorted(artifacts, key=_artifact_source_sort_key) + for artifact in sorted_artifacts: + entry = _as_dict(artifact) + if entry is None: + raise ValueError("manifest artifacts must only contain objects") + + source_pkl = cast(Optional[str], entry.get("source_pkl")) + target_onnx = cast(Optional[str], entry.get("target_onnx")) + source_hash = cast(Optional[str], entry.get("hash")) + status = entry.get("status") + convertible = entry.get("convertible") + opset_value = entry.get("opset") + opset = opset_value if isinstance(opset_value, int) else DEFAULT_OPSET + + if source_pkl is None or source_hash is None: + raise ValueError("manifest artifact is missing required source fields") + + if status != "nn_convertible" or convertible is not True or target_onnx is None: + results.append( + { + "source_pkl": source_pkl, + "target_onnx": target_onnx, + "source_hash": source_hash, + "status": RESULT_SKIPPED_OUT_OF_SCOPE, + "output_hash": None, + "checker_invoked": False, + "error": None, + } + ) + continue + + source_file = os.path.join(models_dir, source_pkl.replace("/", os.sep)) + output_path = os.path.join(models_dir, target_onnx.replace("/", os.sep)) + + try: + converted_model = converter_fn(source_file, entry, opset) + output_hash = write_deterministic_onnx(converted_model, output_path) + except (pickle.UnpicklingError, EOFError) as error: + results.append( + _error_result(entry=entry, status=RESULT_CORRUPT_SOURCE, error=error) + ) + continue + except NotImplementedError as error: + results.append( + _error_result( + entry=entry, status=RESULT_CONVERTER_UNAVAILABLE, error=error + ) + ) + continue + except Exception as error: + results.append( + _error_result(entry=entry, status=RESULT_CONVERSION_ERROR, error=error) + ) + continue + + checker_invoked = False + try: + if run_onnx_checker: + checker_invoked = run_onnx_checker_if_available(output_path) + except Exception as error: + results.append( + _error_result(entry=entry, status=RESULT_CHECK_FAILED, error=error) + ) + continue + + results.append( + { + "source_pkl": source_pkl, + "target_onnx": target_onnx, + "source_hash": source_hash, + "status": RESULT_CONVERTED, + "output_hash": output_hash, + "checker_invoked": checker_invoked, + "error": None, + } + ) + + return { + "mode": MODE_CONVERT, + "schema_version": manifest.get("schema_version", 1), + "results": results, + } + + +def _emit_json(data: object, output_path: str | None) -> None: + if output_path: + _write_json(data, output_path) + return + json.dump(data, sys.stdout, indent=2, sort_keys=True) + _ = sys.stdout.write("\n") + + +def _conversion_exit_code(report: ManifestData) -> int: + results = _as_list(report.get("results")) + if results is None: + return 1 + non_success = { + RESULT_CORRUPT_SOURCE, + RESULT_CONVERTER_UNAVAILABLE, + RESULT_CONVERSION_ERROR, + RESULT_CHECK_FAILED, + } + for item in results: + item_dict = _as_dict(item) + if item_dict is None: + return 1 + status = item_dict.get("status") + if isinstance(status, str) and status in non_success: + return 1 + return 0 + + +def parse_args(args: Sequence[str] | None = None) -> CliArgs: + parser = argparse.ArgumentParser( + description="Inventory, check, and convert shipped pickle artifacts for NN-to-ONNX migration." + ) + _ = parser.add_argument( + "--dry-run-inventory", + action="store_true", + help="enumerate all shipped pickle artifacts and classify conversion scope", + ) + _ = parser.add_argument( + "--convert", + action="store_true", + help="run conversion for nn_convertible artifacts and emit structured results", + ) + _ = parser.add_argument( + "--check", + action="store_true", + help="validate an existing manifest JSON against the schema", + ) + _ = parser.add_argument( + "--models-dir", + default=DEFAULT_MODELS_DIR, + help="directory containing shipped model pickle files", + ) + _ = parser.add_argument( + "--schema", default=DEFAULT_SCHEMA_PATH, help="manifest schema path" + ) + _ = parser.add_argument( + "--manifest", + default=None, + help="output path for dry-run/convert JSON; input path for --check", + ) + _ = parser.add_argument( + "--opset", + type=int, + default=DEFAULT_OPSET, + help="target ONNX opset for nn_convertible artifacts", + ) + namespace = CliArgs() + _ = parser.parse_args(args=args, namespace=namespace) + return namespace + + +def main(args: Sequence[str] | None = None) -> int: + cli_args = parse_args(args=args) + + mode = determine_mode(cli_args) + + if mode == MODE_CHECK: + if cli_args.manifest is None: + raise SystemExit("--check requires --manifest") + schema = _load_json(cli_args.schema) + manifest = _load_json(cli_args.manifest) + errors = validate_check_input_shape(manifest, schema) + if errors: + for error in errors: + print(error, file=sys.stderr) + return 1 + return 0 + + if mode == MODE_CONVERT: + manifest = build_supported_artifact_matrix( + models_dir=cli_args.models_dir, opset=cli_args.opset + ) + report = convert_manifest_entries( + manifest=manifest, + models_dir=cli_args.models_dir, + run_onnx_checker=True, + ) + _emit_json(report, cli_args.manifest) + return _conversion_exit_code(report) + + manifest = build_supported_artifact_matrix( + models_dir=cli_args.models_dir, opset=cli_args.opset + ) + _emit_json(manifest, cli_args.manifest) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main())