Implementation of BEST-RQ - a model for self-supervised learning of speech signals using a random projection quantizer, in Pytorch.
from best_rq_pytorch.best_rq import (
    BestRQ,
    BestRQPretrainWrapper,
    BestRQTrainASRWrapper
)

from best_rq_pytorch.conformer import (
    ConformerWrapper
)

from best_rq_pytorch.data import (
    AudioDataset,
    AudioGraphemeDataset
)

from best_rq_pytorch.trainer import (
    BestRQPretrainer,
    BestRQTrainerASR
)

best_rq_pytorch/__init__.py
lucasnewman/best-rq-pytorch
from collections import namedtuple
from functools import wraps
from packaging import version

import torch
from torch import nn, einsum
import torch.nn.functional as F

from einops import rearrange

# constants

EfficientAttentionConfig = namedtuple('EfficientAttentionConfig', ['enable_flash', 'enable_math', 'enable_mem_efficient'])

# helpers

def exists(val):
    return val is not None

def once(fn):
    called = False
    
    @wraps(fn)
    def inner(x):
        nonlocal called
        if called:
            return
        called = True
        return fn(x)
    return inner


print_once = once(print)

# main class

class Attend(nn.Module):
    def __init__(
        self,
        causal = False,
        dropout = 0.,
        flash = False
    ):
        super().__init__()
        self.dropout = dropout
        self.attn_dropout = nn.Dropout(dropout)

        self.causal = causal
        self.flash = flash
        assert not (flash and version.parse(torch.__version__) < version.parse('2.0.0')), 'in order to use flash attention, you must be using pytorch 2.0 or above'

        # determine efficient attention configs for cuda and cpu

        self.cpu_config = EfficientAttentionConfig(True, True, True)
        self.cuda_config = None

        if not torch.cuda.is_available() or not flash:
            return

        device_properties = torch.cuda.get_device_properties(torch.device('cuda'))

        if device_properties.major == 8 and device_properties.minor == 0:
            print_once('A100 GPU detected, using flash attention if input tensor is on cuda')
            self.cuda_config = EfficientAttentionConfig(True, False, False)
        else:
            print_once('Non-A100 GPU detected, using math or mem efficient attention if input tensor is on cuda')
            self.cuda_config = EfficientAttentionConfig(False, True, True)

    def get_mask(self, i, j, device):
        return torch.ones((i, j), device=device, dtype=torch.bool).triu(j - i + 1)

    def flash_attn(self, q, k, v, mask = None, attn_bias = None):
        _, heads, q_len, _, k_len, is_cuda, device = *q.shape, k.shape[-2], q.is_cuda, q.device

        # single headed key / values

        if k.ndim == 3:
            k = rearrange(k, 'b n d -> b 1 n d')

        if v.ndim == 3:
            v = rearrange(v, 'b n d -> b 1 n d')

        # Check if mask exists and expand to compatible shape
        # The mask is B L, so it would have to be expanded to B H N L

        if exists(mask) and mask.ndim != 4:
            mask = rearrange(mask, 'b j -> b 1 1 j')
            mask = mask.expand(-1, heads, q_len, -1)

        # Check if there is a compatible device for flash attention

        config = self.cuda_config if is_cuda else self.cpu_config

        causal = self.causal

        # handle attention bias

        if exists(attn_bias):
            mask_value = -torch.finfo(q.dtype).max // 2
            causal_mask = self.get_mask(q_len, k_len, device)
            attn_bias = attn_bias.masked_fill(causal_mask, mask_value)

            if exists(mask):
                attn_bias = attn_bias.masked_fill(~mask, mask_value)

            mask = attn_bias
            causal = False

        # pytorch 2.0 flash attn: q, k, v, mask, dropout, causal, softmax_scale

        with torch.backends.cuda.sdp_kernel(**config._asdict()):
            out = F.scaled_dot_product_attention(
                q, k, v,
                attn_mask = mask,
                dropout_p = self.dropout if self.training else 0., 
                is_causal = causal
            )

        return out

    def forward(self, q, k, v, mask = None, attn_bias = None):
        """
        einstein notation
        b - batch
        h - heads
        n, i, j - sequence length (base sequence length, source, target)
        d - feature dimension
        """

        q_len, k_len, device = q.shape[-2], k.shape[-2], q.device

        scale = q.shape[-1] ** -0.5

        kv_einsum_eq = 'b j d' if k.ndim == 3 else 'b h j d'

        if self.flash:
            assert not exists(attn_bias)
            return self.flash_attn(q, k, v, mask = mask)

        # similarity

        sim = einsum(f"b h i d, {kv_einsum_eq} -> b h i j", q, k) * scale

        # attention bias

        if exists(attn_bias):
            sim = sim + attn_bias

        # causal mask

        if self.causal:
            causal_mask = self.get_mask(q_len, k_len, device)
            sim = sim.masked_fill(causal_mask, -torch.finfo(sim.dtype).max)

        # key padding mask

        if exists(mask):
            if mask.ndim != 4:
                mask = rearrange(mask, 'b j -> b 1 1 j')
            sim = sim.masked_fill(~mask, -torch.finfo(sim.dtype).max)

        # attention

        attn = sim.softmax(dim=-1)
        attn = self.attn_dropout(attn)

        # aggregate values

        out = einsum(f"b h i j, {kv_einsum_eq} -> b h i d", attn, v)

        return out

