genhpf 1.0.11__py3-none-any.whl

genhpf/models/genhpf_wav2vec2.py

@@ -0,0 +1,304 @@
+import logging
+from dataclasses import dataclass, field
+from typing import Tuple
+
+import torch
+import torch.nn as nn
+
+from genhpf.models import register_model
+from genhpf.models.genhpf import GenHPF, GenHPFConfig
+from genhpf.modules import GradMultiply, GumbelVectorQuantizer, LayerNorm
+from genhpf.utils import utils
+from genhpf.utils.data_utils import compute_mask_indices
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class GenHPFWav2Vec2Config(GenHPFConfig):
+    logit_temp: float = field(default=0.1, metadata={"help": "temperature to divide logits by"})
+    latent_vars: int = field(
+        default=320, metadata={"help": "number of latent variables V in each group of the codebook"}
+    )
+    latent_groups: int = field(
+        default=2, metadata={"help": "number of groups G of latent variables in the codebook"}
+    )
+    latent_temp: Tuple[float, float, float] = field(
+        default=(2, 0.5, 0.999995),
+        metadata={
+            "help": "temperature for latent variable sampling. "
+            "can be tuple of 3 values (start, end, decay)"
+        },
+    )
+    final_dim: int = field(
+        default=128, metadata={"help": "project final representations and targets to this many dimensions."}
+    )
+
+    num_negatives: int = field(
+        default=25, metadata={"help": "number of negative examples from the same sample"}
+    )
+    codebook_negatives: int = field(default=0, metadata={"help": "number of negative examples codebook"})
+
+    # mask
+    mask_prob: float = field(default=0.65, metadata={"help": "probability of replacing a token with mask"})
+    mask_length: int = field(default=1, metadata={"help": "mask length"})
+    no_mask_overlap: bool = field(default=False, metadata={"help": "whether to allow masks to overlap"})
+    mask_min_space: int = field(
+        default=0, metadata={"help": "min space between spans (if no overlap is enabled)"}
+    )
+
+    feature_grad_mult: float = field(
+        default=0.1, metadata={"help": "multiply event encoder gradients by this"}
+    )
+    dropout_input: float = field(default=0.1, metadata={"help": "dropout to apply to the input"})
+    dropout_features: float = field(default=0.1, metadata={"help": "dropout to apply to the features"})
+
+
+@register_model("genhpf_wav2vec2", dataclass=GenHPFWav2Vec2Config)
+class GenHPFWav2Vec2(GenHPF):
+    def __init__(self, cfg: GenHPFWav2Vec2Config):
+        super().__init__(cfg)
+
+        self.logit_temp = cfg.logit_temp
+
+        self.mask_prob = cfg.mask_prob
+        self.mask_length = cfg.mask_length
+        self.no_mask_overlap = cfg.no_mask_overlap
+        self.mask_min_space = cfg.mask_min_space
+
+        self.feature_grad_mult = cfg.feature_grad_mult
+
+        self.codebook_negatives = cfg.codebook_negatives
+
+        self.quantizer = GumbelVectorQuantizer(
+            dim=cfg.agg_embed_dim,
+            num_vars=cfg.latent_vars,
+            temp=cfg.latent_temp,
+            groups=cfg.latent_groups,
+            combine_groups=False,
+            vq_dim=cfg.agg_embed_dim,
+            time_first=True,
+        )
+
+        self.layer_norm = LayerNorm(cfg.agg_embed_dim)
+        self.dropout_input = nn.Dropout(cfg.dropout_input)
+        self.dropout_features = nn.Dropout(cfg.dropout_features)
+        self.project_q = nn.Linear(cfg.agg_embed_dim, cfg.final_dim)
+        self.final_proj = nn.Linear(cfg.agg_embed_dim, cfg.final_dim)
+        self.mask_emb = nn.Parameter(torch.FloatTensor(cfg.agg_embed_dim).uniform_())
+        self.cross_sample_negatives = 0
+        self.n_negatives = cfg.num_negatives
+
+    @classmethod
+    def build_model(cls, cfg):
+        """Build a new model instance."""
+        return cls(cfg)
+
+    def compute_preds(self, x, y, negatives):
+        neg_is_pos = (y == negatives).all(-1)
+        y = y.unsqueeze(0)
+        targets = torch.cat([y, negatives], dim=0)
+
+        logits = torch.cosine_similarity(x.float(), targets.float(), dim=-1)
+
+        logits = logits / self.logit_temp
+        logits = logits.type_as(x)
+
+        if neg_is_pos.any():
+            logits[1:][neg_is_pos] = float("-inf")
+
+        return logits
+
+    def apply_mask(
+        self,
+        x: torch.Tensor,
+        padding_mask: torch.Tensor,
+        mask_indices=None,
+    ):
+        bsz, tsz, csz = x.shape
+
+        if self.mask_prob > 0:
+            if mask_indices is None:
+                mask_indices = compute_mask_indices(
+                    shape=(bsz, tsz),
+                    padding_mask=padding_mask,
+                    mask_prob=self.mask_prob,
+                    mask_length=self.mask_length,
+                    mask_type="static",
+                    min_masks=2,
+                    no_overlap=self.no_mask_overlap,
+                    min_space=self.mask_min_space,
+                )
+                mask_indices = torch.from_numpy(mask_indices).to(x.device)
+            x[mask_indices] = self.mask_emb
+        else:
+            mask_indices = None
+
+        return x, mask_indices
+
+    def sample_negatives(self, y, num):
+        if self.n_negatives == 0 and self.cross_sample_negatives == 0:
+            return y.new(0)
+
+        batch_size, time_size, feature_size = y.shape
+        y = y.view(-1, feature_size)  # B x T x C -> (B x T) x C
+
+        cross_high = time_size * batch_size
+        high = time_size
+        with torch.no_grad():
+            assert high > 1, f"{batch_size, time_size, feature_size}"
+
+            if self.n_negatives > 0:
+                time_sizes = utils.buffered_arange(num).unsqueeze(-1).expand(-1, self.n_negatives).flatten()
+                neg_idxs = torch.randint(low=0, high=high - 1, size=(batch_size, self.n_negatives * num))
+                neg_idxs[neg_idxs >= time_sizes] += 1
+
+            if self.cross_sample_negatives > 0:
+                time_sizes = (
+                    utils.buffered_arange(num).unsqueeze(-1).expand(-1, self.cross_sample_negatives).flatten()
+                )
+
+                cross_neg_idxs = torch.randint(
+                    low=0, high=cross_high - 1, size=(batch_size, self.cross_sample_negatives * num)
+                )
+                cross_neg_idxs[cross_neg_idxs >= time_sizes] += 1
+
+        if self.n_negatives > 0:
+            for i in range(1, batch_size):
+                neg_idxs[i] += i * high
+        else:
+            neg_idxs = cross_neg_idxs
+
+        if self.cross_sample_negatives > 0 and self.n_negatives > 0:
+            neg_idxs = torch.cat([neg_idxs, cross_neg_idxs], dim=1)
+
+        negs = y[neg_idxs.view(-1)]
+        negs = negs.view(
+            batch_size, num, self.n_negatives + self.cross_sample_negatives, feature_size
+        ).permute(
+            2, 0, 1, 3
+        )  # to N x B x T x C
+
+        return negs, neg_idxs
+
+    def get_logits(self, sample, net_output):
+        logits = net_output["x"]
+        logits = logits.transpose(0, 2)
+        logits = logits.reshape(-1, logits.size(-1))
+
+        return logits
+
+    def get_targets(self, sample, net_output):
+        x = net_output["x"]
+        return x.new_zeros(x.size(1) * x.size(2), dtype=torch.long)
+
+    def get_extra_losses(self, net_output):
+        pen = []
+
+        if "prob_perplexity" in net_output:
+            pen.append((net_output["num_vars"] - net_output["prob_perplexity"]) / net_output["num_vars"])
+
+        if "features_pen" in net_output:
+            pen.append(net_output["features_pen"])
+
+        return pen
+
+    def forward(
+        self,
+        input_ids: torch.Tensor,
+        type_ids: torch.Tensor = None,
+        dpe_ids: torch.Tensor = None,
+        padding_mask: torch.Tensor = None,
+        **kwargs,
+    ):
+        if self.feature_grad_mult > 0:
+            features_ret = super().forward(
+                input_ids=input_ids,
+                type_ids=type_ids,
+                dpe_ids=dpe_ids,
+                padding_mask=padding_mask,
+                encoder_only=True,
+                **kwargs,
+            )
+            features = features_ret["x"]
+            padding_mask = features_ret["padding_mask"]
+            if self.feature_grad_mult != 1.0:
+                features = GradMultiply.apply(features, self.feature_grad_mult)
+        else:
+            with torch.no_grad():
+                features = super().forward(
+                    input_ids=input_ids,
+                    type_ids=type_ids,
+                    dpe_ids=dpe_ids,
+                    padding_mask=padding_mask,
+                    encoder_only=True,
+                    **kwargs,
+                )
+                features = features_ret["x"]
+                padding_mask = features_ret["padding_mask"]
+
+        features_pen = features.float().pow(2).mean()
+
+        features = self.layer_norm(features)
+        unmasked_features = features.clone()
+
+        features = self.dropout_input(features)
+        unmasked_features = self.dropout_features(unmasked_features)
+
+        num_vars = None
+        code_ppl = None
+        prob_ppl = None
+        curr_temp = None
+
+        if padding_mask is None:
+            padding_mask = input_ids[:, :, 1].eq(0).to(features.device)
+
+        x, mask_indices = self.apply_mask(features, padding_mask, mask_indices=None)
+        if mask_indices is not None:
+            y = unmasked_features[mask_indices].view(
+                unmasked_features.size(0), -1, unmasked_features.size(-1)
+            )
+        else:
+            y = unmasked_features
+
+        x = self.event_aggregator(x, src_key_padding_mask=padding_mask)
+
+        q = self.quantizer(y, produce_targets=False)
+        y = q["x"]
+        num_vars = q["num_vars"]
+        code_ppl = q["code_perplexity"]
+        prob_ppl = q["prob_perplexity"]
+        curr_temp = q["temp"]
+
+        y = self.project_q(y)
+
+        negs, _ = self.sample_negatives(y, y.size(1))
+
+        if self.codebook_negatives > 0:
+            cb_negs = self.quantizer.sample_from_codebook(y.size(0) * y.size(1), self.codebook_negatives)
+            cb_negs = cb_negs.view(self.codebook_negatives, y.size(0), y.size(1), -1)
+            cb_negs = self.project_q(cb_negs)
+            negs = torch.cat([negs, cb_negs], dim=0)
+
+        x = x[mask_indices].view(x.size(0), -1, x.size(-1))
+
+        x = self.final_proj(x)
+        x = self.compute_preds(x, y, negs)
+
+        output = {"x": x, "features_pen": features_pen, "mask_indices": mask_indices}
+        if prob_ppl is not None:
+            output["prob_perplexity"] = prob_ppl
+            output["code_perplexity"] = code_ppl
+            output["num_vars"] = num_vars
+            output["temp"] = curr_temp
+
+        return output
+
+    def get_pretraining_parameter_names(self):
+        ret = []
+        ret.append("mask_emb")
+        ret.extend(["quantizer" + "." + x[0] for x in self.quantizer.named_parameters()])
+        ret.extend(["layer_norm" + "." + x[0] for x in self.layer_norm.named_parameters()])
+        ret.extend(["project_q" + "." + x[0] for x in self.project_q.named_parameters()])
+        ret.extend(["final_proj" + "." + x[0] for x in self.final_proj.named_parameters()])
+        return ret

