# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.
#
# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.
import logging
import typing as tp
from functools import partial
import os
from pathlib import Path
import flashy
from omegaconf import DictConfig
import multiprocessing
import numpy as np
import torch
import torch.nn as nn
from . import base, builders
from ..models.builders import get_watermark_model
from ..modules.watermark import pad, mix
from ..metrics.miou import calculate_miou
from ..metrics.pesq import PesqMetric
from ..utils import checkpoint
from ..utils.audio_effects import (
compress_with_encodec,
get_audio_effects,
select_audio_effects,
)
from ..utils.samples.manager import SampleManager
from ..data.audio import save_spectrograms
from ..utils.utils import get_pool_executor
from torchmetrics.audio.snr import ScaleInvariantSignalNoiseRatio
from torchmetrics.audio.stoi import ShortTimeObjectiveIntelligibility
if tp.TYPE_CHECKING:
from ..models.watermark import WMModel
def get_encodec_audio_effect(encodec_cfg: DictConfig, sr: int) -> tp.Dict:
"""
Construct encodec-based compression data agumentation. This method is
is put here instead of in `audiocraft.utils.audio_effects` because
it depends on the package `audiocraft.solvers`, which is one layer
higher than `audiocraft.utils`, so we avoid the circle dependency
from any solvers using `audiocraft.utils.audio_effects` to do the
augmentation
"""
from ..solvers.compression import CompressionSolver
codec_model = CompressionSolver.model_from_checkpoint(encodec_cfg.ckpt)
codec_model.train()
return {
f"encodec_nq={n_q}": partial(
compress_with_encodec,
model=codec_model,
n_q=n_q,
sample_rate=sr,
)
for n_q in encodec_cfg.n_qs
}
def random_message(nbits: int, batch_size: int) -> torch.Tensor:
"""Return random message as 0/1 tensor."""
if nbits == 0:
return torch.tensor([])
return torch.randint(0, 2, (batch_size, nbits))
class WatermarkSolver(base.StandardSolver):
"""Solver for different watermarking models"""
def __init__(self, cfg: DictConfig):
super().__init__(cfg)
self.rng: torch.Generator # set at each epoch
self.model: WMModel
if hasattr(cfg, "fsdp"):
assert not getattr(
cfg.fsdp, "use", False
), "FSDP not supported by WatermarkSolver."
self._init_losses()
self._init_augmentations()
self.balancer = builders.get_balancer(self.loss_weights, self.cfg.balancer)
self.path_specs = os.path.join(self.folder, "spectrograms")
os.makedirs(self.path_specs, exist_ok=True)
def _init_losses(self):
assert hasattr(self.cfg, "losses") and isinstance(
self.cfg.losses, (DictConfig, tp.Mapping)
), "WatermarkSolver must declare training losses in the config"
self.adv_losses = builders.get_adversarial_losses(self.cfg) # noqa
self.register_stateful("adv_losses")
self.aux_losses = nn.ModuleDict() # noqa
self.info_losses = nn.ModuleDict() # noqa
self.wm_losses = nn.ModuleDict() # noqa
loss_weights = {}
for loss_name, weight in self.cfg.losses.items():
# explicitly skip this loss calculation by setting a -1 as weight
# if weight == 0 it will be calculated but kept as info
if weight == -1:
continue
if loss_name in ["adv", "feat"]:
for adv_name, _ in self.adv_losses.items():
loss_weights[f"{loss_name}_{adv_name}"] = weight
elif weight > 0:
if loss_name[:3] == "wm_":
self.wm_losses[loss_name] = builders.get_loss(
loss_name, self.cfg
).to(self.device)
loss_weights[loss_name] = weight
else:
self.aux_losses[loss_name] = builders.get_loss(
loss_name, self.cfg
).to(self.device)
loss_weights[loss_name] = weight
else:
self.info_losses[loss_name] = builders.get_loss(loss_name, self.cfg).to(
self.device
)
self.loss_weights = loss_weights # noqa
def _init_augmentations(self):
if not hasattr(self.cfg, "aug_weights") or not hasattr(
self.cfg, "audio_effects"
):
return
aug_weights = {}
cfg_audio_effects = dict(self.cfg.audio_effects)
# Handle `encodec` augmentation separately as this requires loading a
# CompressionSolver checkpoint
encodec_cfg = cfg_audio_effects.pop("encodec", None)
if encodec_cfg:
encodec_effects = get_encodec_audio_effect(
encodec_cfg, self.cfg.sample_rate
)
for aug_name in encodec_effects.keys():
aug_weights[aug_name] = getattr(self.cfg.aug_weights, "encodec", -1)
else:
encodec_effects = {}
other_effects = get_audio_effects(self.cfg) # noqa
for name in other_effects.keys():
aug_weights[name] = self.cfg.aug_weights.get(name, -1)
self.aug_weights = aug_weights # noqa
self.augmentations = {**encodec_effects, **other_effects} # noqa
@property
def best_metric_name(self) -> tp.Optional[str]:
# best model is the last for the watermark model for now
return None
def build_model(self):
"""Instantiate model and optimizer."""
# Model and optimizer
self.model = get_watermark_model(self.cfg)
# Need two optimizers ?
self.optimizer = builders.get_optimizer(self.model.parameters(), self.cfg.optim)
self.register_stateful("model", "optimizer")
self.register_best_state("model")
self.register_ema("model")
def build_dataloaders(self):
"""Instantiate audio dataloaders for each stage."""
self.dataloaders = builders.get_audio_datasets(self.cfg)
def show(self):
"""Show the Watermark model and employed adversarial loss."""
self.log_model_summary(self.model)
self.logger.info("Sould print losses here:")
def crop(
self, signal: torch.Tensor, watermark: torch.Tensor
) -> tp.Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
"""
Applies a transformation to modify the watermarked signal to train localization.
It can be one of the following:
- zero padding: add zeros at the begining and the end of the signal
- crop: crop the watermark apply a watermark only on some parts of the signal
- shuffle: replace some part of the audio with other non watermarked parts
from the batch
In every cases the function returns a mask that contains indicates the parts that are or
not watermarked
Args:
watermark (torch.Tensor): The watermark to apply on the signal.
signal (torch.Tensor): clean signal
Returns:
watermark (torch.Tensor): modified watermark
signal (torch.Tensor): modified signal
mask (torch.Tensor): mask indicating which portion is still watermarked
"""
assert (
self.cfg.crop.prob + self.cfg.crop.shuffle_prob + self.cfg.crop.pad_prob
<= 1
), f"The sum of the probabilities {self.cfg.crop.prob=} {self.cfg.crop.shuffle_prob=} \
{self.cfg.crop.pad_prob=} should be less than 1"
mask = torch.ones_like(watermark)
p = torch.rand(1)
if p < self.cfg.crop.pad_prob: # Pad with some probability
start = int(torch.rand(1) * 0.33 * watermark.size(-1))
finish = int((0.66 + torch.rand(1) * 0.33) * watermark.size(-1))
mask[:, :, :start] = 0
mask[:, :, finish:] = 0
if torch.rand(1) > 0.5:
mask = 1 - mask
signal *= mask # pad signal
elif (
p < self.cfg.crop.prob + self.cfg.crop.pad_prob + self.cfg.crop.shuffle_prob
):
# Define a mask, then crop or shuffle
mask_size = round(watermark.shape[-1] * self.cfg.crop.size)
n_windows = int(
torch.randint(1, self.cfg.crop.max_n_windows + 1, (1,)).item()
)
window_size = int(mask_size / n_windows)
for _ in range(n_windows): # Create multiple windows in the mask
mask_start = torch.randint(0, watermark.shape[-1] - window_size, (1,))
mask[:, :, mask_start: mask_start + window_size] = (
0 # Apply window to mask
)
# inverse the mask half the time
if torch.rand(1) > 0.5:
mask = 1 - mask
if p < self.cfg.crop.pad_prob + self.cfg.crop.shuffle_prob: # shuffle
# shuffle
signal_cloned = signal.clone().detach() # detach to be sure
shuffle_idx = torch.randint(0, signal.size(0), (signal.size(0),))
signal = signal * mask + signal_cloned[shuffle_idx] * (
1 - mask
) # shuffle signal where not wm
watermark *= mask # Apply mask to the watermark
return signal, watermark, mask
def run_step(self, idx: int, batch: torch.Tensor, metrics: dict):
"""Perform one training or valid step on a given batch."""
x = batch.to(self.device)
y = x.clone()
nbits = getattr(self.model, "nbits")
message = random_message(nbits, y.shape[0]).to(self.device)
watermark = self.model.get_watermark(x, message=message)
y, watermark, mask = self.crop(y, watermark)
y_wm = y + watermark
if (
self.cfg.losses.adv != 0 or self.cfg.losses.feat != 0
) and self.is_training: # train quality adv
d_losses: dict = {}
if (
len(self.adv_losses) > 0
and torch.rand(1, generator=self.rng).item()
<= 1 / self.cfg.adversarial.every
):
for adv_name, adversary in self.adv_losses.items():
disc_loss = adversary.train_adv(y_wm, y)
d_losses[f"d_{adv_name}"] = disc_loss
metrics["d_loss"] = torch.sum(torch.stack(list(d_losses.values())))
metrics.update(d_losses)
balanced_losses: dict = {}
other_losses: dict = {}
# adversarial losses
if self.cfg.losses.adv != 0 or self.cfg.losses.feat != 0:
for adv_name, adversary in self.adv_losses.items():
adv_loss, feat_loss = adversary(y_wm, y)
balanced_losses[f"adv_{adv_name}"] = adv_loss
balanced_losses[f"feat_{adv_name}"] = feat_loss
# auxiliary losses on quality/similarity
for loss_name, criterion in self.aux_losses.items():
loss = criterion(y_wm, y)
balanced_losses[loss_name] = loss
# apply augmentations
mode = "all" if self.cfg.select_aug_mode == "all" else "weighted"
selected_augs = select_audio_effects(
self.augmentations,
self.aug_weights,
mode=mode,
max_length=self.cfg.n_max_aug,
)
N_augs = len(selected_augs)
for (
augmentation_name,
augmentation_method,
) in selected_augs.items():
# concatenate to use the augmentation function only once
y_y_wm = torch.cat([y, y_wm], dim=0)
aug_cat, mask_aug = augmentation_method(y_y_wm, mask=mask)
aug_y = aug_cat[: y.size(0)]
aug_y_wm = aug_cat[y.size(0):]
positive = self.model.detect_watermark(aug_y_wm)
negative = self.model.detect_watermark(aug_y)
for loss_name, criterion in self.wm_losses.items():
loss = criterion(positive, negative, mask_aug, message)
other_losses[f"{loss_name}_{augmentation_name}"] = loss
# weighted losses
metrics.update(balanced_losses)
metrics.update(other_losses)
if self.is_training: # something is weird about the loss balancer not
other_loss = torch.tensor(0.0, device=self.device)
for name, o_loss in other_losses.items():
if "wm_detection" in name:
# here we include the detection losses for augmentation
other_loss += (self.loss_weights["wm_detection"] / N_augs) * o_loss
elif "wm_mb" in name:
other_loss += (self.loss_weights["wm_mb"] / N_augs) * o_loss
else:
other_loss += self.loss_weights[name] * o_loss
if other_loss.requires_grad:
other_loss.backward(retain_graph=True)
ratio1 = sum(
p.grad.data.norm(p=2).pow(2)
for p in self.model.parameters()
if p.grad is not None
)
assert isinstance(ratio1, torch.Tensor)
metrics["ratio1"] = ratio1.sqrt()
# balancer losses backward, returns effective training loss
# with effective weights at the current batch.
metrics["g_loss"] = self.balancer.backward(balanced_losses, y_wm)
# add metrics corresponding to weight ratios
metrics.update(self.balancer.metrics)
ratio2 = sum(
p.grad.data.norm(p=2).pow(2)
for p in self.model.parameters()
if p.grad is not None
)
assert isinstance(ratio2, torch.Tensor)
metrics["ratio2"] = ratio2.sqrt()
# optim
flashy.distrib.sync_model(self.model)
if self.cfg.optim.max_norm:
torch.nn.utils.clip_grad_norm_(
self.model.parameters(), self.cfg.optim.max_norm
)
self.optimizer.step()
self.optimizer.zero_grad()
# informative losses only
info_losses: dict = {}
with torch.no_grad():
for loss_name, criterion in self.info_losses.items():
loss = criterion(y_wm, y)
info_losses[loss_name] = loss
# pesq
metrics["pesq"] = tensor_pesq(y_wm, y, sr=self.cfg.sample_rate)
# max allocated memory
metrics["max_mem"] = torch.cuda.max_memory_allocated() / 1e9
metrics.update(info_losses)
if self.cfg.losses.adv != 0 or self.cfg.losses.feat != 0:
# aggregated GAN losses: this is useful to report adv and feat across different adversarial loss setups
adv_losses = [
loss
for loss_name, loss in metrics.items()
if loss_name.startswith("adv")
]
if len(adv_losses) > 0:
metrics["adv"] = torch.sum(torch.stack(adv_losses))
feat_losses = [
loss
for loss_name, loss in metrics.items()
if loss_name.startswith("feat")
]
if len(feat_losses) > 0:
metrics["feat"] = torch.sum(torch.stack(feat_losses))
return metrics
def run_epoch(self):
# reset random seed at the beginning of the epoch
self.rng = torch.Generator()
self.rng.manual_seed(1234 + self.epoch)
# run epoch
super().run_epoch()
def evaluate(self) -> dict:
"""Evaluate stage. Runs audio reconstruction evaluation."""
self.model.eval()
evaluate_stage_name = str(self.current_stage)
loader = self.dataloaders["evaluate"]
updates = len(loader)
lp = self.log_progress(
f"{evaluate_stage_name} inference",
loader,
total=updates,
updates=self.log_updates,
)
average = flashy.averager()
pendings = []
ctx = multiprocessing.get_context("spawn")
with get_pool_executor(self.cfg.evaluate.num_workers, mp_context=ctx) as pool:
for batch in lp:
x = batch.to(self.device)
with torch.no_grad():
message = random_message(self.model.nbits, x.shape[0])
watermark = self.model.get_watermark(x, message)
x_wm = x + watermark
y_pred = x_wm.cpu()
y = batch.cpu() # should already be on CPU but just in case
pendings.append(
pool.submit(
evaluate_audio_watermark,
y_pred,
y,
self.cfg,
)
)
# evaluate augmentations
# evaluation is run on all the augmentations
for (
augmentation_name,
augmentation_method,
) in self.augmentations.items():
# if (
# "mp3" in augmentation_name
# and idx >= 8
# and self.cfg.evaluate.every <= 2
# ):
# # When evaluating often do not compute mp3 on the full eval dset to make things faster
# continue
with torch.no_grad():
aug_positive = self.model.detect_watermark(
augmentation_method(x_wm)
)
aug_negative = self.model.detect_watermark(
augmentation_method(x)
)
pendings.append(
pool.submit(
evaluate_augmentations,
aug_positive.cpu(),
aug_negative.cpu(),
augmentation_name,
message.cpu(),
)
)
# end eval of augmentations
# evaluate localization cropping
for window_size in np.linspace(0.1, 0.9, 9):
mixed, true_predictions = mix(x, x_wm, window_size=window_size)
model_predictions = self.model.detect_watermark(mixed)
pendings.append(
pool.submit(
evaluate_localizations,
model_predictions.cpu(),
true_predictions.cpu(),
f"crop_{window_size:0.1f}",
)
)
mixed, true_predictions = mix(
x, x_wm, window_size=window_size, shuffle=True
)
model_predictions = self.model.detect_watermark(mixed)
pendings.append(
pool.submit(
evaluate_localizations,
model_predictions.cpu(),
true_predictions.cpu(),
f"shuffle_{window_size:0.1f}",
)
)
# evaluate localization padding
mixed, true_predictions = pad(x_wm)
model_predictions = self.model.detect_watermark(mixed)
pendings.append(
pool.submit(
evaluate_localizations,
model_predictions.cpu(),
true_predictions.cpu(),
"padding",
)
)
mixed, true_predictions = pad(x_wm, central=True)
model_predictions = self.model.detect_watermark(mixed)
pendings.append(
pool.submit(
evaluate_localizations,
model_predictions.cpu(),
true_predictions.cpu(),
"central_padding",
)
)
# end of evaluate localization
metrics_lp = self.log_progress(
f"{evaluate_stage_name} metrics", pendings, updates=self.log_updates
)
for pending in metrics_lp:
metrics = pending.result()
metrics = average(metrics)
metrics = flashy.distrib.average_metrics(metrics, len(loader))
if self.cfg.select_aug_mode == "use_eval_acc":
# Adjust augmentation weights based on evaluation loss.
# Higher accuracy results in lower probability of selecting this augmentation.
for name in self.augmentations.keys():
if (
self.aug_weights[name] != -1
): # keep weight to -1 for unwanted augmentations
# set to 0.05 to ensure that an augmentation is never completely removed during a full epoch.
self.aug_weights[name] = max(1 - metrics[f"aug_{name}_acc"], 0.05)
return metrics
def generate(self):
"""Generate stage."""
self.model.eval()
sample_manager = SampleManager(self.xp, map_reference_to_sample_id=True)
generate_stage_name = str(self.current_stage)
loader = self.dataloaders["generate"]
updates = len(loader)
lp = self.log_progress(
generate_stage_name, loader, total=updates, updates=self.log_updates
)
path_dir = os.path.join(self.path_specs, f"epoch={self.epoch}")
os.makedirs(path_dir, exist_ok=True)
first_batch = True
for batch in lp:
reference, _ = batch
reference = reference.to(self.device)
with torch.no_grad():
message = random_message(self.model.nbits, reference.shape[0])
watermark = self.model.get_watermark(reference, message)
x_wm = reference + watermark
reference = reference.cpu()
sample_manager.add_samples(
x_wm.cpu(), self.epoch, ground_truth_wavs=reference
)
if first_batch and flashy.distrib.is_rank_zero():
for i in range(reference.size(0)):
ys = [
reference.cpu()[i].squeeze(0).numpy(),
x_wm.cpu()[i].squeeze(0).numpy(),
watermark.cpu()[i].squeeze(0).numpy(),
]
path = os.path.join(path_dir, f"spec_{i}.pdf")
save_spectrograms(
ys,
names=["Ground Truth", "Audio Watermarked", "Watermark"],
sr=self.cfg.sample_rate,
path=path,
)
first_batch = False
flashy.distrib.barrier()
def load_from_pretrained(self, name: str) -> dict:
raise ValueError("No pretrained model")
@staticmethod
def model_from_checkpoint(
checkpoint_path: tp.Union[Path, str],
device: tp.Union[torch.device, str] = "cpu",
) -> "WMModel":
"""Instantiate a WatermarkModel from a given checkpoint path or dora sig.
Args:
checkpoint_path (Path or str): Path to checkpoint or dora sig from where the checkpoint is resolved.
device (torch.device or str): Device on which the model is loaded.
"""
checkpoint_path = str(checkpoint_path)
logger = logging.getLogger(__name__)
logger.info(f"Loading WatermarkModel from checkpoint: {checkpoint_path}")
_checkpoint_path = checkpoint.resolve_checkpoint_path(
checkpoint_path, use_fsdp=False
)
assert (
_checkpoint_path is not None
), f"Could not resolve WatermarkModel checkpoint path: {checkpoint_path}"
state = checkpoint.load_checkpoint(_checkpoint_path)
assert (
state is not None and "xp.cfg" in state
), f"Could not load WatermarkModel from ckpt: {checkpoint_path}"
cfg = state["xp.cfg"]
cfg.device = device
watermarking_model = get_watermark_model(cfg).to(device)
assert "best_state" in state and state["best_state"] != {}
assert (
"exported" not in state
), "When loading an exported checkpoint, use the //pretrained/ prefix."
watermarking_model.load_state_dict(state["best_state"]["model"])
watermarking_model.eval()
logger.info("Watermarking model loaded!")
return watermarking_model
def evaluate_localizations(predictions, true_predictions, name):
metrics = {}
# predictions are output of the detector shape [bsz, 2, frames]
# true_predictions is output of the mix method shape [bsz, 2, frames]
metrics[f"localization_acc_{name}"] = (
((predictions[:, 1, :] > 0.5) == true_predictions[:, 1, :])
.float()
.mean()
.item()
)
metrics[f"localization_miou_{name}"] = calculate_miou(
predictions[:, 1, :], true_predictions[:, 1, :]
)
return metrics
def evaluate_augmentations(
positive: torch.Tensor,
negative: torch.Tensor,
augmentation_name: str,
message: torch.Tensor,
) -> dict:
"""calculating evaluation metrics but take name of the augmentation
method that has been done before getting positive and negative results"""
metrics = {}
metrics[f"aug_{augmentation_name}_acc"] = compute_accuracy(positive, negative)
metrics[f"aug_{augmentation_name}_fpr"] = compute_FPR(negative)
metrics[f"aug_{augmentation_name}_fnr"] = compute_FNR(positive)
if message.shape[0] != 0:
metrics[f"aug_{augmentation_name}_bit_acc"] = compute_bit_acc(positive, message)
# add one metric which is average overall score of all augmentations
metrics["all_aug_acc"] = compute_accuracy(positive, negative)
return metrics
def evaluate_audio_watermark(
y_pred: torch.Tensor,
y: torch.Tensor,
cfg: DictConfig,
) -> dict:
"""Audio reconstruction evaluation method that can be conveniently pickled."""
metrics = {}
if cfg.evaluate.metrics.visqol:
visqol = builders.get_visqol(cfg.metrics.visqol)
metrics["visqol"] = visqol(y_pred, y, cfg.sample_rate)
sisnr = ScaleInvariantSignalNoiseRatio().to(y.device)
stoi = ShortTimeObjectiveIntelligibility(fs=cfg.sample_rate)
metrics["sisnr"] = sisnr(y_pred, y)
metrics["stoi"] = stoi(y_pred, y)
metrics["pesq"] = tensor_pesq(y_pred, y, sr=cfg.sample_rate)
return metrics
def tensor_pesq(y_pred: torch.Tensor, y: torch.Tensor, sr: int):
# pesq returns error if no speech is detected, so we catch it
return PesqMetric(sr)(y_pred, y).item()
def compute_accuracy(positive, negative):
N = (positive[:, 1, :].mean(dim=1) > 0.5).sum() + (
negative[:, 0, :].mean(dim=1) > 0.5
).sum()
acc = N / (2 * positive.size(0))
return acc
def compute_FPR(negative):
N = (negative[:, 1, :].mean(dim=1) > 0.5).sum()
fpr = N / (negative.size(0))
return fpr
def compute_FNR(positive):
N = (positive[:, 0, :].mean(dim=1) > 0.5).sum()
fpr = N / (positive.size(0))
return fpr
def _bit_acc(decoded, original):
bit_acc = (decoded == original).float().mean()
return bit_acc
def compute_bit_acc(positive, original, mask=None):
"""Compute bit accuracy.
Args:
positive: detector outputs [bsz, 2+nbits, time_steps]
original: original message (0 or 1) [bsz, nbits]
mask: mask of the watermark [bsz, 1, time_steps]
"""
decoded = positive[:, 2:, :] # b 2+nbits t -> b nbits t
if mask is not None:
# cut last dim of positive to keep only where mask is 1
new_shape = [*decoded.shape[:-1], -1] # b nbits t -> b nbits -1
decoded = torch.masked_select(decoded, mask == 1).reshape(new_shape)
# average decision over time, then threshold
decoded = decoded.mean(dim=-1) > 0 # b nbits
return _bit_acc(decoded, original)