best_rq_pytorch/attend.py
lucasnewman/best-rq-pytorch
from pathlib import Path
from typing import Optional

import torch
from torch import nn, Tensor
import torch.nn.functional as F

from torchaudio import transforms as T

from vector_quantize_pytorch.random_projection_quantizer import RandomProjectionQuantizer

from best_rq_pytorch.conformer import ConformerWrapper

from einops import rearrange

# utilities

def exists(val):
    return val is not None

def default(val, d):
    return val if exists(val) else d

def get_mask_subset_prob(mask, prob, min_mask = 0):
    batch, seq, device = *mask.shape, mask.device
    num_to_mask = (mask.sum(dim = -1, keepdim = True) * prob).clamp(min = min_mask)
    logits = torch.rand((batch, seq), device = device)
    logits = logits.masked_fill(~mask, -1)

    randperm = logits.argsort(dim = -1).float()

    num_padding = (~mask).sum(dim = -1, keepdim = True)
    randperm -= num_padding

    subset_mask = randperm < num_to_mask
    subset_mask.masked_fill_(~mask, False)
    return subset_mask

def set_eos_id(t: Tensor, eos_id: int, pad_id: int):
    eos_indices = ((t == pad_id).cumsum(dim = -1) == 0).sum(dim = -1, keepdim = True).long()

    batch_range = torch.arange(t.shape[0], device = t.device, dtype = torch.long)
    batch_range = rearrange(batch_range, '... -> ... 1')

    t = F.pad(t, (0, 1), value = pad_id)
    t[batch_range, eos_indices] = eos_id
    return t

# BEST-RQ model

