Merge branch 'persephone-refactor' of https://github.com/hyperion-ml/hyperion into persephone-asr

addding changes in persephone-refactor to persephone-asr

Author: jesus-villalba

hyperion/bin/extract_wav2vec2xvectors.py

@@ -19,6 +19,7 @@
 import pandas as pd
 
 import torch
+import torchaudio.transforms as tat
 
 from hyperion.hyp_defs import config_logger, float_cpu, set_float_cpu
 from hyperion.utils import Utt2Info
@@ -30,6 +31,25 @@
 from hyperion.torch.utils import open_device
 from hyperion.torch import TorchModelLoader as TML
 
+resamplers = {}
+
+
+def get_resampler(source_fs, target_fs):
+    if source_fs in resamplers:
+        return resamplers[source_fs]
+
+    resampler = tat.Resample(
+        int(source_fs),
+        int(target_fs),
+        lowpass_filter_width=64,
+        rolloff=0.9475937167399596,
+        resampling_method="kaiser_window",
+        beta=14.769656459379492,
+    )
+    resampler_f = lambda x: resampler(torch.from_numpy(x)).numpy()
+    resamplers[source_fs] = resampler_f
+    return resampler_f
+
 
 def init_device(use_gpu):
     set_float_cpu("float32")
@@ -102,7 +122,7 @@ def extract_xvectors(
     num_augs,
     aug_info_path,
     use_gpu,
-    **kwargs
+    **kwargs,
 ):
 
     rng = np.random.RandomState(seed=1123581321 + kwargs["part_idx"])
