max thesis code

Author: Maximilian Azevedo

users/azevedo/.gitignore

@@ -0,0 +1,2 @@
+**/.DS_Store                                                                            
+**/__pycache__

users/azevedo/__init__.py

@@ -0,0 +1,4 @@
+"""
+Uses as "root" point for hashing related to everything in "pytorch_networks"
+"""
+PACKAGE = __package__

users/azevedo/experiments/__init__.py

@@ -0,0 +1,179 @@
+"""
+Universal helpers to create configuration objects (i6_core ReturnnConfig) for RETURNN training/forwarding
+"""
+import copy
+from typing import Any, Dict, Optional
+
+from i6_core.returnn.config import ReturnnConfig, CodeWrapper
+
+from i6_experiments.common.setups.returnn_pytorch.serialization import (
+    Collection as TorchCollection,
+)
+from i6_experiments.common.setups.serialization import Import
+from .data.common import TrainingDatasets
+from .serializer import serialize_training, serialize_forward, PACKAGE
+
+
+def get_training_config(
+    training_datasets: TrainingDatasets,
+    network_module: str,
+    config: Dict[str, Any],
+    net_args: Dict[str, Any],
+    unhashed_net_args: Optional[Dict[str, Any]] = None,
+    include_native_ops=False,
+    debug: bool = False,
+    use_speed_perturbation: bool = False,
+    post_config: Optional[Dict[str, Any]] = None,
+) -> ReturnnConfig:
+    """
+    Get a generic config for training a model
+
+    :param training_datasets: datasets for training
+    :param network_module: path to the pytorch config file containing Model
+    :param net_args: extra arguments for constructing the PyTorch model
+    :param unhashed_net_args: unhashed extra arguments for constructing the PyTorch model
+    :param config: config arguments for RETURNN
+    :param debug: run training in debug mode (linking from recipe instead of copy)
+    :param use_speed_perturbation: Use speedperturbation in the training
+    :param post_config: Add non-hashed arguments for RETURNN
+    """
+
+    # changing these does not change the hash
+    base_post_config = {"stop_on_nonfinite_train_score": True, "num_workers_per_gpu": 2, "backend": "torch"}
+
+    # TODO: test
+    base_config = {
+        "cleanup_old_models": {
+            "keep_last_n": 4,
+            "keep_best_n": 4,
+            "keep": [1, 10, 125]
+        },
+        #############
+        "train": copy.deepcopy(training_datasets.train.as_returnn_opts()),
+        "dev": training_datasets.cv.as_returnn_opts(),
+        "eval_datasets": {"devtrain": training_datasets.devtrain.as_returnn_opts()},
+    }
+    config = {**base_config, **copy.deepcopy(config)}
+    post_config = {**base_post_config, **copy.deepcopy(post_config or {})}
+
+    serializer = serialize_training(
+        network_module=network_module,
+        net_args=net_args,
+        unhashed_net_args=unhashed_net_args,
+        include_native_ops=include_native_ops,
+        debug=debug,
+    )
+    python_prolog = None
+
+    # TODO: maybe make nice (if capability added to RETURNN itself)
+    if use_speed_perturbation:
+        prolog_serializer = TorchCollection(
+            serializer_objects=[
+                Import(
+                    code_object_path=PACKAGE + ".extra_code.speed_perturbation.legacy_speed_perturbation",
+                    unhashed_package_root=PACKAGE,
+                )
+            ]
+        )
+        python_prolog = [prolog_serializer]
+        config["train"]["datasets"]["zip_dataset"]["audio"]["pre_process"] = CodeWrapper("legacy_speed_perturbation")
+
+    returnn_config = ReturnnConfig(
+        config=config, post_config=post_config, python_prolog=python_prolog, python_epilog=[serializer]
+    )
+    return returnn_config
+
+
+def get_prior_config(
+    training_datasets: TrainingDatasets,  # TODO: replace by single dataset
+    network_module: str,
+    config: Dict[str, Any],
+    net_args: Dict[str, Any],
+    unhashed_net_args: Optional[Dict[str, Any]] = None,
+    debug: bool = False,
+):
+    """
+    Get a generic config for extracting output label priors
+
+    :param training_datasets: datasets for training
+    :param network_module: path to the pytorch config file containing Model
+    :param config: config arguments for RETURNN
+    :param net_args: extra arguments for constructing the PyTorch model
+    :param unhashed_net_args: unhashed extra arguments for constructing the PyTorch model
+    :param debug: run training in debug mode (linking from recipe instead of copy)
+    """
+
+    # changing these does not change the hash
+    post_config = {
+        "num_workers_per_gpu": 2,
+    }
+
+    base_config = {
+        #############
+        "batch_size": 500 * 16000,
+        "max_seqs": 240,
+        #############
+        "forward": copy.deepcopy(training_datasets.prior.as_returnn_opts()),
+    }
+    config = {**base_config, **copy.deepcopy(config)}
+    post_config["backend"] = "torch"
+
+    serializer = serialize_forward(
+        network_module=network_module,
+        net_args=net_args,
+        unhashed_net_args=unhashed_net_args,
+        forward_module=None,  # same as network
+        forward_step_name="prior",
+        forward_init_args=None,
+        unhashed_forward_init_args=None,
+        debug=debug,
+    )
+    returnn_config = ReturnnConfig(config=config, post_config=post_config, python_epilog=[serializer])
+    return returnn_config
+
+
+def get_forward_config(
+    network_module: str,
+    config: Dict[str, Any],
+    net_args: Dict[str, Any],
+    decoder: str,
+    decoder_args: Dict[str, Any],
+    unhashed_decoder_args: Optional[Dict[str, Any]] = None,
+    unhashed_net_args: Optional[Dict[str, Any]] = None,
+    debug: bool = False,
+) -> ReturnnConfig:
+    """
+    Get a generic config for forwarding
+
+    :param network_module: path to the pytorch config file containing Model
+    :param net_args: extra arguments for constructing the PyTorch model
+    :param decoder: which (python) file to load which defines the forward, forward_init and forward_finish functions
+    :param decoder_args: extra arguments to pass to forward_init
+    :param config: config arguments for RETURNN
+    :param unhashed_decoder_args: unhashed extra arguments for the forward init
+    :param unhashed_net_args: unhashed extra arguments for constructing the PyTorch model
+    :param debug: run training in debug mode (linking from recipe instead of copy)
+    """
+
+    # changing these does not change the hash
+    post_config = {}
+
+    # changeing these does change the hash
+    base_config = {
+        "batch_size": 1000 * 16000,
+        "max_seqs": 240,
+    }
+    config = {**base_config, **copy.deepcopy(config)}
+    post_config["backend"] = "torch"
+
+    serializer = serialize_forward(
+        network_module=network_module,
+        net_args=net_args,
+        unhashed_net_args=unhashed_net_args,
+        forward_module=decoder,
+        forward_init_args=decoder_args,
+        unhashed_forward_init_args=unhashed_decoder_args,
+        debug=debug,
+    )
+    returnn_config = ReturnnConfig(config=config, post_config=post_config, python_epilog=[serializer])
+    return returnn_config