class BestRQ(nn.Module):
    def __init__(
        self,
        codebook_size: int,
        codebook_dim: int,
        *,
        conformer: ConformerWrapper,
        sample_rate: int = 24_000,
        n_mels: int = 80,
        win_length: int = 960,
        hop_length: Optional[int] = None,
        n_fft: Optional[int] = None,
    ):
        super().__init__()

        self.rpq = nn.Sequential(
            nn.LayerNorm(n_mels, elementwise_affine = True),
            RandomProjectionQuantizer(
                dim = n_mels,
                codebook_size = codebook_size,
                codebook_dim = codebook_dim,
                norm = False
            )
        )
        self.rpq.requires_grad = False

        hop_length = default(hop_length, win_length // 4)

        self.feature_extractor = nn.Sequential(
            T.MelSpectrogram(
                sample_rate = sample_rate,
                n_mels = n_mels,
                win_length = win_length,
                hop_length = hop_length,
                n_fft = default(n_fft, win_length)
            ),
            T.AmplitudeToDB()
        )
        self.feature_extractor.requires_grad = False

        self.conformer = conformer

        self.pad_id = 0
        self.eos_id = codebook_size + 1

    def load(self, path, strict = True):
        path = Path(path)
        assert path.exists()
        pkg = torch.load(str(path), map_location = 'cpu')
        try:
            self.load_state_dict(pkg['model'], strict = strict)
        except Exception:
            self.load_state_dict(pkg['encoder'], strict = strict)
        return pkg

    @torch.no_grad()
    def extract_features(self, x):
        if x.device.type == 'mps':
            # work around ComplexFloat being unavailable with mps
            return self.feature_extractor.cpu()(x.cpu()).to(x.device)
        else:
            return self.feature_extractor(x)

    def forward(
        self,
        x = None,
        labels = None,
        mask = None,
        return_labels = False,
        return_emb = False,
        return_layer_output: Optional[int] = None
    ):
        assert exists(x) or exists(labels), "either input or labels must be provided"

        if exists(x) and not exists(labels):
            with torch.no_grad():
                features = self.extract_features(x)

                # project labels from features

                features = rearrange(features, "b c n -> b n c")
                labels = self.rpq(features)

                # offset for the pad id

                labels = torch.where(labels != self.pad_id, labels + 1, self.pad_id)

        if return_labels:
            return labels

        if not exists(mask):
            mask = torch.cat([rearrange(x != self.pad_id, "n -> 1 n") for x in labels.unbind(dim = 0)], dim = 0)

        outputs = self.conformer(
            labels,
            mask = mask,
            return_layer_output = return_layer_output,
            return_emb = return_emb
        )

        return outputs
    
# pretraining task wrapper

class BestRQPretrainWrapper(nn.Module):
    def __init__(
        self,
        model: BestRQ,
        mask_prob: float = 0.6
    ):
        super().__init__()

        self.model = model
        self.mask_prob = mask_prob
        self.pad_id = self.model.pad_id

    def forward(self, x):
        # determine labels from the random projection quantizer

        labels = self.model(x, return_labels = True)

        # mask a subset of the labels for pretraining
        # TODO: mask sequences of n milliseconds/frames

        seq_mask = torch.cat([rearrange(x != self.pad_id, "n -> 1 n") for x in labels.unbind(dim = 0)], dim = 0)
        mask = get_mask_subset_prob(seq_mask, self.mask_prob)

        # predict and compute ce loss

        logits = self.model(labels = labels, mask = mask)
        logits = rearrange(logits, "b n c -> b c n")

        masked_labels = labels.masked_fill(~mask, self.pad_id)

        loss = F.cross_entropy(logits, masked_labels, ignore_index = self.pad_id)

        return loss, logits

best_rq_pytorch/best_rq.py
lucasnewman/best-rq-pytorch
import math

import torch
from torch import nn
import torch.nn.functional as F

from einops import rearrange
from einops.layers.torch import Rearrange

from beartype import beartype
from beartype.typing import Union, Dict

from best_rq_pytorch.attend import Attend

# helpers

def exists(val):
    return val is not None

def default(val, d):
    return val if exists(val) else d

def calc_same_padding(kernel_size):
    pad = kernel_size // 2
    return (pad, pad - (kernel_size + 1) % 2)

# rotary embedding

class RotaryEmbedding(nn.Module):
    def __init__(self, dim, theta = 10000):
        super().__init__()
        inv_freq = 1.0 / (theta ** (torch.arange(0, dim, 2).float() / dim))
        self.register_buffer("inv_freq", inv_freq, persistent = False)

    @property
    def device(self):
        return next(self.buffers()).device

    def forward(self, seq_len):
        t = torch.arange(seq_len, device = self.device).type_as(self.inv_freq)
        freqs = torch.einsum('i , j -> i j', t, self.inv_freq)
        freqs = torch.cat((freqs, freqs), dim = -1)
        return freqs

def rotate_half(x):
    x1, x2 = x.chunk(2, dim=-1)
    return torch.cat((-x2, x1), dim=-1)

def apply_rotary_pos_emb(pos, t):
    return (t * pos.cos()) + (rotate_half(t) * pos.sin())

# t5 relative positional bias

class T5RelativePositionBias(nn.Module):
    def __init__(
        self,
        scale = 1.,
        num_buckets = 32,
        max_distance = 128,
        heads = 8
    ):
        super().__init__()
        self.scale = scale
        self.num_buckets = num_buckets
        self.max_distance = max_distance
        self.relative_attention_bias = nn.Embedding(num_buckets, heads)

    @staticmethod
    def _relative_position_bucket(
        relative_position,
        num_buckets = 32,
        max_distance = 128
    ):
        ret = 0
        n = -relative_position

        num_buckets //= 2
        ret += (n < 0).long() * num_buckets
        n = torch.abs(n)

        max_exact = num_buckets // 2
        is_small = n < max_exact

        val_if_large = max_exact + (
            torch.log(n.float() / max_exact) / math.log(max_distance / max_exact) * (num_buckets - max_exact)
        ).long()

        val_if_large = torch.min(
            val_if_large,
            torch.full_like(val_if_large, num_buckets - 1)
        )

        ret += torch.where(is_small, n, val_if_large)
        return ret

    @property
    def device(self):
        return next(self.parameters()).device

    def forward(self, n):
        pos = torch.arange(n, device = self.device).long()
        rel_pos = rearrange(pos, 'j -> 1 j') - rearrange(pos, 'i -> i 1')

        rp_bucket = self._relative_position_bucket(rel_pos, num_buckets = self.num_buckets, max_distance = self.max_distance)
        values = self.relative_attention_bias(rp_bucket)

        bias = rearrange(values, 'i j h -> h i j')
        return bias * self.scale

# conformer

class Swish(nn.Module):
    def forward(self, x):
        return x * x.sigmoid()

class GLU(nn.Module):
    def __init__(self, dim):
        super().__init__()
        self.dim = dim

    def forward(self, x):
        out, gate = x.chunk(2, dim=self.dim)
        return out * gate.sigmoid()

class DepthWiseConv1d(nn.Module):
    def __init__(self, chan_in, chan_out, kernel_size, padding):
        super().__init__()
        self.padding = padding
        self.conv = nn.Conv1d(chan_in, chan_out, kernel_size, groups = chan_in)

    def forward(self, x):
        x = F.pad(x, self.padding)
        return self.conv(x)

# attention, feedforward, and conv module

class Scale(nn.Module):
    def __init__(self, scale, fn):
        super().__init__()
        self.fn = fn
        self.scale = scale

    def forward(self, x, **kwargs):
        return self.fn(x, **kwargs) * self.scale

class ChanLayerNorm(nn.Module):
    def __init__(self, dim):
        super().__init__()
        self.gamma = nn.Parameter(torch.ones(1, dim, 1))

    def forward(self, x):
        eps = 1e-6 if x.dtype == torch.float32 else 1e-4
        var = torch.var(x, dim = 1, unbiased = False, keepdim = True)
        mean = torch.mean(x, dim = 1, keepdim = True)
        return (x - mean) * var.clamp(min = eps).rsqrt() * self.gamma

class PreNorm(nn.Module):
    def __init__(self, dim, fn):
        super().__init__()
        self.fn = fn
        self.norm = nn.LayerNorm(dim)

    def forward(self, x, **kwargs):
        x = self.norm(x)
        return self.fn(x, **kwargs)

class Attention(nn.Module):
    def __init__(
        self,
        dim,
        heads = 8,
        dim_head = 64,
        dropout = 0.,
        flash = True
    ):
        super().__init__()
        inner_dim = dim_head * heads
        self.heads = heads
        self.scale = dim_head ** -0.5

        self.attend = Attend(
            flash = flash,
            dropout = dropout
        )

        self.dropout = nn.Dropout(dropout)

        self.to_q = nn.Linear(dim, inner_dim, bias = False)
        self.to_kv = nn.Linear(dim, inner_dim * 2, bias = False)
        self.to_out = nn.Linear(inner_dim, dim)

    def forward(
        self,
        x,
        context = None,
        mask = None,
        rotary_emb = None,
        attn_bias = None
    ):
        h = self.heads
        context = default(context, x)

        q, k, v = (self.to_q(x), *self.to_kv(context).chunk(2, dim = -1))
        q, k, v = map(lambda t: rearrange(t, 'b n (h d) -> b h n d', h = h), (q, k, v))

        if exists(rotary_emb):
            q = apply_rotary_pos_emb(rotary_emb, q)
            k = apply_rotary_pos_emb(rotary_emb, k)

        out = self.attend(q, k, v, mask = mask, attn_bias = attn_bias)

        out = rearrange(out, 'b h n d -> b n (h d)')
        return self.to_out(out)

class FeedForward(nn.Module):
    def __init__(
        self,
        dim,
        mult = 4,
        dropout = 0.
    ):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(dim, dim * mult),
            Swish(),
            nn.Dropout(dropout),
            nn.Linear(dim * mult, dim),
            nn.Dropout(dropout)
        )

    def forward(self, x):
        return self.net(x)