genhpf/modules/__init__.py

@@ -0,0 +1,15 @@
+from .gather_layer import GatherLayer
+from .grad_multiply import GradMultiply
+from .gumbel_vector_quantizer import GumbelVectorQuantizer
+from .identity_layer import Identity
+from .layer_norm import LayerNorm
+from .positional_encoding import PositionalEncoding
+
+__all__ = [
+    "Identity",
+    "GatherLayer",
+    "GradMultiply",
+    "GumbelVectorQuantizer",
+    "LayerNorm",
+    "PositionalEncoding",
+]

genhpf/modules/gather_layer.py

@@ -0,0 +1,23 @@
+import torch
+import torch.distributed as dist
+
+import genhpf.utils.distributed_utils as dist_utils
+
+class GatherLayer(torch.autograd.Function):
+    """Gather tensors from all process, supporting backward propagation."""
+
+    @staticmethod
+    def forward(ctx, input):
+        ctx.save_for_backward(input)
+
+        group = dist_utils.get_data_parallel_group()
+        output = dist_utils.batch_all_gather(input, group=group)
+
+        return tuple(output)
+    
+    @staticmethod
+    def backward(ctx, *grads):
+        (input, ) = ctx.saved_tensors
+        grad_out = torch.zeros_like(input)
+        grad_out[:] = grads[dist.get_rank()]
+        return grad_out

