Add models for voice conversion

voice_conversion/softvc/hubert/.gitignore

@@ -0,0 +1,132 @@
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+# C extensions
+*.so
+
+# Distribution / packaging
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+pip-wheel-metadata/
+share/python-wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+
+# PyInstaller
+#  Usually these files are written by a python script from a template
+#  before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.nox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+*.py,cover
+.hypothesis/
+.pytest_cache/
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+local_settings.py
+db.sqlite3
+db.sqlite3-journal
+
+# Flask stuff:
+instance/
+.webassets-cache
+
+# Scrapy stuff:
+.scrapy
+
+# Sphinx documentation
+docs/_build/
+
+# PyBuilder
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# IPython
+profile_default/
+ipython_config.py
+
+# pyenv
+.python-version
+
+# pipenv
+#   According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
+#   However, in case of collaboration, if having platform-specific dependencies or dependencies
+#   having no cross-platform support, pipenv may install dependencies that don't work, or not
+#   install all needed dependencies.
+#Pipfile.lock
+
+# PEP 582; used by e.g. github.com/David-OConnor/pyflow
+__pypackages__/
+
+# Celery stuff
+celerybeat-schedule
+celerybeat.pid
+
+# SageMath parsed files
+*.sage.py
+
+# Environments
+.env
+.venv
+env/
+venv/
+ENV/
+env.bak/
+venv.bak/
+
+# VSCode project settings
+.vscode
+
+# Spyder project settings
+.spyderproject
+.spyproject
+
+# Rope project settings
+.ropeproject
+
+# mkdocs documentation
+/site
+
+# mypy
+.mypy_cache/
+.dmypy.json
+dmypy.json
+
+# Pyre type checker
+.pyre/

voice_conversion/softvc/hubert/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2021 Benjamin van Niekerk
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

voice_conversion/softvc/hubert/README.md