class ConformerConvModule(nn.Module):
    def __init__(
        self,
        dim,
        causal = False,
        expansion_factor = 2,
        kernel_size = 31,
        dropout = 0.
    ):
        super().__init__()

        inner_dim = dim * expansion_factor
        padding = calc_same_padding(kernel_size) if not causal else (kernel_size - 1, 0)

        self.net = nn.Sequential(
            nn.LayerNorm(dim),
            Rearrange('b n c -> b c n'),
            nn.Conv1d(dim, inner_dim * 2, 1),
            GLU(dim=1),
            DepthWiseConv1d(inner_dim, inner_dim, kernel_size = kernel_size, padding = padding),
            Swish(),
            ChanLayerNorm(inner_dim),
            nn.Conv1d(inner_dim, dim, 1),
            Rearrange('b c n -> b n c'),
            nn.Dropout(dropout)
        )

    def forward(self, x):
        return self.net(x)

# Conformer Block

class ConformerBlock(nn.Module):
    def __init__(
        self,
        *,
        dim,
        dim_head = 64,
        heads = 8,
        ff_mult = 4,
        conv_expansion_factor = 2,
        conv_kernel_size = 31,
        attn_dropout = 0.,
        attn_flash = True,
        ff_dropout = 0.,
        conv_dropout = 0.,
        conv_causal = False
    ):
        super().__init__()
        self.ff1 = FeedForward(dim = dim, mult = ff_mult, dropout = ff_dropout)
        self.attn = Attention(dim = dim, dim_head = dim_head, heads = heads, dropout = attn_dropout, flash = attn_flash)
        self.conv = ConformerConvModule(dim = dim, causal = conv_causal, expansion_factor = conv_expansion_factor, kernel_size = conv_kernel_size, dropout = conv_dropout)
        self.ff2 = FeedForward(dim = dim, mult = ff_mult, dropout = ff_dropout)

        self.attn = PreNorm(dim, self.attn)
        self.ff1 = Scale(0.5, PreNorm(dim, self.ff1))
        self.ff2 = Scale(0.5, PreNorm(dim, self.ff2))

        self.post_norm = nn.LayerNorm(dim)

    def forward(
        self,
        x,
        mask = None,
        rotary_emb = None,
        attn_bias = None
    ):
        x = self.ff1(x) + x
        x = self.attn(x, mask = mask, rotary_emb = rotary_emb, attn_bias = attn_bias) + x
        x = self.conv(x) + x
        x = self.ff2(x) + x
        x = self.post_norm(x)
        return x

# Conformer

class Conformer(nn.Module):
    def __init__(
        self,
        dim,
        *,
        depth,
        dim_head = 64,
        heads = 8,
        ff_mult = 4,
        conv_expansion_factor = 2,
        conv_kernel_size = 31,
        attn_dropout = 0.,
        ff_dropout = 0.,
        conv_dropout = 0.,
        conv_causal = False,
        attn_flash = True,
        t5_rel_pos_bias = False
    ):
        super().__init__()

        assert not (t5_rel_pos_bias and attn_flash), 'flash attention is not compatible with learned bias'

        self.dim = dim
        self.layers = nn.ModuleList([])

        self.rotary_emb = RotaryEmbedding(dim_head) if not t5_rel_pos_bias else None
        self.rel_pos_bias = T5RelativePositionBias(dim_head ** 0.5, heads = heads) if t5_rel_pos_bias else None

        for _ in range(depth):
            self.layers.append(ConformerBlock(
                dim = dim,
                dim_head = dim_head,
                heads = heads,
                ff_mult = ff_mult,
                conv_expansion_factor = conv_expansion_factor,
                conv_kernel_size = conv_kernel_size,
                attn_dropout = attn_dropout,
                ff_dropout = ff_dropout,
                conv_dropout = conv_dropout,
                conv_causal = conv_causal,
                attn_flash = attn_flash
            ))

    def forward(self, x, mask = None, return_layer_output = None):
        seq_len = x.shape[-2]

        rotary_emb = self.rotary_emb(seq_len) if exists(self.rotary_emb) else None
        attn_bias = self.rel_pos_bias(seq_len) if exists(self.rel_pos_bias) else None

        for index, block in enumerate(self.layers):
            x = block(
                x,
                mask = mask,
                rotary_emb = rotary_emb,
                attn_bias = attn_bias
            )

            # return a specific layer's output early if provided

            if exists(return_layer_output) and return_layer_output == index + 1:
                return x

        return x