genhpf/modules/grad_multiply.py

@@ -0,0 +1,12 @@
+import torch
+
+class GradMultiply(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, x, scale):
+        ctx.scale = scale
+        res = x.new(x)
+        return res
+    
+    @staticmethod
+    def backward(ctx, grad):
+        return grad * ctx.scale, None

genhpf/modules/gumbel_vector_quantizer.py

@@ -0,0 +1,204 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+class GumbelVectorQuantizer(nn.Module):
+    def __init__(
+        self,
+        dim,
+        num_vars,
+        temp,
+        groups,
+        combine_groups,
+        vq_dim,
+        time_first,
+        activation=nn.GELU(),
+        weight_proj_depth=1,
+        weight_proj_factor=1,
+    ):
+        """Vector quantization using gumbel softmax
+
+        Args:
+            dim: input dimension (channels)
+            num_vars: number of quantized vectors per group
+            temp: temperature for training. this should be a tuple of 3 elements: (start, stop, decay factor)
+            groups: number of groups for vector quantization
+            combine_groups: whether to use the vectors for all groups
+            vq_dim: dimensionality of the resulting quantized vector
+            time_first: if true, expect input in BxTxC format, otherwise in BxCxT
+            activation: what activation to use (should be a module). this is only used if weight_proj_depth is > 1
+            weight_proj_depth: number of layers (with activation in between) to project input before computing logits
+            weight_proj_factor: this is used only if weight_proj_depth is > 1. scales the inner dimensionality of
+                                projections by this factor
+        """
+        """
+        num_vars : k
+        groups: number of codebooks (divide channels by ...)
+        combine_groups: whether to map the vectors for overall centroids (over codebooks) => false
+        """
+        super().__init__()
+
+        self.groups = groups
+        self.combine_groups = combine_groups
+        self.input_dim = dim
+        self.num_vars = num_vars
+        self.time_first = time_first
+        self.num_updates = 0
+
+        assert (
+            vq_dim % groups == 0
+        ), f"dim {vq_dim} must be divisible by groups {groups} for concatenation"
+
+        var_dim = vq_dim // groups
+        num_groups = groups if not combine_groups else 1
+
+        self.vars = nn.Parameter(torch.FloatTensor(1, num_groups * num_vars, var_dim))
+        nn.init.uniform_(self.vars)
+
+        if weight_proj_depth > 1:
+
+            def block(input_dim, output_dim):
+                return nn.Sequential(nn.Linear(input_dim, output_dim), activation)
+
+            inner_dim = self.input_dim * weight_proj_factor
+            self.weight_proj = nn.Sequential(
+                *[
+                    block(self.input_dim if i == 0 else inner_dim, inner_dim)
+                    for i in range(weight_proj_depth - 1)
+                ],
+                nn.Linear(inner_dim, groups * num_vars),
+            )
+        else:
+            self.weight_proj = nn.Linear(self.input_dim, groups * num_vars)
+            nn.init.normal_(self.weight_proj.weight, mean=0, std=1)
+            nn.init.zeros_(self.weight_proj.bias)
+
+        if isinstance(temp, str):
+            import ast
+            temp = ast.literal_eval(temp)
+        assert len(temp) == 3, f"{temp}, {len(temp)}"
+
+        self.max_temp, self.min_temp, self.temp_decay = temp
+        self.curr_temp = self.max_temp
+        self.codebook_indices = None
+
+    def set_num_updates(self, num_updates):
+        self.curr_temp = max(
+            self.max_temp * self.temp_decay ** num_updates, self.min_temp
+        )
+        self.num_updates = num_updates
+
+    def get_codebook_indices(self):
+        if self.codebook_indices is None:
+            from itertools import product
+
+            p = [range(self.num_vars)] * self.groups
+            inds = list(product(*p))
+            self.codebook_indices = torch.tensor(
+                inds, dtype=torch.long, device=self.vars.device
+            ).flatten()
+
+            if not self.combine_groups:
+                self.codebook_indices = self.codebook_indices.view(
+                    self.num_vars ** self.groups, -1
+                )
+                for b in range(1, self.groups):
+                    self.codebook_indices[:, b] += self.num_vars * b
+                self.codebook_indices = self.codebook_indices.flatten()
+        return self.codebook_indices
+
+    def codebook(self):
+        indices = self.get_codebook_indices()
+        return (
+            self.vars.squeeze(0)
+            .index_select(0, indices)
+            .view(self.num_vars ** self.groups, -1)
+        )
+
+    def sample_from_codebook(self, b, n):
+        indices = self.get_codebook_indices()
+        indices = indices.view(-1, self.groups)
+        cb_size = indices.size(0)
+        assert (
+            n < cb_size
+        ), f"sample size {n} is greater than size of codebook {cb_size}"
+        sample_idx = torch.randint(low=0, high=cb_size, size=(b * n,))
+        indices = indices[sample_idx]
+
+        z = self.vars.squeeze(0).index_select(0, indices.flatten()).view(b, n, -1)
+        return z
+
+    def to_codebook_index(self, indices):
+        res = indices.new_full(indices.shape[:-1], 0)
+        for i in range(self.groups):
+            exponent = self.groups - i - 1
+            res += indices[..., i] * (self.num_vars ** exponent)
+        return res
+
+    def forward_idx(self, x):
+        res = self.forward(x, produce_targets=True)
+        return res["x"], res["targets"]
+
+    def forward(self, x, produce_targets=False):
+
+        result = {"num_vars": self.num_vars * self.groups}
+        if not self.time_first:
+            x = x.transpose(1, 2)
+        
+        bsz, tsz, fsz = x.shape
+        x = x.reshape(-1, fsz)
+        x = self.weight_proj(x)
+        x = x.view(bsz * tsz * self.groups, -1)
+
+        _, k = x.max(-1)
+
+        hard_x = (
+            x.new_zeros(*x.shape)
+            .scatter_(-1, k.view(-1, 1), 1.0)
+            .view(bsz * tsz, self.groups, -1)
+        )
+        hard_probs = torch.mean(hard_x.float(), dim=0)
+        result["code_perplexity"] = torch.exp(
+            -torch.sum(hard_probs * torch.log(hard_probs + 1e-7), dim=-1)
+        ).sum()
+
+        avg_probs = torch.softmax(
+            x.view(bsz * tsz, self.groups, -1).float(), dim=-1
+        ).mean(dim=0)
+        result["prob_perplexity"] = torch.exp(
+            -torch.sum(avg_probs * torch.log(avg_probs + 1e-7), dim=-1)
+        ).sum()
+
+        result["temp"] = self.curr_temp
+
+        if self.training:
+            x = F.gumbel_softmax(x.float(), tau=self.curr_temp, hard=True).type_as(x)
+        else:
+            x = hard_x
+
+        x = x.view(bsz * tsz, -1)
+
+        vars = self.vars
+
+        if self.combine_groups:
+            vars = vars.repeat(1, self.groups, 1)
+
+        if produce_targets:
+            result["targets"] = (
+                x.view(bsz * tsz * self.groups, -1)
+                .argmax(dim=-1)
+                .view(bsz, tsz, self.groups)
+                .detach()
+            )
+
+        x = x.unsqueeze(-1) * vars
+        x = x.view(bsz * tsz, self.groups, self.num_vars, -1)
+        x = x.sum(-2)
+        x = x.view(bsz, tsz, -1)
+
+        if not self.time_first:
+            x = x.transpose(1, 2)  # BTC -> BCT
+
+        result["x"] = x
+
+        return result