@@ -0,0 +1,163 @@
+<p align="center">
+    <a target="_blank" href="https://colab.research.google.com/github/bshall/soft-vc/blob/main/soft-vc-demo.ipynb">
+        <img src="https://colab.research.google.com/assets/colab-badge.svg" alt="Open In Colab"/>
+    </a>
+</p>
+
+# HuBERT
+
+Training and inference scripts for the HuBERT content encoders in [A Comparison of Discrete and Soft Speech Units for Improved Voice Conversion](https://ieeexplore.ieee.org/abstract/document/9746484).
+For more details see [soft-vc](https://github.com/bshall/soft-vc). Audio samples can be found [here](https://bshall.github.io/soft-vc/). Colab demo can be found [here](https://colab.research.google.com/github/bshall/soft-vc/blob/main/soft-vc-demo.ipynb).
+
+<div align="center">
+    <img width="100%" alt="Soft-VC"
+      src="https://raw.githubusercontent.com/bshall/hubert/main/content-encoder.png">
+</div>
+<div>
+  <sup>
+    <strong>Fig 1:</strong> Architecture of the voice conversion system. a) The <strong>discrete</strong> content encoder clusters audio features to produce a sequence of discrete speech units. b) The <strong>soft</strong> content encoder is trained to predict the discrete units. The acoustic model transforms the discrete/soft speech units into a target spectrogram. The vocoder converts the spectrogram into an audio waveform.
+  </sup>
+</div>
+
+## Example Usage
+
+### Programmatic Usage
+
+```python
+import torch, torchaudio
+
+# Load checkpoint (either hubert_soft or hubert_discrete)
+hubert = torch.hub.load("bshall/hubert:main", "hubert_soft").cuda()
+
+# Load audio
+wav, sr = torchaudio.load("path/to/wav")
+assert sr == 16000
+wav = wav.unsqueeze(0).cuda()
+
+# Extract speech units
+units = hubert.units(x)
+```
+
+### Script-Based Usage
+
+```
+usage: encode.py [-h] [--extension EXTENSION] {soft,discrete} in-dir out-dir
+
+Encode an audio dataset.
+
+positional arguments:
+  {soft,discrete}       available models (HuBERT-Soft or HuBERT-Discrete)
+  in-dir                path to the dataset directory.
+  out-dir               path to the output directory.
+
+optional arguments:
+  -h, --help            show this help message and exit
+  --extension EXTENSION
+                        extension of the audio files (defaults to .flac).
+```
+
+## Training
+
+### Step 1: Dataset Preparation
+
+Download and extract the [LibriSpeech](https://www.openslr.org/12) corpus. The training script expects the following tree structure for the dataset directory:
+
+```
+│   lengths.json
+│
+└───wavs
+    ├───dev-*
+    │   ├───84
+    │   ├───...
+    │   └───8842
+    └───train-*
+        ├───19
+        ├───...
+        └───8975
+```
+
+The `train-*` and `dev-*` directories should contain the training and validation splits respectively. Note that there can be multiple `train` and `dev` folders e.g., `train-clean-100`, `train-other-500`, etc. Finally, the `lengths.json` file should contain key-value pairs with the file path and number of samples:
+
+```json
+{
+    "dev-clean/1272/128104/1272-128104-0000": 93680,
+    "dev-clean/1272/128104/1272-128104-0001": 77040,
+}
+```
+
+### Step 2: Extract Discrete Speech Units
+
+Encode LibriSpeech using the HuBERT-Discrete model and `encode.py` script:
+
+```
+usage: encode.py [-h] [--extension EXTENSION] {soft,discrete} in-dir out-dir
+
+Encode an audio dataset.
+
+positional arguments:
+  {soft,discrete}       available models (HuBERT-Soft or HuBERT-Discrete)
+  in-dir                path to the dataset directory.
+  out-dir               path to the output directory.
+
+optional arguments:
+  -h, --help            show this help message and exit
+  --extension EXTENSION
+                        extension of the audio files (defaults to .flac).
+```
+
+for example:
+
+```
+python encode.py discrete path/to/LibriSpeech/wavs path/to/LibriSpeech/units
+```
+
+At this point the directory tree should look like:
+
+```
+│   lengths.json
+│
+├───discrete
+│   ├───...
+└───wavs
+    ├───...
+```
+
+### Step 3: Train the HuBERT-Soft Content Encoder
+
+```
+usage: train.py [-h] [--resume RESUME] [--warmstart] [--mask] [--alpha ALPHA] dataset-dir checkpoint-dir
+
+Train HuBERT soft content encoder.
+
+positional arguments:
+  dataset-dir      path to the data directory.
+  checkpoint-dir   path to the checkpoint directory.
+
+optional arguments:
+  -h, --help       show this help message and exit
+  --resume RESUME  path to the checkpoint to resume from.
+  --warmstart      whether to initialize from the fairseq HuBERT checkpoint.
+  --mask           whether to use input masking.
+  --alpha ALPHA    weight for the masked loss.
+```
+
+## Links
+
+- [Soft-VC repo](https://github.com/bshall/soft-vc)
+- [Soft-VC paper](https://ieeexplore.ieee.org/abstract/document/9746484)
+- [Official HuBERT repo](https://github.com/pytorch/fairseq)
+- [HuBERT paper](https://arxiv.org/abs/2106.07447)
+
+## Citation
+
+If you found this work helpful please consider citing our paper:
+
+```
+@inproceedings{
+    soft-vc-2022,
+    author={van Niekerk, Benjamin and Carbonneau, Marc-André and Zaïdi, Julian and Baas, Matthew and Seuté, Hugo and Kamper, Herman},
+    booktitle={ICASSP}, 
+    title={A Comparison of Discrete and Soft Speech Units for Improved Voice Conversion}, 
+    year={2022}
+}
+```

voice_conversion/softvc/hubert/cluster.py

@@ -0,0 +1,66 @@
+from pathlib import Path
+import logging
+import argparse
+
+import torch
+import numpy as np
+from sklearn.cluster import KMeans
+
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+
+def cluster(args):
+    with open(args.subset) as file:
+        subset = [line.strip() for line in file]
+
+    logger.info(f"Loading features from {args.in_dir}")
+    features = []
+    for path in subset:
+        in_path = args.in_dir / path
+        features.append(np.load(in_path.with_suffix(".npy")))
+    features = np.concatenate(features, axis=0)
+
+    logger.info(f"Clustering features of shape: {features.shape}")
+    kmeans = KMeans(n_clusters=args.n_clusters).fit(features)
+
+    checkpoint_path = args.checkpoint_dir / f"kmeans_{args.n_clusters}.pt"
+    checkpoint_path.parent.mkdir(exist_ok=True, parents=True)
+    torch.save(
+        checkpoint_path,
+        {
+            "n_features_in_": kmeans.n_features_in_,
+            "_n_threads": kmeans._n_threads,
+            "cluster_centers_": kmeans.cluster_centers_,
+        },
+    )
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="Cluster speech features features.")
+    parser.add_argument(
+        "in_dir",
+        metavar="in-dir",
+        help="path to the encoded dataset",
+        type=Path,
+    )
+    parser.add_argument(
+        "subset",
+        matavar="subset",
+        help="path to the .txt file containing the list of files to cluster",
+        type=Path,
+    )
+    parser.add_argument(
+        "checkpoint_dir",
+        metavar="checkpoint-dir",
+        help="path to the checkpoint directory",
+        type=Path,
+    )
+    parser.add_argument(
+        "--n-clusters",
+        help="number of clusters",
+        type=int,
+        default=100,
+    )
+    args = parser.parse_args()
+    cluster(args)

voice_conversion/softvc/hubert/encode.py

@@ -0,0 +1,56 @@
+import argparse
+import logging
+import numpy as np
+from pathlib import Path
+from tqdm import tqdm
+
+import torch
+import torchaudio
+from torchaudio.functional import resample
+
+
+def encode_dataset(args):
+    print(f"Loading hubert checkpoint")
+    hubert = torch.hub.load("bshall/hubert:main", f"hubert_{args.model}").cuda()
+
+    print(f"Encoding dataset at {args.in_dir}")
+    for in_path in tqdm(list(args.in_dir.rglob(f"*{args.extension}"))):
+        wav, sr = torchaudio.load(in_path)
+        wav = resample(wav, sr, 16000)
+        wav = wav.unsqueeze(0).cuda()
+
+        with torch.inference_mode():
+            units = hubert.units(wav)
+
+        out_path = args.out_dir / in_path.relative_to(args.in_dir)
+        out_path.parent.mkdir(parents=True, exist_ok=True)
+        np.save(out_path.with_suffix(".npy"), units.squeeze().cpu().numpy())
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="Encode an audio dataset.")
+    parser.add_argument(
+        "model",
+        help="available models (HuBERT-Soft or HuBERT-Discrete)",
+        choices=["soft", "discrete"],
+    )
+    parser.add_argument(
+        "in_dir",
+        metavar="in-dir",
+        help="path to the dataset directory.",
+        type=Path,
+    )
+    parser.add_argument(
+        "out_dir",
+        metavar="out-dir",
+        help="path to the output directory.",
+        type=Path,
+    )
+    parser.add_argument(
+        "--extension",
+        help="extension of the audio files (defaults to .flac).",
+        default=".flac",
+        type=str,
+    )
+    args = parser.parse_args()
+    encode_dataset(args)

voice_conversion/softvc/hubert/hubconf.py

@@ -0,0 +1,3 @@
+dependencies = ["torch", "torchaudio", "sklearn"]
+
+from hubert import hubert_discrete, hubert_soft, kmeans100

voice_conversion/softvc/hubert/hubert/__init__.py

@@ -0,0 +1,8 @@
+from .model import (
+    Hubert,
+    HubertDiscrete,
+    HubertSoft,
+    hubert_discrete,
+    hubert_soft,
+    kmeans100,
+)

voice_conversion/softvc/hubert/hubert/dataset.py

@@ -0,0 +1,91 @@
+import random
+from pathlib import Path
+import numpy as np
+import json
+
+import torch
+import torch.nn.functional as F
+from torch.utils.data import Dataset
+import torchaudio
+
+
+class AcousticUnitsDataset(Dataset):
+    def __init__(
+        self,
+        root: Path,
+        sample_rate: int = 16000,
+        label_rate: int = 50,
+        min_samples: int = 32000,
+        max_samples: int = 250000,
+        train: bool = True,
+    ):
+        self.wavs_dir = root / "wavs"
+        self.units_dir = root / "units"
+
+        with open(root / "lengths.json") as file:
+            self.lenghts = json.load(file)
+
+        pattern = "train-*/**/*.flac" if train else "dev-*/**/*.flac"
+        metadata = (
+            (path, path.relative_to(self.wavs_dir).with_suffix("").as_posix())
+            for path in self.wavs_dir.rglob(pattern)
+        )
+        metadata = ((path, key) for path, key in metadata if key in self.lenghts)
+        self.metadata = [
+            path for path, key in metadata if self.lenghts[key] > min_samples
+        ]
+
+        self.sample_rate = sample_rate
+        self.label_rate = label_rate
+        self.min_samples = min_samples
+        self.max_samples = max_samples
+        self.train = train
+
+    def __len__(self):
+        return len(self.metadata)
+
+    def __getitem__(self, index):
+        wav_path = self.metadata[index]
+        units_path = self.units_dir / wav_path.relative_to(self.wavs_dir)
+
+        wav, _ = torchaudio.load(wav_path)
+        wav = F.pad(wav, ((400 - 320) // 2, (400 - 320) // 2))
+        codes = np.load(units_path.with_suffix(".npy"))
+
+        return wav, torch.from_numpy(codes).long()
+
+    def collate(self, batch):
+        wavs, codes = zip(*batch)
+        wavs, codes = list(wavs), list(codes)
+
+        wav_lengths = [wav.size(-1) for wav in wavs]
+        code_lengths = [code.size(-1) for code in codes]
+
+        wav_frames = min(self.max_samples, *wav_lengths)
+
+        collated_wavs, wav_offsets = [], []
+        for wav in wavs:
+            wav_diff = wav.size(-1) - wav_frames
+            wav_offset = random.randint(0, wav_diff)
+            wav = wav[:, wav_offset : wav_offset + wav_frames]
+
+            collated_wavs.append(wav)
+            wav_offsets.append(wav_offset)
+
+        rate = self.label_rate / self.sample_rate
+        code_offsets = [round(wav_offset * rate) for wav_offset in wav_offsets]
+        code_frames = round(wav_frames * rate)
+        remaining_code_frames = [
+            length - offset for length, offset in zip(code_lengths, code_offsets)
+        ]
+        code_frames = min(code_frames, *remaining_code_frames)
+
+        collated_codes = []
+        for code, code_offset in zip(codes, code_offsets):
+            code = code[code_offset : code_offset + code_frames]
+            collated_codes.append(code)
+
+        wavs = torch.stack(collated_wavs, dim=0)
+        codes = torch.stack(collated_codes, dim=0)
+
+        return wavs, codes

voice_conversion/softvc/hubert/hubert/model.py

@@ -0,0 +1,293 @@
+import copy
+from typing import Optional, Tuple
+import random
+
+from sklearn.cluster import KMeans
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.nn.modules.utils import consume_prefix_in_state_dict_if_present
+
+URLS = {
+    "hubert-discrete": "https://github.com/bshall/hubert/releases/download/v0.1/hubert-discrete-e9416457.pt",
+    "hubert-soft": "https://github.com/bshall/hubert/releases/download/v0.1/hubert-soft-0d54a1f4.pt",
+    "kmeans100": "https://github.com/bshall/hubert/releases/download/v0.1/kmeans100-50f36a95.pt",
+}
+
+
+class Hubert(nn.Module):
+    def __init__(self, num_label_embeddings: int = 100, mask: bool = True):
+        super().__init__()
+        self._mask = mask
+        self.feature_extractor = FeatureExtractor()
+        self.feature_projection = FeatureProjection()
+        self.positional_embedding = PositionalConvEmbedding()
+        self.norm = nn.LayerNorm(768)
+        self.dropout = nn.Dropout(0.1)
+        self.encoder = TransformerEncoder(
+            nn.TransformerEncoderLayer(
+                768, 12, 3072, activation="gelu", batch_first=True
+            ),
+            12,
+        )
+        self.proj = nn.Linear(768, 256)
+
+        self.masked_spec_embed = nn.Parameter(torch.FloatTensor(768).uniform_())
+        self.label_embedding = nn.Embedding(num_label_embeddings, 256)
+
+    def mask(self, x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
+        mask = None
+        if self.training and self._mask:
+            mask = _compute_mask((x.size(0), x.size(1)), 0.8, 10, x.device, 2)
+            x[mask] = self.masked_spec_embed.to(x.dtype)
+        return x, mask
+
+    def encode(
+        self, x: torch.Tensor, layer: Optional[int] = None
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        x = self.feature_extractor(x)
+        x = self.feature_projection(x.transpose(1, 2))
+        x, mask = self.mask(x)
+        x = x + self.positional_embedding(x)
+        x = self.dropout(self.norm(x))
+        x = self.encoder(x, output_layer=layer)
+        return x, mask
+
+    def logits(self, x: torch.Tensor) -> torch.Tensor:
+        logits = torch.cosine_similarity(
+            x.unsqueeze(2),
+            self.label_embedding.weight.unsqueeze(0).unsqueeze(0),
+            dim=-1,
+        )
+        return logits / 0.1
+
+    def forward(self, x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
+        x, mask = self.encode(x)
+        x = self.proj(x)
+        logits = self.logits(x)
+        return logits, mask
+
+
+class HubertSoft(Hubert):
+    def __init__(self):
+        super().__init__()
+
+    @torch.inference_mode()
+    def units(self, wav: torch.Tensor) -> torch.Tensor:
+        wav = F.pad(wav, ((400 - 320) // 2, (400 - 320) // 2))
+        x, _ = self.encode(wav)
+        return self.proj(x)
+
+
+class HubertDiscrete(Hubert):
+    def __init__(self, kmeans):
+        super().__init__(504)
+        self.kmeans = kmeans
+
+    @torch.inference_mode()
+    def units(self, wav: torch.Tensor) -> torch.LongTensor:
+        wav = F.pad(wav, ((400 - 320) // 2, (400 - 320) // 2))
+        x, _ = self.encode(wav, layer=7)
+        x = self.kmeans.predict(x.squeeze().cpu().numpy())
+        return torch.tensor(x, dtype=torch.long, device=wav.device)
+
+
+class FeatureExtractor(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.conv0 = nn.Conv1d(1, 512, 10, 5, bias=False)
+        self.norm0 = nn.GroupNorm(512, 512)
+        self.conv1 = nn.Conv1d(512, 512, 3, 2, bias=False)
+        self.conv2 = nn.Conv1d(512, 512, 3, 2, bias=False)
+        self.conv3 = nn.Conv1d(512, 512, 3, 2, bias=False)
+        self.conv4 = nn.Conv1d(512, 512, 3, 2, bias=False)
+        self.conv5 = nn.Conv1d(512, 512, 2, 2, bias=False)
+        self.conv6 = nn.Conv1d(512, 512, 2, 2, bias=False)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = F.gelu(self.norm0(self.conv0(x)))
+        x = F.gelu(self.conv1(x))
+        x = F.gelu(self.conv2(x))
+        x = F.gelu(self.conv3(x))
+        x = F.gelu(self.conv4(x))
+        x = F.gelu(self.conv5(x))
+        x = F.gelu(self.conv6(x))
+        return x
+
+
+class FeatureProjection(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.norm = nn.LayerNorm(512)
+        self.projection = nn.Linear(512, 768)
+        self.dropout = nn.Dropout(0.1)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = self.norm(x)
+        x = self.projection(x)
+        x = self.dropout(x)
+        return x
+
+
+class PositionalConvEmbedding(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.conv = nn.Conv1d(
+            768,
+            768,
+            kernel_size=128,
+            padding=128 // 2,
+            groups=16,
+        )
+        self.conv = nn.utils.weight_norm(self.conv, name="weight", dim=2)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = self.conv(x.transpose(1, 2))
+        x = F.gelu(x[:, :, :-1])
+        return x.transpose(1, 2)
+
+
+class TransformerEncoder(nn.Module):
+    def __init__(
+        self, encoder_layer: nn.TransformerEncoderLayer, num_layers: int
+    ) -> None:
+        super(TransformerEncoder, self).__init__()
+        self.layers = nn.ModuleList(
+            [copy.deepcopy(encoder_layer) for _ in range(num_layers)]
+        )
+        self.num_layers = num_layers
+
+    def forward(
+        self,
+        src: torch.Tensor,
+        mask: torch.Tensor = None,
+        src_key_padding_mask: torch.Tensor = None,
+        output_layer: Optional[int] = None,
+    ) -> torch.Tensor:
+        output = src
+        for layer in self.layers[:output_layer]:
+            output = layer(
+                output, src_mask=mask, src_key_padding_mask=src_key_padding_mask
+            )
+        return output
+
+
+def _compute_mask(
+    shape: Tuple[int, int],
+    mask_prob: float,
+    mask_length: int,
+    device: torch.device,
+    min_masks: int = 0,
+) -> torch.Tensor:
+    batch_size, sequence_length = shape
+
+    if mask_length < 1:
+        raise ValueError("`mask_length` has to be bigger than 0.")
+
+    if mask_length > sequence_length:
+        raise ValueError(
+            f"`mask_length` has to be smaller than `sequence_length`, but got `mask_length`: {mask_length} and `sequence_length`: {sequence_length}`"
+        )
+
+    # compute number of masked spans in batch
+    num_masked_spans = int(mask_prob * sequence_length / mask_length + random.random())
+    num_masked_spans = max(num_masked_spans, min_masks)
+
+    # make sure num masked indices <= sequence_length
+    if num_masked_spans * mask_length > sequence_length:
+        num_masked_spans = sequence_length // mask_length
+
+    # SpecAugment mask to fill
+    mask = torch.zeros((batch_size, sequence_length), device=device, dtype=torch.bool)
+
+    # uniform distribution to sample from, make sure that offset samples are < sequence_length
+    uniform_dist = torch.ones(
+        (batch_size, sequence_length - (mask_length - 1)), device=device
+    )
+
+    # get random indices to mask
+    mask_indices = torch.multinomial(uniform_dist, num_masked_spans)
+
+    # expand masked indices to masked spans
+    mask_indices = (
+        mask_indices.unsqueeze(dim=-1)
+        .expand((batch_size, num_masked_spans, mask_length))
+        .reshape(batch_size, num_masked_spans * mask_length)
+    )
+    offsets = (
+        torch.arange(mask_length, device=device)[None, None, :]
+        .expand((batch_size, num_masked_spans, mask_length))
+        .reshape(batch_size, num_masked_spans * mask_length)
+    )
+    mask_idxs = mask_indices + offsets
+
+    # scatter indices to mask
+    mask = mask.scatter(1, mask_idxs, True)
+
+    return mask
+
+
+def hubert_discrete(
+    pretrained: bool = True,
+    progress: bool = True,
+) -> HubertDiscrete:
+    r"""HuBERT-Discrete from `"A Comparison of Discrete and Soft Speech Units for Improved Voice Conversion"`.
+    Args:
+        pretrained (bool): load pretrained weights into the model
+        progress (bool): show progress bar when downloading model
+    """
+    kmeans = kmeans100(pretrained=pretrained, progress=progress)
+    hubert = HubertDiscrete(kmeans)
+    if pretrained:
+        checkpoint = torch.hub.load_state_dict_from_url(
+            URLS["hubert-discrete"], progress=progress
+        )
+        consume_prefix_in_state_dict_if_present(checkpoint, "module.")
+        hubert.load_state_dict(checkpoint)
+        hubert.eval()
+    return hubert
+
+
+def hubert_soft(
+    pretrained: bool = True,
+    progress: bool = True,
+) -> HubertSoft:
+    r"""HuBERT-Soft from `"A Comparison of Discrete and Soft Speech Units for Improved Voice Conversion"`.
+    Args:
+        pretrained (bool): load pretrained weights into the model
+        progress (bool): show progress bar when downloading model
+    """
+    hubert = HubertSoft()
+    if pretrained:
+        checkpoint = torch.hub.load_state_dict_from_url(
+            URLS["hubert-soft"], progress=progress
+        )
+        consume_prefix_in_state_dict_if_present(checkpoint, "module.")
+        hubert.load_state_dict(checkpoint)
+        hubert.eval()
+    return hubert
+
+
+def _kmeans(
+    num_clusters: int, pretrained: bool = True, progress: bool = True
+) -> KMeans:
+    kmeans = KMeans(num_clusters)
+    if pretrained:
+        checkpoint = torch.hub.load_state_dict_from_url(
+            URLS[f"kmeans{num_clusters}"], progress=progress
+        )
+        kmeans.__dict__["n_features_in_"] = checkpoint["n_features_in_"]
+        kmeans.__dict__["_n_threads"] = checkpoint["_n_threads"]
+        kmeans.__dict__["cluster_centers_"] = checkpoint["cluster_centers_"].numpy()
+    return kmeans
+
+
+def kmeans100(pretrained: bool = True, progress: bool = True) -> KMeans:
+    r"""
+    k-means checkpoint for HuBERT-Discrete with 100 clusters.
+    Args:
+        pretrained (bool): load pretrained weights into the model
+        progress (bool): show progress bar when downloading model
+    """
+    return _kmeans(100, pretrained, progress)

voice_conversion/softvc/hubert/hubert/utils.py

@@ -0,0 +1,58 @@
+import torch
+
+
+class Metric:
+    def __init__(self):
+        self.steps = 0
+        self.value = 0
+
+    def update(self, value):
+        self.steps += 1
+        self.value += (value - self.value) / self.steps
+        return self.value
+
+    def reset(self):
+        self.steps = 0
+        self.value = 0
+
+
+def save_checkpoint(
+    checkpoint_dir,
+    hubert,
+    optimizer,
+    scaler,
+    step,
+    loss,
+    best,
+    logger,
+):
+    state = {
+        "hubert": hubert.state_dict(),
+        "optimizer": optimizer.state_dict(),
+        "scaler": scaler.state_dict(),
+        "step": step,
+        "loss": loss,
+    }
+    checkpoint_dir.mkdir(exist_ok=True, parents=True)
+    checkpoint_path = checkpoint_dir / f"model-{step}.pt"
+    torch.save(state, checkpoint_path)
+    if best:
+        best_path = checkpoint_dir / "model-best.pt"
+        torch.save(state, best_path)
+    logger.info(f"Saved checkpoint: {checkpoint_path.stem}")
+
+
+def load_checkpoint(
+    load_path,
+    hubert,
+    optimizer,
+    scaler,
+    rank,
+    logger,
+):
+    logger.info(f"Loading checkpoint from {load_path}")
+    checkpoint = torch.load(load_path, map_location={"cuda:0": f"cuda:{rank}"})
+    hubert.load_state_dict(checkpoint["hubert"])
+    scaler.load_state_dict(checkpoint["scaler"])
+    optimizer.load_state_dict(checkpoint["optimizer"])
+    return checkpoint["step"], checkpoint["loss"]

voice_conversion/softvc/hubert/train.py

@@ -0,0 +1,453 @@
+import argparse
+import logging
+from pathlib import Path
+
+import torch
+import torch.cuda.amp as amp
+import torch.nn as nn
+import torch.nn.functional as F
+import torch.optim as optim
+from torch.utils.data import DataLoader
+from torch.utils.tensorboard import SummaryWriter
+import torch.distributed as dist
+from torch.utils.data.distributed import DistributedSampler
+import torch.multiprocessing as mp
+from torch.nn.parallel import DistributedDataParallel as DDP
+from torch.nn.modules.utils import consume_prefix_in_state_dict_if_present
+
+from hubert.model import Hubert, URLS
+from hubert.dataset import AcousticUnitsDataset
+from hubert.utils import Metric, save_checkpoint, load_checkpoint
+
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+########################################################################################
+# Define hyperparameters for training:
+########################################################################################
+
+BATCH_SIZE = 64
+LEARNING_RATE = 2e-5
+BETAS = (0.9, 0.98)
+EPS = 1e-06
+WEIGHT_DECAY = 1e-2
+MAX_NORM = 10
+STEPS = 25000
+LOG_INTERVAL = 5
+VALIDATION_INTERVAL = 1000
+CHECKPOINT_INTERVAL = 1000
+BACKEND = "nccl"
+INIT_METHOD = "tcp://localhost:54321"
+
+
+def train(rank, world_size, args):
+    dist.init_process_group(
+        BACKEND,
+        rank=rank,
+        world_size=world_size,
+        init_method=INIT_METHOD,
+    )
+
+    ####################################################################################
+    # Setup logging utilities:
+    ####################################################################################
+
+    log_dir = args.checkpoint_dir / "logs"
+    log_dir.mkdir(exist_ok=True, parents=True)
+
+    if rank == 0:
+        logger.setLevel(logging.INFO)
+        handler = logging.FileHandler(log_dir / f"{args.checkpoint_dir.stem}.log")
+        handler.setLevel(logging.INFO)
+        formatter = logging.Formatter(
+            "%(asctime)s [%(levelname)s] %(message)s", datefmt="%m/%d/%Y %I:%M:%S"
+        )
+        handler.setFormatter(formatter)
+        logger.addHandler(handler)
+    else:
+        logger.setLevel(logging.ERROR)
+
+    writer = SummaryWriter(log_dir) if rank == 0 else None
+
+    ####################################################################################
+    # Initialize models
+    ####################################################################################
+
+    hubert = Hubert(mask=args.mask).to(rank)
+
+    if args.warmstart:
+        checkpoint = torch.hub.load_state_dict_from_url(
+            URLS["hubert-discrete"], map_location={"cuda:0": f"cuda:{rank}"}
+        )
+        consume_prefix_in_state_dict_if_present(checkpoint, "module.")
+        hubert.load_state_dict(checkpoint, strict=False)
+
+    hubert = DDP(hubert, device_ids=[rank])
+
+    ####################################################################################
+    # Initialze optimizer and grad scaler
+    ####################################################################################
+
+    optimizer = optim.AdamW(
+        hubert.parameters(),
+        lr=LEARNING_RATE,
+        betas=BETAS,
+        eps=EPS,
+        weight_decay=WEIGHT_DECAY,
+    )
+    scaler = amp.GradScaler()
+
+    ####################################################################################
+    # Initialize datasets and dataloaders
+    ####################################################################################
+
+    train_dataset = AcousticUnitsDataset(
+        root=args.dataset_dir,
+        train=True,
+    )
+    train_sampler = DistributedSampler(train_dataset, drop_last=True)
+    train_loader = DataLoader(
+        train_dataset,
+        collate_fn=train_dataset.collate,
+        batch_size=BATCH_SIZE,
+        sampler=train_sampler,
+        num_workers=8,
+        pin_memory=True,
+        shuffle=False,
+        drop_last=True,
+    )
+
+    validation_dataset = AcousticUnitsDataset(
+        root=args.dataset_dir,
+        train=False,
+    )
+    validation_loader = DataLoader(
+        validation_dataset,
+        batch_size=1,
+        shuffle=False,
+        num_workers=8,
+        pin_memory=True,
+    )
+
+    ####################################################################################
+    # Load checkpoint if args.resume is set
+    ####################################################################################
+
+    if args.resume is not None:
+        global_step, best_loss = load_checkpoint(
+            load_path=args.resume,
+            hubert=hubert,
+            optimizer=optimizer,
+            scaler=scaler,
+            rank=rank,
+            logger=logger,
+        )
+    else:
+        global_step, best_loss = 0, float("inf")
+
+    # =================================================================================#
+    # Start training loop
+    # =================================================================================#
+
+    n_epochs = STEPS // len(train_loader) + 1
+    start_epoch = global_step // len(train_loader) + 1
+
+    logger.info("**" * 40)
+    logger.info(f"PyTorch version: {torch.__version__}")
+    logger.info(f"CUDA version: {torch.version.cuda}")
+    logger.info(f"CUDNN version: {torch.backends.cudnn.version()}")
+    logger.info(f"CUDNN enabled: {torch.backends.cudnn.enabled}")
+    logger.info(f"CUDNN deterministic: {torch.backends.cudnn.deterministic}")
+    logger.info(f"CUDNN benchmark: {torch.backends.cudnn.benchmark}")
+    logger.info(f"# of GPUS: {torch.cuda.device_count()}")
+    logger.info(f"batch size: {BATCH_SIZE}")
+    logger.info(f"iterations per epoch: {len(train_loader)}")
+    logger.info(f"# of epochs: {n_epochs}")
+    logger.info(f"started at epoch: {start_epoch}")
+    logger.info("**" * 40 + "\n")
+
+    if args.mask:
+        average_masked_loss = Metric()
+        average_unmasked_loss = Metric()
+        average_masked_accuracy = Metric()
+        average_unmasked_accuracy = Metric()
+
+        epoch_masked_loss = Metric()
+        epoch_unmasked_loss = Metric()
+        epoch_masked_accuracy = Metric()
+        epoch_unmasked_accuracy = Metric()
+    else:
+        average_loss = Metric()
+        average_accuracy = Metric()
+
+        epoch_loss = Metric()
+        epoch_accuracy = Metric()
+
+    validation_loss = Metric()
+    validation_accuracy = Metric()
+
+    for epoch in range(start_epoch, n_epochs + 1):
+        train_sampler.set_epoch(epoch)
+
+        hubert.train()
+        if args.mask:
+            epoch_masked_loss.reset()
+            epoch_unmasked_loss.reset()
+            epoch_masked_accuracy.reset()
+            epoch_unmasked_accuracy.reset()
+        else:
+            epoch_loss.reset()
+            epoch_accuracy.reset()
+
+        for wavs, codes in train_loader:
+            global_step += 1
+            wavs, codes = wavs.to(rank), codes.to(rank)
+
+            ############################################################################
+            # Compute training loss
+            ############################################################################
+
+            optimizer.zero_grad()
+
+            with amp.autocast():
+                logits, mask = hubert(wavs)
+                length = min(
+                    mask.size(-1) if args.mask else float("inf"), codes.size(-1)
+                )
+                logits = logits[:, :length, :]
+                codes = codes[:, :length]
+                if args.mask:
+                    mask = mask[:, :length]
+
+                if args.mask:
+                    masked_loss = F.cross_entropy(logits[mask], codes[mask])
+                    unmasked_loss = F.cross_entropy(logits[~mask], codes[~mask])
+                    loss = args.alpha * masked_loss + (1 - args.alpha) * unmasked_loss
+                else:
+                    loss = F.cross_entropy(logits.transpose(1, 2), codes)
+
+            scaler.scale(loss).backward()
+            scaler.unscale_(optimizer)
+
+            nn.utils.clip_grad_norm_(hubert.parameters(), MAX_NORM)
+
+            scaler.step(optimizer)
+            scaler.update()
+
+            if args.mask:
+                masked_accuracy = logits[mask].argmax(dim=-1) == codes[mask]
+                masked_accuracy = torch.mean(masked_accuracy.float())
+
+                unmasked_accuracy = logits[~mask].argmax(dim=-1) == codes[~mask]
+                unmasked_accuracy = torch.mean(unmasked_accuracy.float())
+            else:
+                accuracy = logits.argmax(dim=-1) == codes
+                accuracy = torch.mean(accuracy.float())
+
+            ############################################################################
+            # Update and log training metrics
+            ############################################################################
+
+            if args.mask:
+                average_masked_loss.update(masked_loss.item())
+                average_unmasked_loss.update(unmasked_loss.item())
+                average_masked_accuracy.update(masked_accuracy.item())
+                average_unmasked_accuracy.update(unmasked_accuracy.item())
+
+                epoch_masked_loss.update(masked_loss.item())
+                epoch_unmasked_loss.update(unmasked_loss.item())
+                epoch_masked_accuracy.update(masked_accuracy.item())
+                epoch_unmasked_accuracy.update(unmasked_accuracy.item())
+            else:
+                average_loss.update(loss.item())
+                average_accuracy.update(accuracy.item())
+
+                epoch_loss.update(loss.item())
+                epoch_accuracy.update(accuracy.item())
+
+            if rank == 0 and global_step % LOG_INTERVAL == 0:
+                if args.mask:
+                    writer.add_scalar(
+                        "train/masked_loss",
+                        average_masked_loss.value,
+                        global_step,
+                    )
+                    writer.add_scalar(
+                        "train/unmasked_loss",
+                        average_unmasked_loss.value,
+                        global_step,
+                    )
+                    writer.add_scalar(
+                        "train/masked_accuracy",
+                        average_masked_accuracy.value * 100,
+                        global_step,
+                    )
+                    writer.add_scalar(
+                        "train/unmasked_accuracy",
+                        average_unmasked_accuracy.value * 100,
+                        global_step,
+                    )
+                    average_masked_loss.reset()
+                    average_unmasked_loss.reset()
+                    average_masked_accuracy.reset()
+                    average_unmasked_accuracy.reset()
+                else:
+                    writer.add_scalar(
+                        "train/loss",
+                        average_loss.value,
+                        global_step,
+                    )
+                    writer.add_scalar(
+                        "train/accuracy",
+                        average_accuracy.value,
+                        global_step,
+                    )
+                    average_loss.reset()
+                    average_accuracy.reset()
+
+            # --------------------------------------------------------------------------#
+            # Start validation loop
+            # --------------------------------------------------------------------------#
+
+            if global_step % VALIDATION_INTERVAL == 0:
+                hubert.eval()
+                validation_loss.reset()
+                validation_accuracy.reset()
+                for wavs, codes in validation_loader:
+                    wavs, codes = wavs.to(rank), codes.to(rank)
+
+                    with torch.no_grad():
+                        logits, _ = hubert(wavs)
+                        logits = logits.transpose(1, 2)
+
+                    loss = F.cross_entropy(logits, codes)
+
+                    accuracy = logits.argmax(dim=1) == codes
+                    accuracy = torch.mean(accuracy.float())
+
+                    ####################################################################
+                    # Update validation metrics
+                    ####################################################################
+
+                    validation_loss.update(loss.item())
+                    validation_accuracy.update(accuracy.item())
+
+                hubert.train()
+
+                ############################################################################
+                # Log validation metrics
+                ############################################################################
+
+                if rank == 0:
+                    writer.add_scalar(
+                        "validation/unit_loss",
+                        validation_loss.value,
+                        global_step,
+                    )
+                    writer.add_scalar(
+                        "validation/unit_accuracy",
+                        validation_accuracy.value * 100,
+                        global_step,
+                    )
+                    logger.info(
+                        f"valid -- epoch: {epoch}, loss: {validation_loss.value:.4f}, accuracy: {validation_accuracy.value * 100:.2f}"
+                    )
+
+                ############################################################################
+                # Save model checkpoint
+                ############################################################################
+
+                new_best = best_loss > validation_loss.value
+                if new_best or global_step % CHECKPOINT_INTERVAL == 0:
+                    if new_best:
+                        logger.info("-------- new best model found!")
+                        best_loss = validation_loss.value
+
+                    if rank == 0:
+                        save_checkpoint(
+                            checkpoint_dir=args.checkpoint_dir,
+                            hubert=hubert,
+                            optimizer=optimizer,
+                            scaler=scaler,
+                            step=global_step,
+                            loss=validation_loss.value,
+                            best=new_best,
+                            logger=logger,
+                        )
+
+            # -----------------------------------------------------------------------------#
+            # End validation loop
+            # -----------------------------------------------------------------------------#
+
+        ####################################################################################
+        # Log training metrics
+        ####################################################################################
+
+        logger.info(
+            f"""
+            train -- epoch: {epoch}, masked loss: {epoch_masked_loss.value:.4f}, unmasked loss: {epoch_unmasked_loss.value:.4f}, 
+                     masked accuracy: {epoch_masked_accuracy.value * 100:.2f}, umasked accuracy: {epoch_unmasked_accuracy.value * 100:.2f}
+            """
+        )
+
+        # ==================================================================================#
+        # End training loop
+        # ==================================================================================#
+
+    dist.destroy_process_group()
+
+
+def train_hubert(args):
+    world_size = torch.cuda.device_count()
+    mp.spawn(
+        train,
+        args=(world_size, args),
+        nprocs=world_size,
+        join=True,
+    )
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="Train HuBERT soft content encoder.")
+    parser.add_argument(
+        "dataset_dir",
+        metavar="dataset-dir",
+        help="path to the data directory.",
+        type=Path,
+    )
+    parser.add_argument(
+        "checkpoint_dir",
+        metavar="checkpoint-dir",
+        help="path to the checkpoint directory.",
+        type=Path,
+    )
+    parser.add_argument(
+        "--resume",
+        help="path to the checkpoint to resume from.",
+        type=Path,
+    )
+    parser.add_argument(
+        "--warmstart",
+        help="whether to initialize from the fairseq HuBERT checkpoint.",
+        action="store_true",
+    )
+    parser.add_argument(
+        "--mask",
+        help="whether to use input masking.",
+        action="store_true",
+    )
+    parser.add_argument(
+        "--alpha",
+        help="weight for the masked loss.",
+        default=1,
+        type=float,
+    )
+    args = parser.parse_args()
+
+    world_size = torch.cuda.device_count()
+    mp.spawn(
+        train,
+        args=(world_size, args),
+        nprocs=world_size,
+        join=True,
+    )