class ConformerWrapper(nn.Module):
    @beartype
    def __init__(
        self,
        *,
        num_tokens,
        conformer: Union[Conformer, Dict[str, any]]
    ):
        super().__init__()
        self.conformer = conformer

        if isinstance(conformer, dict):
            self.conformer = Conformer(**self.conformer)

        embed_dim = self.conformer.dim
        self.num_token_ids = num_tokens + 2 # add pad and eos tokens

        self.pad_id = 0
        self.eos_id = num_tokens + 1

        self.code_embeds = nn.Embedding(self.num_token_ids, embed_dim, padding_idx = self.pad_id)

        self.to_logits = nn.Sequential(
            nn.LayerNorm(embed_dim),
            nn.Linear(embed_dim, self.num_token_ids)
        )

    def forward(
        self,
        x,
        *,
        mask = None,
        cond = None,
        return_layer_output = None,
        return_emb = False
    ):
        x = self.code_embeds(x)

        if exists(cond):
            if cond.ndim == 2:
                cond = rearrange(cond, 'b d -> b 1 d')

            x = x + cond

        x = self.conformer(x, mask = mask, return_layer_output = return_layer_output)

        if exists(return_layer_output) or return_emb:
            return x
        else:
            return self.to_logits(x)

best_rq_pytorch/conformer.py
lucasnewman/best-rq-pytorch
from pathlib import Path
from functools import wraps

from beartype import beartype
from beartype.typing import Optional, Tuple
from beartype.door import is_bearable

import torch
from torch.nn import functional as F
from torch.nn.utils.rnn import pad_sequence
from torch.utils.data import Dataset, DataLoader

import torchaudio

from einops import rearrange

def exists(val):
    return val is not None

class AudioDataset(Dataset):
    def __init__(
        self,
        folder,
        max_length_in_seconds: Optional[int] = None,
        pad_to_max_length = True
    ):
        super().__init__()
        path = Path(folder)
        assert path.exists(), 'folder does not exist'

        files = list(path.glob('**/*.flac'))
        assert len(files) > 0, 'no files found'

        self.files = files
        self.max_length = (max_length_in_seconds * 24_000) if exists(max_length_in_seconds) else None
        self.pad_to_max_length = pad_to_max_length
    
    def __len__(self):
        return len(self.files)
    
    def __getitem__(self, idx):
        file = self.files[idx]
        
        # this assumes 24kHz single-channel audio

        wave, _ = torchaudio.load(file)
        wave = rearrange(wave, '1 n -> n')
        
        if exists(self.max_length):
            wave_length = wave.shape[0]
            
            if wave_length > self.max_length:
                # take a random crop to the max length
                start = torch.randint(0, wave_length - self.max_length, (1, )).item()
                wave = wave[start:start + self.max_length]
            elif self.pad_to_max_length:
                # work around variable length sequences recomputing the graph on mps
                wave = F.pad(wave, (0, self.max_length - wave_length))

        return wave
    
# data loader utilities

def collate_one_or_multiple_tensors(fn):
    @wraps(fn)
    def inner(data):
        is_one_data = not isinstance(data[0], tuple)

        if is_one_data:
            data = fn(data)
            return (data,)

        outputs = []
        for datum in zip(*data):
            if is_bearable(datum, Tuple[str, ...]):
                output = list(datum)
            else:
                output = fn(datum)

            outputs.append(output)

        return tuple(outputs)

    return inner

@collate_one_or_multiple_tensors
def curtail_to_shortest_collate(data):
    min_len = min(*[datum.shape[0] for datum in data])
    data = [datum[:min_len] for datum in data]
    return torch.stack(data)

@collate_one_or_multiple_tensors
def pad_to_longest_fn(data):
    return pad_sequence(data, batch_first = True)
    
def get_dataloader(ds, pad_to_longest = True, **kwargs):
    collate_fn = pad_to_longest_fn if pad_to_longest else curtail_to_shortest_collate
    return DataLoader(ds, collate_fn = collate_fn, **kwargs)

best_rq_pytorch/data.py
lucasnewman/best-rq-pytorch
from torch.optim import AdamW, Adam

def separate_weight_decayable_params(params):
    wd_params, no_wd_params = [], []
    for param in params:
        param_list = no_wd_params if param.ndim < 2 else wd_params
        param_list.append(param)
    return wd_params, no_wd_params