users/azevedo/experiments/librispeech/__init__.py

@@ -0,0 +1,106 @@
+"""
+Dataset helpers for the BPE-based training
+"""
+from sisyphus import tk
+
+from i6_core.g2p.convert import BlissLexiconToG2PLexiconJob
+from i6_core.lexicon.bpe import CreateBPELexiconJob
+
+from i6_experiments.common.datasets.librispeech import get_ogg_zip_dict, get_bliss_lexicon
+from i6_experiments.common.datasets.librispeech.vocab import get_subword_nmt_bpe_v2
+from i6_experiments.common.setups.returnn.datastreams.vocabulary import BpeDatastream
+
+from .common import DatasetSettings, TrainingDatasets, build_training_datasets
+from ..default_tools import MINI_RETURNN_ROOT, RETURNN_EXE, SUBWORD_NMT_REPO
+
+
+def get_bpe_datastream(librispeech_key: str, bpe_size: int, is_recog: bool, use_postfix: bool) -> BpeDatastream:
+    """
+    Returns the datastream for the bpe labels
+
+    Uses the legacy BPE setup that is compatible with old LM models
+
+    :param librispeech_key: which librispeech corpus to use for bpe training
+    :param bpe_size: size for the bpe labels
+    :param is_recog: removes the UNK label when not in training
+    :param use_postfix: True for RNN-T or Attention, False for CTC
+    """
+    bpe_settings = get_subword_nmt_bpe_v2(corpus_key=librispeech_key, bpe_size=bpe_size, unk_label="<unk>")
+
+    bpe_targets = BpeDatastream(
+        available_for_inference=False,
+        bpe_settings=bpe_settings,
+        use_unk_label=is_recog,
+        seq_postfix=0 if use_postfix else None,
+    )
+    return bpe_targets
+
+
+def get_bpe_lexicon(librispeech_key: str, bpe_size: int) -> tk.Path:
+    """
+    Create BPE lexicon without unknown and silence
+
+    :param librispeech_key: which librispeech corpus to use for bpe training
+    :param bpe_size: number of BPE splits
+    :return: path to a lexicon bliss xml file
+    """
+    bpe_settings = get_subword_nmt_bpe_v2(corpus_key=librispeech_key, bpe_size=bpe_size, unk_label="<unk>")
+    bpe_lexicon = CreateBPELexiconJob(
+        base_lexicon_path=get_bliss_lexicon(add_unknown_phoneme_and_mapping=False, add_silence=False),
+        bpe_codes=bpe_settings.bpe_codes,
+        bpe_vocab=bpe_settings.bpe_vocab,
+        subword_nmt_repo=SUBWORD_NMT_REPO,
+        unk_label="<unk>",
+    ).out_lexicon
+
+    return bpe_lexicon
+
+
+def get_text_lexicon(prefix: str, librispeech_key: str, bpe_size: int) -> tk.Path:
+    """
+    Get a bpe lexicon in line-based text format to be used for torchaudio/Flashlight decoding
+
+    :param prefix:
+    :param librispeech_key: which librispeech corpus to use for bpe training
+    :param bpe_size: number of BPE splits
+    :return: path to a lexicon text file
+    """
+    bliss_lex = get_bpe_lexicon(librispeech_key=librispeech_key, bpe_size=bpe_size)
+    word_lexicon = BlissLexiconToG2PLexiconJob(
+        bliss_lex,
+        include_pronunciation_variants=True,
+        include_orthography_variants=True,
+    ).out_g2p_lexicon
+    return word_lexicon
+
+
+def build_bpe_training_datasets(
+    prefix: str,
+    librispeech_key: str,
+    bpe_size: int,
+    settings: DatasetSettings,
+    use_postfix: bool,
+) -> TrainingDatasets:
+    """
+
+    :param librispeech_key: which librispeech corpus to use for bpe training
+    :param bpe_size: number of BPE splits
+    :param settings: configuration object for the dataset pipeline
+    :param use_postfix: True for RNN-T or Attention, False for CTC
+    """
+    label_datastream = get_bpe_datastream(
+        librispeech_key=librispeech_key, bpe_size=bpe_size, is_recog=False, use_postfix=use_postfix
+    )
+
+    ogg_zip_dict = get_ogg_zip_dict(prefix, returnn_root=MINI_RETURNN_ROOT, returnn_python_exe=RETURNN_EXE)
+    train_ogg = ogg_zip_dict[librispeech_key]
+    dev_clean_ogg = ogg_zip_dict["dev-clean"]
+    dev_other_ogg = ogg_zip_dict["dev-other"]
+
+    return build_training_datasets(
+        train_ogg=train_ogg,
+        dev_clean_ogg=dev_clean_ogg,
+        dev_other_ogg=dev_other_ogg,
+        settings=settings,
+        label_datastream=label_datastream,
+    )