standalone 2024 setup add LSTM lm pipeline

Author: JackTemaki

users/rossenbach/experiments/librispeech/ctc_rnnt_standalone_2024/data/bpe_lm.py

@@ -0,0 +1,186 @@
+from sisyphus import tk
+from sisyphus.delayed_ops import DelayedFormat
+
+from dataclasses import dataclass
+import os
+from typing import Any, Dict, Optional
+
+from i6_core.text.label.subword_nmt.apply import ApplyBPEToTextJob
+from i6_core.corpus.convert import CorpusToTxtJob
+from i6_core.text.processing import ConcatenateJob
+from i6_core.returnn.config import CodeWrapper
+
+from i6_experiments.common.setups.returnn.datasets import MetaDataset, ControlDataset, Dataset
+from i6_experiments.common.setups.returnn.datastreams.base import Datastream
+from i6_experiments.common.setups.returnn.datastreams.vocabulary import BpeDatastream
+from i6_experiments.common.helpers.text_labels.subword_nmt_bpe import get_returnn_subword_nmt
+
+from i6_experiments.common.datasets.librispeech import get_bliss_corpus_dict
+from i6_experiments.common.datasets.librispeech.vocab import get_subword_nmt_bpe_v2
+from i6_experiments.common.datasets.librispeech.language_model import get_librispeech_normalized_lm_data
+
+
+
+SOURCE_DATASTREAM_KEY = "data"
+TARGET_DATASTREAN_KEY = "delayed"
+
+
+@dataclass(frozen=True)
+class TrainingDatasets:
+    train: Dataset
+    cv: Dataset
+    devtrain: Dataset
+    datastreams: Dict[str, Datastream]
+
+
+class LmDataset(ControlDataset):
+
+    def __init__(
+        self,
+        *,
+        corpus_file: tk.Path,
+        vocab_file: tk.Path,
+        # super parameters
+        partition_epoch: Optional[int] = None,
+        segment_file: Optional[tk.Path] = None,
+        seq_ordering: Optional[str] = None,
+        random_subset: Optional[int] = None,
+        additional_options: Optional[Dict] = None,
+    ):
+        super().__init__(
+            partition_epoch=partition_epoch,
+            segment_file=segment_file,
+            seq_ordering=seq_ordering,
+            random_subset=random_subset,
+            additional_options=additional_options
+        )
+
+        self.corpus_file = corpus_file
+        self.vocab_file = vocab_file
+
+    def as_returnn_opts(self) -> Dict[str, Any]:
+        d = {
+            "class": "LmDataset",
+            "corpus_file": CodeWrapper(DelayedFormat('lambda: cf("{}")', self.corpus_file)),
+            "orth_symbols_map_file": self.vocab_file,
+            "orth_replace_map_file": "",
+            "word_based": True,
+            "seq_end_symbol": "</s>",
+            "auto_replace_unknown_symbol": False,
+            "unknown_symbol": "<unk>",
+            "add_delayed_seq_data": True,
+            "delayed_seq_data_start_symbol": "<s>",
+        }
+        sd = super().as_returnn_opts()
+        assert all([k not in sd.keys() for k in d.keys()]), (
+            "conflicting keys in %s and %s"
+            % (str(list(sd.keys())), str(list(d.keys()))),
+        )
+        d.update(sd)
+
+        return d
+
+@dataclass()
+class LMDatasetSettings:
+    train_partition_epoch: int
+    train_seq_ordering: str
+
+
+def get_subword_repo():
+    """
+    This is a for now very ugly helper to get the same subword_nmt repo
+    as the get_subword_nmt_bpe_v2 is using
+    :return: 
+    """
+    subword_nmt_repo = get_returnn_subword_nmt(
+        commit_hash="5015a45e28a958f800ef1c50e7880c0c9ef414cf", output_prefix=""
+    )
+    # overwrite hash for future bugfixes, it is unlikely the logic will ever be changed
+    subword_nmt_repo.hash_overwrite = "I6_SUBWORD_NMT_V2"
+    return subword_nmt_repo
+
+def build_lm_training_datasets(prefix, librispeech_key, bpe_size, settings: LMDatasetSettings):
+    
+    #data_map = {SOURCE_DATASTREAM_KEY: ("lm_dataset", "data"), TARGET_DATASTREAN_KEY: ("lm_dataset", "delayed")}
+    #def make_meta(dataset: LmDataset):
+    #    return MetaDataset(
+    #        data_map=data_map, datasets={"lm_dataset": dataset}, seq_order_control_dataset="lm_dataset"
+    #    )
+    
+    bpe_settings = get_subword_nmt_bpe_v2(corpus_key=librispeech_key, bpe_size=bpe_size, unk_label='<unk>')
+    ls_bliss_corpus_dict = get_bliss_corpus_dict()
+    bpe_datastream = BpeDatastream(available_for_inference=False, bpe_settings=bpe_settings)
+
+     #### Training Data ####
+
+    lm_data = get_librispeech_normalized_lm_data()
+    ls_train_bliss = ls_bliss_corpus_dict["train-other-960"]
+    ls_train_text = CorpusToTxtJob(
+        bliss_corpus=ls_train_bliss,
+        gzip=True,
+    ).out_txt
+    full_train_text = ConcatenateJob(
+        text_files=[lm_data, ls_train_text],
+        zip_out=True,
+    ).out
+    lm_bpe_data_job = ApplyBPEToTextJob(
+        text_file=full_train_text,
+        bpe_codes=bpe_settings.bpe_codes,
+        bpe_vocab=bpe_settings.bpe_count_vocab,
+        gzip_output=True,
+        subword_nmt_repo=get_subword_repo(),
+        mini_task=False,  # this is a large file, so run in cluster
+    )
+    lm_bpe_data_job.add_alias(os.path.join(prefix, "apply_bpe_to_train"))
+
+    #### Dev Data ####
+
+    dev_clean_text = CorpusToTxtJob(bliss_corpus=ls_bliss_corpus_dict["dev-clean"], gzip=True).out_txt
+    dev_other_text = CorpusToTxtJob(bliss_corpus=ls_bliss_corpus_dict["dev-other"], gzip=True).out_txt
+    cv_text = ConcatenateJob(
+        text_files=[dev_clean_text, dev_other_text],
+        zip_out=True,
+    ).out
+    cv_bpe_data_job = ApplyBPEToTextJob(
+        text_file=cv_text,
+        bpe_codes=bpe_settings.bpe_codes,
+        bpe_vocab=bpe_settings.bpe_count_vocab,
+        gzip_output=True,
+        subword_nmt_repo=get_subword_repo(),
+    )
+
+    #### datasets ####
+    lm_train_dataset = LmDataset(
+        corpus_file=lm_bpe_data_job.out_bpe_text,
+        vocab_file=bpe_settings.bpe_vocab,
+        partition_epoch=settings.train_partition_epoch,
+        segment_file=None,
+        seq_ordering=settings.train_seq_ordering
+    )
+
+    lm_cv_dataset = LmDataset(
+        corpus_file=cv_bpe_data_job.out_bpe_text,
+        vocab_file=bpe_settings.bpe_vocab,
+        partition_epoch=1,
+        segment_file=None,
+        seq_ordering="sorted"
+    )
+
+    lm_devtrain_dataset = LmDataset(
+        corpus_file=lm_bpe_data_job.out_bpe_text,
+        vocab_file=bpe_settings.bpe_vocab,
+        partition_epoch=1,
+        segment_file=None,
+        seq_ordering="sorted",
+        random_subset=3000,
+    )
+
+    return TrainingDatasets(
+        train=lm_train_dataset,
+        cv=lm_cv_dataset,
+        # devtrain=lm_devtrain_dataset,
+        # TODO: Ultra hack for now
+        devtrain=lm_cv_dataset,
+        datastreams={"data": bpe_datastream, "delayed": bpe_datastream},
+    )
+