def get_optimizer(
    params,
    lr = 1e-4,
    wd = 1e-2,
    betas = (0.9, 0.99),
    eps = 1e-8,
    filter_by_requires_grad = False,
    group_wd_params = True
):
    has_wd = wd > 0

    if filter_by_requires_grad:
        params = list(filter(lambda t: t.requires_grad, params))

    if group_wd_params and has_wd:
        wd_params, no_wd_params = separate_weight_decayable_params(params)

        params = [
            {'params': wd_params},
            {'params': no_wd_params, 'weight_decay': 0},
        ]

    if not has_wd:
        return Adam(params, lr = lr, betas = betas, eps = eps)

    return AdamW(params, lr = lr, weight_decay = wd, betas = betas, eps = eps)

best_rq_pytorch/optimizer.py
lucasnewman/best-rq-pytorch
import re
from pathlib import Path
from shutil import rmtree

from beartype import beartype
from beartype.typing import Optional

import torch
from torch import nn
import torch.nn.functional as F
from torch.optim.lr_scheduler import CosineAnnealingLR
from torch.utils.data import Dataset, random_split

from best_rq_pytorch.best_rq import BestRQ, BestRQPretrainWrapper
from best_rq_pytorch.optimizer import get_optimizer
from best_rq_pytorch.data import get_dataloader

from accelerate import Accelerator, DistributedType

# utilities

def exists(val):
    return val is not None

def noop(*args, **kwargs):
    pass

def cycle(dl):
    while True:
        for data in dl:
            yield data

def cast_tuple(t):
    return t if isinstance(t, (tuple, list)) else (t,)

def yes_or_no(question):
    answer = input(f'{question} (y/n) ')
    return answer.lower() in ('yes', 'y')

def accum_log(log, new_logs):
    for key, new_value in new_logs.items():
        old_value = log.get(key, 0.)
        log[key] = old_value + new_value
    return log

def checkpoint_num_steps(checkpoint_path):
    """Returns the number of steps trained from a checkpoint based on the filename.

    Filename format assumed to be something like "/path/to/speech.speech.20000.pt" which is
    for 20k train steps. Returns 20000 in that case.
    """
    results = re.findall(r'\d+', str(checkpoint_path))

    if len(results) == 0:
        return 0

    return int(results[-1])

def mask_after_eos(target, eos_id, pad_id):
    mask = (target == eos_id).cumsum(dim = -1) > 0
    mask = F.pad(mask, (1, -1), value = False)
    return target.masked_fill(mask, pad_id)

# pretraining