@@ -122,12 +142,11 @@ def extract_xvectors(
         num_augs = 1
 
     ar_args = AR.filter_args(**kwargs)
+    ar_args["wav_scale"] = 1.0
     logging.info("opening output stream: %s", output_spec)
     with DWF.create(output_spec, scp_sep=scp_sep) as writer:
 
-        logging.info(
-            "opening input stream: {} with args={}".format(input_spec, ar_args)
-        )
+        logging.info(f"opening input stream: {input_spec} with args={ar_args}")
         with AR(input_spec, **ar_args) as reader:
 
             if vad_spec is not None:
@@ -146,6 +165,11 @@ def extract_xvectors(
                 key0 = key[0]
                 fs = fs[0]
                 t2 = time.time()
+                if fs != model.sample_frequency:
+                    resampler = get_resampler(fs, model.sample_frequency)
+                    print(f"x01 {x0.shape} {np.max(x0)}")
+                    x0 = resampler(x0)
+                    print(f"x01 {x0.shape} {np.max(x0)}")
 
                 logging.info("processing utt %s", key0)
                 for aug_id in range(num_augs):

hyperion/bin/finetune_wav2vec2xvector.py

@@ -25,7 +25,8 @@
 from hyperion.torch.utils import ddp
 from hyperion.torch.trainers import XVectorTrainer as Trainer
 from hyperion.torch.data import AudioDataset as AD
-from hyperion.torch.data import ClassWeightedSeqSampler as Sampler
+from hyperion.torch.data import SegSamplerFactory
+
 from hyperion.torch.metrics import CategoricalAccuracy
 from hyperion.torch.models import (
     HFWav2Vec2ResNet1dXVector,
@@ -45,19 +46,21 @@ def init_data(partition, rank, num_gpus, **kwargs):
 
     kwargs = kwargs["data"][partition]
     ad_args = AD.filter_args(**kwargs["dataset"])
-    sampler_args = Sampler.filter_args(**kwargs["sampler"])
+    sampler_args = kwargs["sampler"]
     if rank == 0:
         logging.info("{} audio dataset args={}".format(partition, ad_args))
         logging.info("{} sampler args={}".format(partition, sampler_args))
         logging.info("init %s dataset", partition)
 
-    ad_args["is_val"] = partition == "val"
+    is_val = partition == "val"
+    ad_args["is_val"] = is_val
+    sampler_args["shuffle"] = not is_val
     dataset = AD(**ad_args)
 
     if rank == 0:
         logging.info("init %s samplers", partition)
 
-    sampler = Sampler(dataset, **sampler_args)
+    sampler = SegSamplerFactory.create(dataset, **sampler_args)
 
     if rank == 0:
         logging.info("init %s dataloader", partition)
@@ -71,18 +74,6 @@ def init_data(partition, rank, num_gpus, **kwargs):
     return data_loader
 
 
-# def init_model(num_classes, in_model_file, rank, **kwargs):
-#     xvec_args = kwargs["model"]["xvector"]
-#     if rank == 0:
-#         logging.info("xvector network ft args={}".format(xvec_args))
-#     xvec_args["num_classes"] = num_classes
-#     model = TML.load(in_model_file)
-#     model.rebuild_output_layer(**xvec_args)
-#     if rank == 0:
-#         logging.info("model={}".format(model))
-#     return model
-
-
 def init_model(num_classes, in_model_file, rank, **kwargs):
     model_args = kwargs["model"]
     if rank == 0:
@@ -127,19 +118,15 @@ def train_model(gpu_id, args):
 
     train_loader = init_data(partition="train", **kwargs)
     val_loader = init_data(partition="val", **kwargs)
-    model = init_model(train_loader.dataset.num_classes, **kwargs)
+    model = init_model(list(train_loader.dataset.num_classes.values())[0], **kwargs)
     init_hard_prototype_mining(model, train_loader, val_loader, rank)
 
     trn_args = Trainer.filter_args(**kwargs["trainer"])
     if rank == 0:
         logging.info("trainer args={}".format(trn_args))
     metrics = {"acc": CategoricalAccuracy()}
     trainer = Trainer(
-        model,
-        device=device,
-        metrics=metrics,
-        ddp=world_size > 1,
-        **trn_args,
+        model, device=device, metrics=metrics, ddp=world_size > 1, **trn_args,
     )
     trainer.load_last_checkpoint()
     trainer.fit(train_loader, val_loader)
@@ -153,7 +140,7 @@ def make_parser(model_class):
     parser.add_argument("--cfg", action=ActionConfigFile)
     train_parser = ArgumentParser(prog="")
     AD.add_class_args(train_parser, prefix="dataset", skip={})
-    Sampler.add_class_args(train_parser, prefix="sampler")
+    SegSamplerFactory.add_class_args(train_parser, prefix="sampler")
     train_parser.add_argument(
         "--data_loader.num-workers",
         type=int,
@@ -163,7 +150,7 @@ def make_parser(model_class):
 
     val_parser = ArgumentParser(prog="")
     AD.add_class_args(val_parser, prefix="dataset", skip={})
-    Sampler.add_class_args(val_parser, prefix="sampler")
+    SegSamplerFactory.add_class_args(val_parser, prefix="sampler")
     val_parser.add_argument(
         "--data_loader.num-workers",
         type=int,
@@ -175,14 +162,11 @@ def make_parser(model_class):
     data_parser.add_argument("--val", action=ActionParser(parser=val_parser))
     parser.add_argument("--data", action=ActionParser(parser=data_parser))
     parser.link_arguments(
-        "data.train.dataset.class_file", "data.val.dataset.class_file"
+        "data.train.dataset.class_files", "data.val.dataset.class_files"
     )
     parser.link_arguments(
         "data.train.data_loader.num_workers", "data.val.data_loader.num_workers"
     )
-    parser.link_arguments(
-        "data.train.sampler.batch_size", "data.val.sampler.batch_size"
-    )
 
     parser.add_argument("--in-model-file", required=True)
     model_class.add_finetune_args(parser, prefix="model")

hyperion/bin/finetune_xvector_from_wav.py

@@ -21,11 +21,9 @@
 from hyperion.hyp_defs import config_logger, set_float_cpu
 from hyperion.torch.utils import ddp
 
-# from hyperion.torch.models import XVector as XVec
 from hyperion.torch.trainers import XVectorTrainerFromWav as Trainer
 from hyperion.torch.data import AudioDataset as AD
 
-# from hyperion.torch.data import ClassWeightedSeqSampler as Sampler
 from hyperion.torch import TorchModelLoader as TML
 from hyperion.torch.data import SegSamplerFactory
 from hyperion.torch.metrics import CategoricalAccuracy

hyperion/bin/train_wav2vec2xvector.py

@@ -27,7 +27,6 @@
 from hyperion.torch.data import AudioDataset as AD
 from hyperion.torch.data import SegSamplerFactory
 
-# from hyperion.torch.data import ClassWeightedSeqSampler as Sampler
 from hyperion.torch.metrics import CategoricalAccuracy
 from hyperion.torch.models import (
     HFWav2Vec2ResNet1dXVector,
@@ -74,36 +73,6 @@ def init_data(partition, rank, num_gpus, **kwargs):
     return data_loader
 
 
-# def init_data(partition, rank, num_gpus, **kwargs):
-
-#     kwargs = kwargs["data"][partition]
-#     ad_args = AD.filter_args(**kwargs["dataset"])
-#     sampler_args = Sampler.filter_args(**kwargs["sampler"])
-#     if rank == 0:
-#         logging.info("{} audio dataset args={}".format(partition, ad_args))
-#         logging.info("{} sampler args={}".format(partition, sampler_args))
-#         logging.info("init %s dataset", partition)
-
-#     ad_args["is_val"] = partition == "val"
-#     dataset = AD(**ad_args)
-
-#     if rank == 0:
-#         logging.info("init %s samplers", partition)
-
-#     sampler = Sampler(dataset, **sampler_args)
-
-#     if rank == 0:
-#         logging.info("init %s dataloader", partition)
-
-#     num_workers = kwargs["data_loader"]["num_workers"]
-#     num_workers_per_gpu = int((num_workers + num_gpus - 1) / num_gpus)
-#     largs = (
-#         {"num_workers": num_workers_per_gpu, "pin_memory": True} if num_gpus > 0 else {}
-#     )
-#     data_loader = torch.utils.data.DataLoader(dataset, batch_sampler=sampler, **largs)
-#     return data_loader
-
-
 def init_model(num_classes, rank, model_class, **kwargs):
     model_args = model_class.filter_args(**kwargs["model"])
     if rank == 0:

hyperion/torch/data/audio_dataset.py

@@ -734,6 +734,14 @@ def __getitem__(self, segment):
         else:
             r = [x]
 
+        # try:
+        #     import soundfile as sf
+
+        #     for i, z in enumerate(r):
+        #         sf.write(f"file_{seg_id}.wav", z, fs, "PCM_16")
+        # except:
+        #     print("soundfile failed", flush=True)
+
         # adds the segment labels
         seg_info = self._get_segment_info(seg_id)
         r.extend(seg_info)

hyperion/torch/data/class_weighted_seg_chunk_sampler.py

@@ -235,7 +235,6 @@ def set_hard_prototypes(self, affinity_matrix):
             if np.all(mask_i == 0):
                 affinity_matrix[:, i] = -1000
 
-        # affinity_matrix[np.diag(affinity_matrix.shape[0])] = -1.0
         # hard prototypes for a class are itself and k-1 closest to it.
         self.hard_prototypes = torch.topk(
             affinity_matrix, self.num_hard_prototypes, dim=-1

hyperion/torch/models/wav2xvectors/hf_wav2xvector.py

@@ -83,6 +83,10 @@ def _fuse_hid_feats(self, hid_feats):
 
         return feats
 
+    @property
+    def sample_frequency(self):
+        return self.hf_feats.sample_frequency
+
     def compute_prototype_affinity(self):
         return self.xvector.compute_prototype_affinity()
 

hyperion/torch/models/xvectors/xvector.py

@@ -892,7 +892,7 @@ def add_finetune_args(parser, prefix=None):
         parser.add_argument(
             "--num-subcenters",
             default=2,
-            type=float,
+            type=int,
             help="number of subcenters in subcenter losses",
         )