users/rossenbach/experiments/librispeech/ctc_rnnt_standalone_2024/experiments/lm_bpe/__init__.py

@@ -0,0 +1,77 @@
+from sisyphus import tk
+
+from dataclasses import asdict
+from typing import cast
+
+from i6_experiments.common.setups.returnn.datastreams.vocabulary import LabelDatastream
+
+from ...data.bpe_lm import build_lm_training_datasets, LMDatasetSettings
+from ...default_tools import RETURNN_EXE, MINI_RETURNN_ROOT
+from ...pipeline import training
+
+
+def bpe_kazuki_lstm():
+    prefix_name = "experiments/librispeech/ctc_rnnt_standalone_2024/kazuki_lstm/"
+
+    train_settings = LMDatasetSettings(
+        train_partition_epoch=4,
+        train_seq_ordering="laplace:.100",
+    )
+
+    # build the training datasets object containing train, cv, dev-train and the extern_data dict
+    train_data_bpe10k = build_lm_training_datasets(
+        prefix=prefix_name,
+        librispeech_key="train-other-960",
+        bpe_size=10000,
+        settings=train_settings,
+    )
+    label_datastream_bpe5000 = cast(LabelDatastream, train_data_bpe10k.datastreams["data"])
+    vocab_size_without_blank = label_datastream_bpe5000.vocab_size
+
+    default_returnn = {
+        "returnn_exe": RETURNN_EXE,
+        "returnn_root": MINI_RETURNN_ROOT,
+    }
+
+    from ...pytorch_networks.lm.lstm.kazuki_lstm_zijian_variant_v1_cfg import ModelConfig
+
+    default_init_args = {
+        'init_args_w': {'func': 'normal', 'arg': {'mean': 0.0, 'std': 0.1}},
+        'init_args_b': {'func': 'normal', 'arg': {'mean': 0.0, 'std': 0.1}}
+    }
+
+    lstm_base_config = ModelConfig(
+        vocab_dim=vocab_size_without_blank,
+        embed_dim=512,
+        hidden_dim=2048,
+        n_lstm_layers=2,
+        use_bottle_neck=False,
+        dropout=0.2,
+        init_args=default_init_args,
+    )
+
+    train_config_24gbgpu = {
+        "optimizer": {"class": "SGD"},
+        #############
+        "batch_size": 1280,  # BPE tokens
+        "accum_grad_multiple_step": 1,
+        "learning_rate": 1.0,
+        "decay": 0.8,
+        "multi_num_epochs": train_settings.train_partition_epoch,
+        "relative_error_threshold": 0,
+        "multi_update_interval": 1,
+        "error_measure": "dev_ce",
+    }
+
+    network_module = "lm.lstm.kazuki_lstm_zijian_variant_v1"
+    train_args = {
+        "config": train_config_24gbgpu,
+        "network_module": network_module,
+        "net_args": {"model_config_dict": asdict(lstm_base_config)},
+        "debug": False,
+        "add_cache_manager": True,
+    }
+
+    training_name = prefix_name + "/" + network_module + ".512dim_sub6_24gbgpu_50eps"
+    train_job = training(training_name, train_data_bpe10k, train_args, num_epochs=30, **default_returnn)
+    train_job.rqmt["gpu_mem"] = 24

