初始提交：3 款纯音频游戏 + 共享引擎 + 完整文档与资产管线
- 共享引擎（app/src/main）：GestureEngine 6 手势识别、AudioEngine 三层音频 （ExoPlayer BGM 交叉淡入 / 对白 ducking / SoundPool 低延迟）、Game/GameEngine 状态机接口、MainActivity 全屏封面 + TalkBack - 《声命线》storyline：末世互动广播剧，3 幕 8 抉择点 4 结局 - 《声东击西》reflex：听声辨位反应，12 波战斗 + BGM 三档 BPM 变速 - 《暗室》escape：3x3 网格密室逃脱，钥匙→铁柜→保险丝→电箱→大门谜题链 - 文档：setup（从零搭建）/ architecture（架构）/ engine-api（引擎 API）/ asset-pipeline（资产流水线）/ games/*（三款游戏设计文档） - 资产管线：Qwen3-TTS 对白 + 程序合成音效/BGM + Mureka BGM 的 gen 脚本与台账， 全部 ogg 资产入库可复现 - env.d 环境脚本：android/tts 入库；sfx/mureka 含密钥不入库，仅提交 .example.sh 模板
This commit is contained in:
SoundGame Dev
@@ -0,0 +1,504 @@
|
||||
from pytorch_lightning.loggers import WandbLogger, CometLogger
|
||||
from ..interface.aeiou import pca_point_cloud
|
||||
|
||||
import math
|
||||
import wandb
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
import os
|
||||
import typing as tp
|
||||
|
||||
def get_rank():
|
||||
"""Get rank of current process."""
|
||||
|
||||
if "SLURM_PROCID" in os.environ:
|
||||
return int(os.environ["SLURM_PROCID"])
|
||||
|
||||
if not torch.distributed.is_available() or not torch.distributed.is_initialized():
|
||||
return 0
|
||||
|
||||
return torch.distributed.get_rank()
|
||||
|
||||
class InverseLR(torch.optim.lr_scheduler._LRScheduler):
|
||||
"""Implements an inverse decay learning rate schedule with an optional exponential
|
||||
warmup. When last_epoch=-1, sets initial lr as lr.
|
||||
inv_gamma is the number of steps/epochs required for the learning rate to decay to
|
||||
(1 / 2)**power of its original value.
|
||||
Args:
|
||||
optimizer (Optimizer): Wrapped optimizer.
|
||||
inv_gamma (float): Inverse multiplicative factor of learning rate decay. Default: 1.
|
||||
power (float): Exponential factor of learning rate decay. Default: 1.
|
||||
warmup (float): Exponential warmup factor (0 <= warmup < 1, 0 to disable)
|
||||
Default: 0.
|
||||
final_lr (float): The final learning rate. Default: 0.
|
||||
last_epoch (int): The index of last epoch. Default: -1.
|
||||
"""
|
||||
|
||||
def __init__(self, optimizer, inv_gamma=1., power=1., warmup=0., final_lr=0.,
|
||||
last_epoch=-1):
|
||||
self.inv_gamma = inv_gamma
|
||||
self.power = power
|
||||
if not 0. <= warmup < 1:
|
||||
raise ValueError('Invalid value for warmup')
|
||||
self.warmup = warmup
|
||||
self.final_lr = final_lr
|
||||
super().__init__(optimizer, last_epoch)
|
||||
|
||||
def get_lr(self):
|
||||
if not self._get_lr_called_within_step:
|
||||
import warnings
|
||||
warnings.warn("To get the last learning rate computed by the scheduler, "
|
||||
"please use `get_last_lr()`.")
|
||||
|
||||
return self._get_closed_form_lr()
|
||||
|
||||
def _get_closed_form_lr(self):
|
||||
warmup = 1 - self.warmup ** (self.last_epoch + 1)
|
||||
lr_mult = (1 + self.last_epoch / self.inv_gamma) ** -self.power
|
||||
return [warmup * max(self.final_lr, base_lr * lr_mult)
|
||||
for base_lr in self.base_lrs]
|
||||
|
||||
def create_optimizer_from_config(optimizer_config, parameters):
|
||||
"""Create optimizer from config.
|
||||
|
||||
Args:
|
||||
parameters (iterable): parameters to optimize.
|
||||
optimizer_config (dict): optimizer config.
|
||||
|
||||
Returns:
|
||||
torch.optim.Optimizer: optimizer.
|
||||
"""
|
||||
|
||||
optimizer_type = optimizer_config["type"]
|
||||
|
||||
if optimizer_type == "FusedAdam":
|
||||
from deepspeed.ops.adam import FusedAdam
|
||||
optimizer = FusedAdam(parameters, **optimizer_config["config"])
|
||||
elif optimizer_type == "CAdamW":
|
||||
from stable_audio_tools.training.optims import CAdamW
|
||||
optimizer = CAdamW(parameters, **optimizer_config["config"])
|
||||
elif optimizer_type == "CLion":
|
||||
from stable_audio_tools.training.optims import CLion
|
||||
optimizer = CLion(parameters, **optimizer_config["config"])
|
||||
elif optimizer_type == "AdamW8bit":
|
||||
from bitsandbytes.optim import AdamW8bit
|
||||
optimizer = AdamW8bit(parameters, **optimizer_config["config"])
|
||||
elif optimizer_type == "MuonAdamW":
|
||||
from stable_audio_tools.training.optims import MuonAdamW
|
||||
optimizer = MuonAdamW(parameters, **optimizer_config["config"])
|
||||
else:
|
||||
optimizer_fn = getattr(torch.optim, optimizer_type)
|
||||
optimizer = optimizer_fn(parameters, **optimizer_config["config"])
|
||||
return optimizer
|
||||
|
||||
def create_scheduler_from_config(scheduler_config, optimizer):
|
||||
"""Create scheduler from config.
|
||||
|
||||
Args:
|
||||
scheduler_config (dict): scheduler config.
|
||||
optimizer (torch.optim.Optimizer): optimizer.
|
||||
|
||||
Returns:
|
||||
torch.optim.lr_scheduler._LRScheduler: scheduler.
|
||||
"""
|
||||
if scheduler_config["type"] == "InverseLR":
|
||||
scheduler_fn = InverseLR
|
||||
else:
|
||||
scheduler_fn = getattr(torch.optim.lr_scheduler, scheduler_config["type"])
|
||||
scheduler = scheduler_fn(optimizer, **scheduler_config["config"])
|
||||
return scheduler
|
||||
|
||||
def logger_project_name(logger) -> str:
|
||||
if isinstance(logger, WandbLogger):
|
||||
return logger.experiment.project
|
||||
elif isinstance(logger, CometLogger):
|
||||
return logger.name
|
||||
|
||||
def log_metric(logger, key, value, step=None):
|
||||
from pytorch_lightning.loggers import WandbLogger, CometLogger
|
||||
if isinstance(logger, WandbLogger):
|
||||
logger.experiment.log({key: value})
|
||||
elif isinstance(logger, CometLogger):
|
||||
logger.experiment.log_metrics({key: value}, step=step)
|
||||
|
||||
def log_audio(logger, key, audio_path, sample_rate, caption=None, step=None):
|
||||
if isinstance(logger, WandbLogger):
|
||||
logger.experiment.log({key: wandb.Audio(audio_path, sample_rate=sample_rate, caption=caption)})
|
||||
elif isinstance(logger, CometLogger):
|
||||
logger.experiment.log_audio(audio_path, file_name=key, sample_rate=sample_rate, step=step)
|
||||
|
||||
def log_image(logger, key, img_data, step=None):
|
||||
if isinstance(logger, WandbLogger):
|
||||
logger.experiment.log({key: wandb.Image(img_data)})
|
||||
elif isinstance(logger, CometLogger):
|
||||
logger.experiment.log_image(img_data, name=key, step=step)
|
||||
|
||||
def log_point_cloud(logger, key, tokens, caption=None):
|
||||
if isinstance(logger, WandbLogger):
|
||||
point_cloud = pca_point_cloud(tokens)
|
||||
logger.experiment.log({key: point_cloud})
|
||||
elif isinstance(logger, CometLogger):
|
||||
point_cloud = pca_point_cloud(tokens, rgb_float=True, output_type="points")
|
||||
#logger.experiment.log_points_3d(scene_name=key, points=point_cloud)
|
||||
|
||||
|
||||
def compute_per_elem_trim(conditioning, sample_rate, margin_seconds=5.0):
|
||||
"""Compute per-element trim lengths from seconds_total in conditioning dicts.
|
||||
|
||||
Returns a list of trim lengths (in audio samples) or None if no seconds_total found.
|
||||
"""
|
||||
if not any("seconds_total" in c for c in conditioning):
|
||||
return None
|
||||
margin_samples = int(margin_seconds * sample_rate)
|
||||
per_elem_trim = []
|
||||
for c in conditioning:
|
||||
if "seconds_total" in c:
|
||||
per_elem_trim.append(int(c["seconds_total"] * sample_rate) + margin_samples)
|
||||
else:
|
||||
per_elem_trim.append(None)
|
||||
return per_elem_trim
|
||||
|
||||
def trim_and_concat(x, per_elem_trim):
|
||||
"""Per-element trim and concatenate along time axis.
|
||||
|
||||
Trims each batch element to its own length (from seconds_total + margin),
|
||||
removing trailing padding silence before concatenation.
|
||||
|
||||
Args:
|
||||
x: (b, d, n) tensor or list of (d, n) tensors
|
||||
per_elem_trim: list of trim lengths per element, or None for no trimming
|
||||
"""
|
||||
items = [x[i] for i in range(x.shape[0])] if isinstance(x, torch.Tensor) and x.dim() == 3 else x
|
||||
if per_elem_trim is None:
|
||||
return torch.cat(items, dim=-1)
|
||||
parts = []
|
||||
for i, elem in enumerate(items):
|
||||
if per_elem_trim[i] is not None:
|
||||
parts.append(elem[..., :min(per_elem_trim[i], elem.shape[-1])])
|
||||
else:
|
||||
parts.append(elem)
|
||||
return torch.cat(parts, dim=-1)
|
||||
|
||||
|
||||
class StaggeredLogger:
|
||||
"""Accumulates log values over N steps and flushes averaged metrics.
|
||||
|
||||
Avoids per-step CUDA syncs (.item()) and prevents PL from caching
|
||||
stale values between flush intervals by logging directly to the
|
||||
experiment logger.
|
||||
|
||||
Usage in a PL module::
|
||||
|
||||
# In __init__:
|
||||
self._staggered_logger = StaggeredLogger(every_n_steps=10)
|
||||
|
||||
# At end of training_step (pass detached tensors or Python scalars):
|
||||
self._staggered_logger.log(log_dict, self)
|
||||
"""
|
||||
|
||||
def __init__(self, every_n_steps: int = 10):
|
||||
self.every_n_steps = every_n_steps
|
||||
self._accum: dict = {}
|
||||
self._counts: dict = {}
|
||||
self._steps: int = 0
|
||||
|
||||
def log(self, log_dict: dict, module) -> None:
|
||||
"""Accumulate *log_dict* and flush every *every_n_steps* calls.
|
||||
|
||||
NaN/inf values are handled at flush time: if an accumulated sum
|
||||
is non-finite (due to a NaN step poisoning the window), that
|
||||
metric is simply omitted from the flush. This avoids the
|
||||
expensive per-step CUDA syncs that an eager ``isfinite`` check
|
||||
would require.
|
||||
|
||||
Args:
|
||||
log_dict: Mapping of metric names to detached GPU tensors or
|
||||
Python scalars. Tensors are kept on-device until flush.
|
||||
module: The ``pl.LightningModule`` instance (used for
|
||||
``log_dict``, ``logger``, and ``global_step``).
|
||||
"""
|
||||
for k, v in log_dict.items():
|
||||
val = v.detach() if torch.is_tensor(v) else v
|
||||
if k not in self._accum:
|
||||
self._accum[k] = val
|
||||
self._counts[k] = 1
|
||||
else:
|
||||
self._accum[k] = self._accum[k] + val
|
||||
self._counts[k] += 1
|
||||
self._steps += 1
|
||||
|
||||
if self._steps >= self.every_n_steps:
|
||||
flushed = {}
|
||||
for k, v in self._accum.items():
|
||||
c = self._counts[k]
|
||||
if c > 0:
|
||||
val = v.item() / c if torch.is_tensor(v) else v / c
|
||||
# Skip non-finite values (NaN from a bad step poisoned this window)
|
||||
if math.isfinite(val):
|
||||
flushed[k] = val
|
||||
# Update progress bar only (logger=False avoids PL caching stale values)
|
||||
if flushed:
|
||||
module.log_dict(flushed, prog_bar=True, on_step=True, logger=False)
|
||||
# Log directly to experiment (bypasses PL metric caching)
|
||||
if module.logger is not None:
|
||||
module.logger.log_metrics(flushed, step=module.global_step)
|
||||
self._accum = {}
|
||||
self._counts = {}
|
||||
self._steps = 0
|
||||
|
||||
|
||||
def resize_padding_mask(padding_mask: torch.Tensor, target_length: int) -> torch.Tensor:
|
||||
"""Resize a padding mask to target_length using ceiling-based length scaling.
|
||||
|
||||
Unlike F.interpolate(mode="nearest"), this ensures any target position
|
||||
that partially overlaps valid audio is marked valid (rounds up).
|
||||
|
||||
Args:
|
||||
padding_mask: (B, T) or (T,) tensor where 1/True = valid, 0/False = padding
|
||||
target_length: desired output length
|
||||
Returns:
|
||||
Resized boolean mask with same ndim as input
|
||||
"""
|
||||
if padding_mask.ndim == 1:
|
||||
valid_length = padding_mask.sum()
|
||||
source_length = padding_mask.shape[0]
|
||||
valid_target_length = torch.ceil(
|
||||
valid_length.float() * target_length / source_length
|
||||
).long().clamp(max=target_length)
|
||||
positions = torch.arange(target_length, device=padding_mask.device)
|
||||
return positions < valid_target_length
|
||||
else:
|
||||
valid_lengths = padding_mask.sum(dim=-1) # (B,)
|
||||
source_length = padding_mask.shape[-1]
|
||||
valid_target_lengths = torch.ceil(
|
||||
valid_lengths.float() * target_length / source_length
|
||||
).long().clamp(max=target_length)
|
||||
positions = torch.arange(target_length, device=padding_mask.device).unsqueeze(0)
|
||||
return positions < valid_target_lengths.unsqueeze(1)
|
||||
|
||||
|
||||
def create_augmented_padding_mask(
|
||||
padding_mask: torch.Tensor,
|
||||
silence_extension_scale_seconds: float = 0.0,
|
||||
sample_rate: int = 44100,
|
||||
downsampling_ratio: int = 1,
|
||||
) -> torch.Tensor:
|
||||
"""
|
||||
Augment padding mask by randomly extending the valid region to include silence.
|
||||
|
||||
This helps the model learn to handle silence at the end of audio by randomly
|
||||
including some of the padded (silent) region as valid.
|
||||
|
||||
Uses an exponential distribution so most extensions are small (within ~scale seconds)
|
||||
but can occasionally extend to the full sequence length.
|
||||
|
||||
Args:
|
||||
padding_mask: Boolean tensor of shape (batch, seq_len) where True = valid audio
|
||||
silence_extension_scale_seconds: Scale (mean) of exponential distribution in seconds.
|
||||
~63% of samples within 1x scale, ~86% within 2x, ~95% within 3x.
|
||||
sample_rate: Audio sample rate (used to convert seconds to tokens)
|
||||
downsampling_ratio: Downsampling ratio from pretransform (to convert samples to latent tokens)
|
||||
|
||||
Returns:
|
||||
Augmented padding mask of shape (batch, seq_len) where True = valid region (audio + extended silence)
|
||||
"""
|
||||
if silence_extension_scale_seconds <= 0:
|
||||
return padding_mask
|
||||
|
||||
batch_size, seq_len = padding_mask.shape
|
||||
device = padding_mask.device
|
||||
|
||||
# Find the last valid position for each batch element
|
||||
valid_lengths = padding_mask.sum(dim=-1) # (batch,)
|
||||
|
||||
# Calculate scale in tokens for exponential distribution
|
||||
scale_tokens = silence_extension_scale_seconds * sample_rate / downsampling_ratio
|
||||
|
||||
# Sample from exponential distribution (biased toward smaller values)
|
||||
# Exponential: ~63% within 1x scale, ~86% within 2x, ~95% within 3x
|
||||
random_extensions = torch.empty(batch_size, device=device).exponential_(lambd=1.0/scale_tokens)
|
||||
random_extensions = random_extensions.long()
|
||||
|
||||
# Extend valid lengths, clamping to sequence length (allows extension to full padding)
|
||||
augmented_lengths = torch.clamp(valid_lengths + random_extensions, max=seq_len)
|
||||
|
||||
# Create augmented mask: True where position < augmented_length
|
||||
positions = torch.arange(seq_len, device=device).unsqueeze(0) # (1, seq_len)
|
||||
augmented_mask = positions < augmented_lengths.unsqueeze(1) # (batch, seq_len)
|
||||
|
||||
return augmented_mask
|
||||
|
||||
|
||||
def compute_normalized_mse(
|
||||
pred: torch.Tensor,
|
||||
target: torch.Tensor,
|
||||
loss_mask: torch.Tensor,
|
||||
loss_normalization: str = "none",
|
||||
loss_norm_eps: float = 1e-6,
|
||||
) -> torch.Tensor:
|
||||
"""
|
||||
Compute MSE normalized by detached target magnitude.
|
||||
|
||||
Prevents high-magnitude latents from dominating the loss,
|
||||
ensuring quiet passages receive equal learning signal.
|
||||
|
||||
Normalization is only applied to valid (non-padded) regions.
|
||||
Padded regions retain raw MSE for compatibility with mask_loss_weight.
|
||||
|
||||
Args:
|
||||
pred: [B, C, T] predicted tensor
|
||||
target: [B, C, T] target tensor
|
||||
loss_mask: [B, T] boolean, True = valid (non-padded), or None
|
||||
loss_normalization: One of "none", "timestep", "sample", "sample_channel"
|
||||
loss_norm_eps: Epsilon for numerical stability
|
||||
|
||||
Returns:
|
||||
Per-element MSE [B, C, T], normalized for signal, raw for padding
|
||||
"""
|
||||
mse = (pred - target) ** 2 # [B, C, T]
|
||||
|
||||
if loss_normalization == "none":
|
||||
return mse
|
||||
|
||||
if loss_mask is None:
|
||||
mask_expanded = torch.ones(pred.shape[0], 1, pred.shape[2], device=pred.device, dtype=torch.bool)
|
||||
else:
|
||||
mask_expanded = loss_mask.unsqueeze(1) # [B, 1, T]
|
||||
|
||||
with torch.no_grad():
|
||||
if loss_normalization == "timestep":
|
||||
# Var across channels at each timestep [B, 1, T]
|
||||
mag_sq = torch.mean((target - torch.mean(target, dim=1, keepdim=True)) ** 2, dim=1, keepdim=True) + loss_norm_eps
|
||||
else:
|
||||
# For sample/sample_channel modes, exclude padding via NaN masking
|
||||
masked_targets = torch.where(mask_expanded, target, float('nan'))
|
||||
|
||||
if loss_normalization == "sample":
|
||||
sample_mean = torch.nanmean(masked_targets, dim=(1, 2), keepdim=True)
|
||||
mag_sq = torch.nanmean((masked_targets - sample_mean) ** 2, dim=(1, 2), keepdim=True) + loss_norm_eps
|
||||
elif loss_normalization == "sample_channel":
|
||||
channel_mean = torch.nanmean(masked_targets, dim=2, keepdim=True)
|
||||
mag_sq = torch.nanmean((masked_targets - channel_mean) ** 2, dim=2, keepdim=True) + loss_norm_eps
|
||||
else:
|
||||
raise ValueError(f"Unknown loss normalization mode: {loss_normalization}")
|
||||
|
||||
# Handle edge case where entire sample is padding (all NaN -> NaN variance)
|
||||
mag_sq = torch.where(torch.isnan(mag_sq), torch.ones_like(mag_sq), mag_sq)
|
||||
|
||||
# Only normalize signal regions; keep padding MSE raw for mask_loss_weight compatibility
|
||||
normalized_mse = mse / mag_sq
|
||||
return torch.where(mask_expanded, normalized_mse, mse)
|
||||
|
||||
|
||||
def compute_masked_loss(
|
||||
loss_full: torch.Tensor,
|
||||
loss_mask: torch.Tensor,
|
||||
mask_padding_attention: bool,
|
||||
mask_loss_weight: float = 0.0,
|
||||
) -> tuple:
|
||||
"""
|
||||
Compute loss with separate signal and padding contributions.
|
||||
|
||||
Args:
|
||||
loss_full: Full loss tensor of shape (B, C, T)
|
||||
loss_mask: Boolean mask of shape (B, T) where True = signal, False = padding
|
||||
mask_padding_attention: If True, only compute loss on signal region
|
||||
mask_loss_weight: Weight for padding loss when mask_padding_attention is False
|
||||
|
||||
Returns:
|
||||
Tuple of (loss, signal_mean, padding_mean) where:
|
||||
- loss: The final scalar loss
|
||||
- signal_mean: Mean loss on signal region (for logging)
|
||||
- padding_mean: Mean loss on padding region (for logging)
|
||||
"""
|
||||
# Compute separate signal and padding losses
|
||||
signal = torch.where(loss_mask.unsqueeze(1), loss_full, 0.0)
|
||||
signal_sum = signal.sum(dim=(1, 2))
|
||||
n_channels = loss_full.shape[1]
|
||||
signal_count = loss_mask.sum(dim=1) * n_channels
|
||||
|
||||
padding = torch.where(~loss_mask.unsqueeze(1), loss_full, 0.0)
|
||||
padding_sum = padding.sum(dim=(1, 2))
|
||||
padding_count = (~loss_mask).sum(dim=1) * n_channels
|
||||
|
||||
if mask_padding_attention:
|
||||
# When attention masking is on, only compute loss on signal region
|
||||
# (padding saw no attention, so loss there is meaningless)
|
||||
per_sample_loss = signal_sum / (signal_count + 1e-8)
|
||||
loss = per_sample_loss.mean()
|
||||
else:
|
||||
# Weighted loss: signal + mask_loss_weight * padding
|
||||
w = mask_loss_weight
|
||||
denom = signal_count + w * padding_count + 1e-8
|
||||
per_sample_loss = (signal_sum + w * padding_sum) / denom
|
||||
loss = per_sample_loss.mean()
|
||||
|
||||
# Compute separate means for logging
|
||||
signal_mean = signal.sum() / (signal_count.sum() + 1e-8)
|
||||
padding_mean = padding.sum() / (padding_count.sum() + 1e-8)
|
||||
|
||||
return loss, signal_mean.detach(), padding_mean.detach()
|
||||
|
||||
|
||||
def masked_mean(
|
||||
values: torch.Tensor,
|
||||
mask: torch.Tensor,
|
||||
) -> torch.Tensor:
|
||||
"""
|
||||
Compute mean over masked (valid) positions only.
|
||||
|
||||
Args:
|
||||
values: Tensor of shape (B, C, T) or (B, T)
|
||||
mask: Boolean mask of shape (B, T) where True = valid, False = padding
|
||||
|
||||
Returns:
|
||||
Scalar mean over valid positions only
|
||||
"""
|
||||
if values.ndim == 2:
|
||||
# (B, T) -> add channel dim
|
||||
values = values.unsqueeze(1)
|
||||
|
||||
# Expand mask to match values shape: (B, T) -> (B, 1, T)
|
||||
mask_expanded = mask.unsqueeze(1)
|
||||
|
||||
# Zero out padding positions
|
||||
masked_values = torch.where(mask_expanded, values, torch.zeros_like(values))
|
||||
|
||||
# Count valid positions (per sample, times channels)
|
||||
n_channels = values.shape[1]
|
||||
valid_count = mask.sum(dim=1) * n_channels # (B,)
|
||||
|
||||
# Sum over all dimensions and divide by valid count
|
||||
per_sample_sum = masked_values.sum(dim=(1, 2)) # (B,)
|
||||
per_sample_mean = per_sample_sum / (valid_count + 1e-8)
|
||||
|
||||
return per_sample_mean.mean()
|
||||
|
||||
|
||||
def masked_sum(
|
||||
values: torch.Tensor,
|
||||
mask: torch.Tensor,
|
||||
dim: tp.List[int],
|
||||
) -> torch.Tensor:
|
||||
"""
|
||||
Compute sum over masked (valid) positions only.
|
||||
|
||||
Args:
|
||||
values: Tensor of shape (B, C, T) or (B, T)
|
||||
mask: Boolean mask of shape (B, T) where True = valid, False = padding
|
||||
dim: Dimensions to sum over (should include the time dimension)
|
||||
|
||||
Returns:
|
||||
Sum over valid positions, reduced over specified dims
|
||||
"""
|
||||
if values.ndim == 2:
|
||||
# (B, T) -> add channel dim for consistency
|
||||
mask_expanded = mask
|
||||
else:
|
||||
# (B, C, T) -> expand mask
|
||||
mask_expanded = mask.unsqueeze(1)
|
||||
|
||||
# Zero out padding positions
|
||||
masked_values = torch.where(mask_expanded, values, torch.zeros_like(values))
|
||||
|
||||
return masked_values.sum(dim=tuple(dim))
|
||||
Reference in New Issue
Block a user