class BestRQPretrainer(nn.Module):
    @beartype
    def __init__(
        self,
        model: BestRQ,
        *,
        num_train_steps,
        num_warmup_steps,
        batch_size,
        dataset: Optional[Dataset] = None,
        mask_prob: float = 0.6,
        lr = 1e-4,
        initial_lr = 1e-5,
        grad_accum_every = 1,
        wd = 0.,
        max_grad_norm = 0.5,
        valid_frac = 0.05,
        random_split_seed = 42,
        log_every = 10,
        save_results_every = 100,
        save_model_every = 1000,
        results_folder = './results',
        force_clear_prev_results = None,
        split_batches = False,
        drop_last = False,
        accelerate_kwargs: dict = dict(),
    ):
        super().__init__()

        self.accelerator = Accelerator(
            split_batches = split_batches,
            **accelerate_kwargs
        )

        self.model = model

        self.train_wrapper = BestRQPretrainWrapper(
            model = model,
            mask_prob = mask_prob
        )

        self.register_buffer('steps', torch.Tensor([0]))

        self.num_train_steps = num_train_steps
        self.num_warmup_steps = num_warmup_steps
        self.batch_size = batch_size
        self.grad_accum_every = grad_accum_every

        # optimizers
        self.lr = lr
        self.initial_lr = initial_lr
        self.optim = get_optimizer(model.parameters(), lr = lr, wd = wd)
        self.scheduler = CosineAnnealingLR(self.optim, T_max = num_train_steps)

        # max grad norm

        self.max_grad_norm = max_grad_norm

        # create dataset

        self.ds = dataset

        # split for validation

        if valid_frac > 0:
            train_size = int((1 - valid_frac) * len(self.ds))
            valid_size = len(self.ds) - train_size
            self.ds, self.valid_ds = random_split(self.ds, [train_size, valid_size], generator = torch.Generator().manual_seed(random_split_seed))
            self.print(f'training with dataset of {len(self.ds)} samples and validating with randomly splitted {len(self.valid_ds)} samples')
        else:
            self.valid_ds = self.ds
            self.print(f'training with shared training and valid dataset of {len(self.ds)} samples')

        assert len(self.ds) >= batch_size, 'dataset must have sufficient samples for training'
        assert len(self.valid_ds) >= batch_size, f'validation dataset must have sufficient number of samples (currently {len(self.valid_ds)}) for training'

        # dataloader

        self.dl = get_dataloader(self.ds, batch_size = batch_size, shuffle = True, drop_last = drop_last)

        self.valid_dl = get_dataloader(self.valid_ds, batch_size = batch_size, shuffle = True, drop_last = drop_last)

        # prepare with accelerator

        (
            self.train_wrapper,
            self.optim,
            self.scheduler,
            self.dl,
            self.valid_dl
        ) = self.accelerator.prepare(
            self.train_wrapper,
            self.optim,
            self.scheduler,
            self.dl,
            self.valid_dl
        )

        # dataloader iterators

        self.dl_iter = cycle(self.dl)
        self.valid_dl_iter = cycle(self.valid_dl)

        self.log_every = log_every
        self.save_model_every = save_model_every
        self.save_results_every = save_results_every

        self.results_folder = Path(results_folder)

        if self.is_main and force_clear_prev_results is True or (not exists(force_clear_prev_results) and len([*self.results_folder.glob('**/*')]) > 0 and yes_or_no('do you want to clear previous experiment checkpoints and results?')):
            rmtree(str(self.results_folder))

        self.results_folder.mkdir(parents = True, exist_ok = True)

        hps = {"num_train_steps": num_train_steps, "num_warmup_steps": num_warmup_steps, "learning_rate": lr, "initial_learning_rate": lr}
        self.accelerator.init_trackers("speechspeech", config=hps)

    def save(self, path):
        pkg = dict(
            model = self.accelerator.get_state_dict(self.model),
            optim = self.optim.state_dict(),
            scheduler = self.scheduler.state_dict()
        )
        torch.save(pkg, path)

    def load(self, path, restore_optimizer = True):
        model = self.accelerator.unwrap_model(self.model)
        pkg = model.load(path)

        if restore_optimizer:
            self.optim.load_state_dict(pkg['optim'])
            self.scheduler.load_state_dict(pkg['scheduler'])

            # + 1 to start from the next step and avoid overwriting the last checkpoint
            self.steps = torch.tensor([checkpoint_num_steps(path) + 1], device=self.device)

    def print(self, msg):
        self.accelerator.print(msg)

    @property
    def device(self):
        return self.accelerator.device

    @property
    def is_distributed(self):
        return not (self.accelerator.distributed_type == DistributedType.NO and self.accelerator.num_processes == 1)

    @property
    def is_main(self):
        return self.accelerator.is_main_process

    @property
    def is_local_main(self):
        return self.accelerator.is_local_main_process

    def warmup(self, step):
        if step < self.num_warmup_steps:
            return self.initial_lr + (self.lr - self.initial_lr) * step / self.num_warmup_steps
        else:
            return self.lr
    
    def train_step(self):
        steps = int(self.steps.item())

        self.train_wrapper.train()

        # adjust the lr according to the schedule

        if steps < self.num_warmup_steps:
            # Apply warmup
            lr = self.warmup(steps)
            for param_group in self.optim.param_groups:
                param_group['lr'] = lr
        else:
            # After warmup period, start to apply CosineAnnealingLR
            self.scheduler.step()

        # logs

        logs = {}

        # update encoder

        for _ in range(self.grad_accum_every):
            x, = next(self.dl_iter)

            loss, _ = self.train_wrapper(x)

            self.accelerator.backward(loss / self.grad_accum_every)

            accum_log(logs, {'loss': loss.item() / self.grad_accum_every})

        if exists(self.max_grad_norm):
            self.accelerator.clip_grad_norm_(self.model.parameters(), self.max_grad_norm)

        self.optim.step()
        self.optim.zero_grad()

        # log

        if not (steps % self.log_every):
            self.print(f"{steps}: loss: {logs['loss']:0.3f}")

        self.accelerator.log({"train_loss": logs['loss']}, step=steps)

        # sample results every so often

        self.accelerator.wait_for_everyone()

        if self.is_main and not (steps % self.save_results_every):
            x, = next(self.valid_dl_iter)

            with torch.inference_mode():
                self.train_wrapper.eval()
                valid_loss, _ = self.train_wrapper(x)

            self.print(f'{steps}: valid loss {valid_loss:0.3f}')
            self.accelerator.log({"valid_loss": valid_loss}, step=steps)

        # save model every so often

        if self.is_main and not (steps % self.save_model_every):
            model_path = str(self.results_folder / f'bestrq.{steps}.pt')
            self.save(model_path)

            self.print(f'{steps}: saving model to {str(self.results_folder)}')

        self.steps += 1
        return logs

    def train(self, log_fn = noop):
        while self.steps < self.num_train_steps:
            logs = self.train_step()
            log_fn(logs)

        self.print('training complete')

best_rq_pytorch/trainer.py
lucasnewman/best-rq-pytorch
import joblib
from pathlib import Path
from tqdm import tqdm

import torch

from einops import rearrange

from best_rq_pytorch.best_rq import BestRQ
from best_rq_pytorch.conformer import ConformerWrapper
from best_rq_pytorch.data import AudioDataset

from sklearn.cluster import MiniBatchKMeans

pretrained_checkpoint = "results/bestrq.XXXX.pt"
km_output_path = "results/bestrq_km1024.bin"
accelerator = "cuda"

# k-means utilities