genhpf/modules/identity_layer.py

@@ -0,0 +1,8 @@
+import torch.nn as nn
+
+class Identity(nn.Module):
+    def __init__(self):
+        super(Identity, self).__init__()
+
+    def forward(self, x):
+        return x

genhpf/modules/layer_norm.py

@@ -0,0 +1,27 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+has_fused_layernorm = False
+
+def LayerNorm(normalized_shape, eps=1e-5, elementwise_affine=True, export=False):
+    if torch.jit.is_scripting():
+        export = True
+    if not export and torch.cuda.is_available() and has_fused_layernorm:
+        return FusedLayerNorm(normalized_shape, eps, elementwise_affine)
+    return torch.nn.LayerNorm(normalized_shape, eps, elementwise_affine)
+
+
+class Fp32LayerNorm(nn.LayerNorm):
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+
+    def forward(self, input):
+        output = F.layer_norm(
+            input.float(),
+            self.normalized_shape,
+            self.weight.float() if self.weight is not None else None,
+            self.bias.float() if self.bias is not None else None,
+            self.eps,
+        )
+        return output.type_as(input)

genhpf/modules/positional_encoding.py

@@ -0,0 +1,24 @@
+import math
+
+import torch
+import torch.nn as nn
+
+class PositionalEncoding(nn.Module):
+    def __init__(self, d_model, dropout, max_len):
+        super().__init__()
+        self.dropout = nn.Dropout(p=dropout)
+        position = torch.arange(max_len).unsqueeze(1)
+        div_term = torch.exp(torch.arange(0, d_model, 2) * (-math.log(10000.0) / d_model)) 
+        pe = torch.zeros(1, max_len, d_model)
+        pe[0, :, 0::2] = torch.sin(position * div_term)
+        pe[0, :, 1::2] = torch.cos(position * div_term)
+        self.register_buffer("pe", pe)
+
+    def forward(self, x): 
+        """
+        Args:
+            x: Tensor, shape [batch_size, seq_len, embedding_dim]
+        """
+
+        x = x + self.pe[:, :x.size(1)]
+        return self.dropout(x)