users/rossenbach/experiments/librispeech/ctc_rnnt_standalone_2024/experiments/lm_bpe/kazuki_lstm.py

@@ -0,0 +1,108 @@
+import torch
+from torch import nn
+
+from .kazuki_lstm_zijian_variant_v1_cfg import ModelConfig
+
+def mask_tensor(tensor: torch.Tensor, seq_len: torch.Tensor) -> torch.Tensor:
+    """
+    mask a tensor with a "positive" mask (boolean true means position is used)
+
+    :param tensor: [B,T,....]
+    :param seq_len: [B]
+    :return: [B,T] as boolean
+    """
+    seq_len = seq_len.to(device=tensor.device)
+    r = torch.arange(tensor.shape[1], device=tensor.device)  # [T]
+    seq_mask = torch.less(r[None, :], seq_len[:, None])  # broadcast to [B,T]
+    return seq_mask
+
+
+class Model(nn.Module):
+    """
+    Simple LSTM LM with an embedding, an LSTM, and a final linear
+    """
+    def __init__(self, model_config_dict, **kwargs):
+        super().__init__()
+        self.cfg = ModelConfig(**model_config_dict)
+        if self.cfg.dropout > 0:
+            self.dropout = nn.Dropout(p=self.cfg.dropout)
+        else:
+            self.dropout = None
+        self.use_bottle_neck = self.cfg.use_bottle_neck
+        self.embed = nn.Embedding(self.cfg.vocab_dim, self.cfg.embed_dim)
+        self.lstm = nn.LSTM(
+            input_size=self.cfg.embed_dim,
+            hidden_size=self.cfg.hidden_dim,
+            num_layers=self.cfg.n_lstm_layers,
+            bias=self.cfg.bias,
+            batch_first=True,
+            dropout=self.cfg.dropout,
+            bidirectional=False,
+        )
+        if self.cfg.use_bottle_neck:
+            self.bottle_neck = nn.Linear(self.cfg.hidden_dim,self.cfg.bottle_neck_dim, bias=True)
+            self.final_linear = nn.Linear(self.cfg.bottle_neck_dim, self.cfg.vocab_dim, bias=True)
+        else:
+            self.final_linear = nn.Linear(self.cfg.hidden_dim, self.cfg.vocab_dim, bias=True)
+        self._param_init(**self.cfg.init_args)
+
+
+    def _param_init(self, init_args_w=None, init_args_b=None):
+        if init_args_w is None:
+            init_args_w = {'func': 'normal', 'arg': {'mean': 0.0, 'std': 0.1}}
+        if init_args_b is None:
+            init_args_b = {'func': 'normal', 'arg': {'mean': 0.0, 'std': 0.1}}
+
+        for m in self.modules():
+
+            for name, param in m.named_parameters():
+                if 'bias' in name:
+                    if init_args_b['func'] == 'normal':
+                        init_func = nn.init.normal_
+                    else:
+                        NotImplementedError
+                    hyp = init_args_b['arg']
+                else:
+                    if init_args_w['func'] == 'normal':
+                        init_func = nn.init.normal_
+                    else:
+                        NotImplementedError
+                    hyp = init_args_w['arg']
+                init_func(param, **hyp)
+
+    def forward(self, x):
+        """
+        Return logits of each batch at each time step
+        x: (B, S, F)
+        """
+        x = self.embed(x)
+        if self.dropout:
+            x = self.dropout(x)
+        batch_size = x.shape[0]
+        h0 = torch.zeros((self.cfg.n_lstm_layers, batch_size, self.cfg.hidden_dim), device=x.device).detach()
+        c0 = torch.zeros_like(h0, device=x.device).detach()
+        # This is a uni-directional LSTM, so sequence masking is not necessary
+        x, _ = self.lstm(x, (h0, c0))
+        if self.dropout:
+            x = self.dropout(x)
+        if self.use_bottle_neck:
+            x = self.bottle_neck(x)
+            if self.dropout:
+                x = self.dropout(x)
+        x = self.final_linear(x)
+        return x
+    
+    
+def train_step(*, model: Model, data, run_ctx, **kwargs):
+    labels = data["ldata"]
+    labels_len = data["data:size1"]
+    delayed_labels = data["delayed"]
+
+    lm_logits = model(delayed_labels)  # (B, S, F)
+
+    ce_loss = torch.nn.functional.cross_entropy(lm_logits.transpose(1, 2), labels, reduction='none')
+    seq_mask = mask_tensor(labels, labels_len)
+    ce_loss = (ce_loss * seq_mask).sum()
+    total_length = torch.sum(labels_len)
+    
+    run_ctx.mark_as_loss(name="ce", loss=ce_loss, inv_norm_factor=total_length)

users/rossenbach/experiments/librispeech/ctc_rnnt_standalone_2024/pytorch_networks/lm/__init__.py

@@ -0,0 +1,15 @@
+from dataclasses import dataclass
+
+from i6_models.config import ModelConfiguration
+
+@dataclass
+class ModelConfig():
+    vocab_dim: int
+    embed_dim: int
+    hidden_dim: int
+    n_lstm_layers: int
+    init_args: dict
+    bias: bool = True
+    use_bottle_neck: bool = False
+    bottle_neck_dim: int = 512
+    dropout: float = 0.0