def get_km_model(
    n_clusters,
    init,
    max_iter,
    batch_size,
    tol,
    max_no_improvement,
    n_init,
    reassignment_ratio,
):
    return MiniBatchKMeans(
        n_clusters = n_clusters,
        init = init,
        max_iter = max_iter,
        batch_size = batch_size,
        verbose = 1,
        compute_labels = False,
        tol = tol,
        max_no_improvement = max_no_improvement,
        init_size = None,
        n_init = n_init,
        reassignment_ratio = reassignment_ratio,
    )


# set up the encoder

brq = BestRQ(
    codebook_size = 1024,
    codebook_dim = 16,
    sample_rate = 24_000,
    n_mels = 80,
    win_length = 960,
    hop_length = 960 // 3,
    conformer = ConformerWrapper(
        num_tokens = 1024,
        conformer = dict(
            dim = 1024,
            depth = 24,
            heads = 16,
            conv_kernel_size = 5,
            ff_mult = 4,
            attn_dropout = 0.1,
            ff_dropout = 0.1,
            conv_dropout = 0.1,
            attn_flash = False
        )
    )
).to(accelerator)

brq.load(pretrained_checkpoint)

# load the dataset

dataset_folder = "..."
ds = AudioDataset(dataset_folder)

# collect activations for k-means

output_layer = 13

samples = 0
activations = []

for wave in tqdm(ds):
    with torch.no_grad():
        activation = brq(rearrange(wave, "n -> 1 n").to(accelerator), return_layer_output = output_layer)
        activations.append(activation.cpu())

activations = rearrange(torch.cat(activations, dim = 1), "1 n d -> n d")

# run k-means -- the centroid indicies will be the semantic token ids

n_clusters = 1024
max_iter = 100
batch_size = 10000

km_model = get_km_model(
    n_clusters,
    init = "k-means++",
    max_iter = max_iter,
    batch_size = batch_size,
    tol = 0.,
    max_no_improvement = 100,
    n_init = 20,
    reassignment_ratio = 0.,
)
km_model.fit(activations)

joblib.dump(km_model, km_output_path)

print(f"saved k-means quantizer to path: {km_output_path}")

# generate semantic tokens from the model + k-means quantizer

for file_index, wave in tqdm(enumerate(ds)):
    path = Path(ds.files[file_index])
    output_path = path.with_suffix(".semantic.pt")
    
    with torch.no_grad():
        activations = brq(rearrange(wave, "n -> 1 n"), return_layer_output = output_layer)
        activations = rearrange(activations, "1 n d -> n d")
        
        semantic_labels = km_model.predict(activations)
        torch.save(torch.tensor(semantic_labels, dtype = torch.long), output_path)

examples/fit_kmeans.py
lucasnewman/best-rq-pytorch
from best_rq_pytorch.best_rq import BestRQ
from best_rq_pytorch.trainer import BestRQPretrainer
from best_rq_pytorch.conformer import ConformerWrapper
from best_rq_pytorch.data import AudioDataset

# load the dataset

dataset_folder = "..."

ds = AudioDataset(
    dataset_folder,
    max_length_in_seconds = 10
)

# set up the model

brq = BestRQ(
    codebook_size = 1024,
    codebook_dim = 16,
    sample_rate = 24_000,
    n_mels = 80,
    win_length = 960,
    hop_length = 960 // 3,
    conformer = ConformerWrapper(
        num_tokens = 1024,
        conformer = dict(
            dim = 1024,
            depth = 24,
            heads = 16,
            conv_kernel_size = 5,
            ff_mult = 4,
            attn_dropout = 0.1,
            ff_dropout = 0.1,
            conv_dropout = 0.1,
            attn_flash = False
        )
    )
)

print(f"trainable parameters: {sum(p.numel() for p in brq.parameters() if p.requires_grad)}")

# train the model

trainer = BestRQPretrainer(
    model = brq,
    dataset = ds,
    num_train_steps = 100_000,
    lr = 5e-4,
    num_warmup_steps = 1000,
    initial_lr = 1e-5,
    batch_size = 8,
    grad_accum_every = 1,
    mask_prob = 0.6
)

trainer.train()

examples/pretrain.py
lucasnewman/best-rq-pytorch
from setuptools import setup, find_packages

setup(
    name = 'best-rq-pytorch',
    packages = find_packages(exclude=[]),
    version = '0.0.2',
    license='MIT',
    description = 'BEST-RQ - Pytorch',
    author = 'Lucas Newman',
    author_email = 'lucasnewman@me.com',
    long_description_content_type = 'text/markdown',
    url = 'https://github.com/lucasnewman/best-rq-pytorch',
    keywords = [
        'artificial intelligence',
        'asr',
        'audio-generation,'
        'deep learning',
        'transformers',
        'text-to-speech'
    ],
    install_requires=[
        'beartype',
        'einops>=0.6.1',
        'rotary-embedding-torch>=0.3.0',
        'torch>=2.0',
        'torchaudio>=2.0',
        'tqdm',
        'vector-quantize-pytorch>=1.7.0'
    ],
    classifiers=[
        'Development Status :: 4 - Beta',
        'Intended Audience :: Developers',
        'Topic :: Scientific/Engineering :: Artificial Intelligence',
        'License :: OSI Approved :: MIT License',
        'Programming Language :: Python :: 3.9',
    ],
)

setup.py
lucasnewman/best-rq-pytorch