nvidia-nemo--speech
wehub-resource-sync
ba4be087d5
Create PR to main with cherry-pick from release / cherry-pick (push) Failing after 0s
CICD NeMo / pre-flight (push) Failing after 0s
CICD NeMo / configure (push) Has been skipped
Build, validate, and release Neural Modules / pre-flight (push) Failing after 1s
CICD NeMo / code-linting (push) Has been skipped
Build, validate, and release Neural Modules / release (push) Has been skipped
Build, validate, and release Neural Modules / release-summary (push) Has been cancelled
CICD NeMo / cicd-test-container-build (push) Has been cancelled
CICD NeMo / cicd-import-tests (push) Has been cancelled
CICD NeMo / L0_Setup_Test_Data_And_Models (push) Has been cancelled
CICD NeMo / cicd-main-unit-tests (push) Has been cancelled
CICD NeMo / cicd-main-speech (push) Has been cancelled
CICD NeMo / Nemo_CICD_Test (push) Has been cancelled
CICD NeMo / Coverage (e2e) (push) Has been cancelled
CICD NeMo / Coverage (unit-test) (push) Has been cancelled
CodeQL / Analyze (python) (push) Has been cancelled
CICD NeMo / cicd-wait-in-queue (push) Has been cancelled
820 行
29 KiB
Python
820 行
29 KiB
Python
# Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
|
|
from abc import ABC, abstractmethod
|
|
from typing import Optional, Tuple, Type
|
|
|
|
import numpy as np
|
|
import torch
|
|
|
|
from nemo.collections.common.parts.utils import mask_sequence_tensor
|
|
from nemo.core.classes import NeuralModule, typecheck
|
|
from nemo.core.neural_types import FloatType, LengthsType, NeuralType, SpectrogramType, VoidType
|
|
from nemo.utils import logging
|
|
|
|
|
|
class StochasticDifferentialEquation(NeuralModule, ABC):
|
|
"""Base class for stochastic differential equations."""
|
|
|
|
def __init__(self, time_min: float, time_max: float, num_steps: int):
|
|
super().__init__()
|
|
|
|
# min and max time
|
|
if time_min <= 0:
|
|
raise ValueError(f'time_min should be positive, current value {time_min}')
|
|
|
|
if time_max <= time_min:
|
|
raise ValueError(f'time_max should be larger than time_min, current max {time_max} and min {time_min}')
|
|
|
|
self.time_min = time_min
|
|
self.time_max = time_max
|
|
|
|
# number of steps
|
|
if num_steps <= 0:
|
|
raise ValueError(f'num_steps needs to be positive: current value {num_steps}')
|
|
|
|
self.num_steps = num_steps
|
|
|
|
@property
|
|
def dt(self) -> float:
|
|
"""Time step for this SDE.
|
|
This denotes the step size between `0` and `self.time_max` when using `self.num_steps`.
|
|
"""
|
|
return self.time_max / self.num_steps
|
|
|
|
@property
|
|
def time_delta(self) -> float:
|
|
"""Time range for this SDE."""
|
|
return self.time_max - self.time_min
|
|
|
|
def generate_time(self, size: int, device: torch.device) -> torch.Tensor:
|
|
"""Generate random time steps in the valid range.
|
|
|
|
Time steps are generated between `self.time_min` and `self.time_max`.
|
|
|
|
Args:
|
|
size: number of samples
|
|
device: device to use
|
|
|
|
Returns:
|
|
A tensor of floats with shape (size,)
|
|
"""
|
|
time = torch.rand(size, device=device) * self.time_delta + self.time_min
|
|
return time
|
|
|
|
@abstractmethod
|
|
def coefficients(self, state: torch.Tensor, time: torch.Tensor, **kwargs) -> Tuple[torch.Tensor, torch.Tensor]:
|
|
"""
|
|
Args:
|
|
state: tensor of shape (B, C, D, T)
|
|
time: tensor of shape (B,)
|
|
|
|
Returns:
|
|
Tuple with drift and diffusion coefficients.
|
|
"""
|
|
pass
|
|
|
|
@typecheck(
|
|
input_types={
|
|
"prior_mean": NeuralType(('B', 'C', 'D', 'T'), VoidType()),
|
|
},
|
|
output_types={
|
|
"sample": NeuralType(('B', 'C', 'D', 'T'), VoidType()),
|
|
},
|
|
)
|
|
@abstractmethod
|
|
def prior_sampling(self, prior_mean: torch.Tensor) -> torch.Tensor:
|
|
"""Generate a sample from the prior distribution p_T.
|
|
|
|
Args:
|
|
prior_mean: Mean of the prior distribution
|
|
|
|
Returns:
|
|
A sample from the prior distribution.
|
|
"""
|
|
pass
|
|
|
|
def discretize(
|
|
self, *, state: torch.Tensor, time: torch.Tensor, state_length: Optional[torch.Tensor] = None, **kwargs
|
|
) -> Tuple[torch.Tensor, torch.Tensor]:
|
|
"""Assume we have the following SDE:
|
|
|
|
dx = drift(x, t) * dt + diffusion(x, t) * dwt
|
|
|
|
where `wt` is the standard Wiener process.
|
|
|
|
We assume the following discretization:
|
|
|
|
new_state = current_state + total_drift + total_diffusion * z_norm
|
|
|
|
where `z_norm` is sampled from normal distribution with zero mean and unit variance.
|
|
|
|
Args:
|
|
state: current state of the process, shape (B, C, D, T)
|
|
time: current time of the process, shape (B,)
|
|
state_length: length of the valid time steps for each example in the batch, shape (B,)
|
|
**kwargs: other parameters
|
|
|
|
Returns:
|
|
Drift and diffusion.
|
|
"""
|
|
# Get coefficients
|
|
drift_coefficient, diffusion_coefficient = self.coefficients(
|
|
state=state, time=time, state_length=state_length, **kwargs
|
|
)
|
|
|
|
# Discretized drift
|
|
drift = drift_coefficient * self.dt
|
|
|
|
# Note:
|
|
# Scale with sqrt(dt) because z_norm is sampled from a normal distribution with zero mean and
|
|
# unit variance and dwt is normally distributed with zero mean and variance dt
|
|
diffusion = diffusion_coefficient * np.sqrt(self.dt)
|
|
|
|
return drift, diffusion
|
|
|
|
@abstractmethod
|
|
def copy(self):
|
|
"""Create a copy of this SDE."""
|
|
pass
|
|
|
|
def __repr__(self):
|
|
desc = f'{self.__class__.__name__}(time_min={self.time_min}, time_max={self.time_max}, num_steps={self.num_steps})'
|
|
desc += f'\n\tdt: {self.dt}'
|
|
desc += f'\n\ttime_delta: {self.time_delta}'
|
|
return desc
|
|
|
|
|
|
class OrnsteinUhlenbeckVarianceExplodingSDE(StochasticDifferentialEquation):
|
|
"""This class implements the Ornstein-Uhlenbeck SDE with variance exploding noise schedule.
|
|
|
|
The SDE is given by:
|
|
|
|
dx = theta * (y - x) dt + g(t) dw
|
|
|
|
where `theta` is the stiffness parameter and `g(t)` is the diffusion coefficient:
|
|
|
|
g(t) = std_min * (std_max/std_min)^t * sqrt(2 * log(std_max/std_min))
|
|
|
|
References:
|
|
Richter et al., Speech Enhancement and Dereverberation with Diffusion-based Generative Models, Tr. ASLP 2023
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
stiffness: float,
|
|
std_min: float,
|
|
std_max: float,
|
|
num_steps: int = 100,
|
|
time_min: float = 3e-2,
|
|
time_max: float = 1.0,
|
|
eps: float = 1e-8,
|
|
):
|
|
super().__init__(time_min=time_min, time_max=time_max, num_steps=num_steps)
|
|
|
|
# Small regularization
|
|
if eps <= 0:
|
|
raise ValueError(f'eps should be positive, current value {eps}')
|
|
self.eps = eps
|
|
|
|
# stifness
|
|
self.stiffness = stiffness
|
|
|
|
# noise schedule
|
|
if std_min <= 0:
|
|
raise ValueError(f'std_min should be positive, current value {std_min}')
|
|
|
|
if std_max <= std_min:
|
|
raise ValueError(f'std_max should be larger than std_min, current max {std_max} and min {std_min}')
|
|
|
|
self.std_min = std_min
|
|
self.std_max = std_max
|
|
|
|
logging.debug('Initialized %s with', self.__class__.__name__)
|
|
logging.debug('\tstiffness: %s', self.stiffness)
|
|
logging.debug('\tstd_min: %s', self.std_min)
|
|
logging.debug('\tstd_max: %s', self.std_max)
|
|
logging.debug('\tnum_steps: %s', self.num_steps)
|
|
logging.debug('\ttime_min: %s', self.time_min)
|
|
logging.debug('\ttime_max: %s', self.time_max)
|
|
logging.debug('\teps: %s', self.eps)
|
|
|
|
@property
|
|
def std_ratio(self) -> float:
|
|
return self.std_max / (self.std_min + self.eps)
|
|
|
|
@property
|
|
def log_std_ratio(self) -> float:
|
|
return np.log(self.std_ratio + self.eps)
|
|
|
|
@typecheck(
|
|
input_types={
|
|
"state": NeuralType(('B', 'C', 'D', 'T'), VoidType()),
|
|
"prior_mean": NeuralType(('B', 'C', 'D', 'T'), VoidType()),
|
|
"time": NeuralType(tuple('B'), FloatType()),
|
|
},
|
|
output_types={
|
|
"mean": NeuralType(('B', 'C', 'D', 'T'), FloatType()),
|
|
},
|
|
)
|
|
def perturb_kernel_mean(self, state: torch.Tensor, prior_mean: torch.Tensor, time: torch.Tensor) -> torch.Tensor:
|
|
"""Return the mean of the perturbation kernel for this SDE.
|
|
|
|
Args:
|
|
state: current state of the process, shape (B, C, D, T)
|
|
prior_mean: mean of the prior distribution
|
|
time: current time of the process, shape (B,)
|
|
|
|
Returns:
|
|
A tensor of shape (B, C, D, T)
|
|
"""
|
|
# exponential weighting
|
|
weight = torch.exp(-self.stiffness * time)
|
|
|
|
# view as [B, C, D, T]
|
|
weight = weight.view(-1, 1, 1, 1)
|
|
|
|
# closed-form mean
|
|
mean = weight * state + (1 - weight) * prior_mean
|
|
|
|
return mean
|
|
|
|
@typecheck(
|
|
input_types={
|
|
"time": NeuralType(tuple('B'), FloatType()),
|
|
},
|
|
output_types={
|
|
"std": NeuralType(tuple('B'), FloatType()),
|
|
},
|
|
)
|
|
def perturb_kernel_std(self, time: torch.Tensor) -> torch.Tensor:
|
|
"""Return the standard deviation of the perturbation kernel for this SDE.
|
|
|
|
Note that the standard deviation depends on the time and the noise schedule,
|
|
which is parametrized using `self.stiffness`, `self.std_min` and `self.std_max`.
|
|
|
|
Args:
|
|
time: current time of the process, shape (B,)
|
|
|
|
Returns:
|
|
A tensor of shape (B,)
|
|
"""
|
|
var = (self.std_min**2) * self.log_std_ratio
|
|
var *= torch.pow(self.std_ratio, 2 * time) - torch.exp(-2 * self.stiffness * time)
|
|
var /= self.stiffness + self.log_std_ratio
|
|
std = torch.sqrt(var)
|
|
return std
|
|
|
|
@typecheck(
|
|
input_types={
|
|
"state": NeuralType(('B', 'C', 'D', 'T'), VoidType()),
|
|
"prior_mean": NeuralType(('B', 'C', 'D', 'T'), VoidType()),
|
|
"time": NeuralType(tuple('B'), FloatType()),
|
|
},
|
|
output_types={
|
|
"mean": NeuralType(('B', 'C', 'D', 'T'), FloatType()),
|
|
"std": NeuralType(('B', 'C', 'D', 'T'), FloatType()),
|
|
},
|
|
)
|
|
def perturb_kernel_params(self, state: torch.Tensor, prior_mean: torch.Tensor, time: torch.Tensor) -> torch.Tensor:
|
|
"""Return the mean and standard deviation of the perturbation kernel for this SDE.
|
|
|
|
Args:
|
|
state: current state of the process, shape (B, C, D, T)
|
|
prior_mean: mean of the prior distribution
|
|
time: current time of the process, shape (B,)
|
|
"""
|
|
assert torch.all(time <= self.time_max)
|
|
assert torch.all(time >= self.time_min)
|
|
|
|
# compute the mean
|
|
mean = self.perturb_kernel_mean(state=state, prior_mean=prior_mean, time=time)
|
|
|
|
# compute the standard deviation
|
|
std = self.perturb_kernel_std(time=time)
|
|
# view as [B, C, D, T]
|
|
std = std.view(-1, 1, 1, 1)
|
|
|
|
return mean, std
|
|
|
|
@typecheck(
|
|
input_types={
|
|
"state": NeuralType(('B', 'C', 'D', 'T'), VoidType()),
|
|
"time": NeuralType(tuple('B'), VoidType()),
|
|
"prior_mean": NeuralType(('B', 'C', 'D', 'T'), VoidType()),
|
|
"state_length": NeuralType(tuple('B'), LengthsType(), optional=True),
|
|
},
|
|
output_types={
|
|
"drift_coefficient": NeuralType(('B', 'C', 'D', 'T'), FloatType()),
|
|
"diffusion_coefficient": NeuralType(('B', 'C', 'D', 'T'), FloatType()),
|
|
},
|
|
)
|
|
def coefficients(
|
|
self,
|
|
state: torch.Tensor,
|
|
time: torch.Tensor,
|
|
prior_mean: torch.Tensor,
|
|
state_length: Optional[torch.Tensor] = None,
|
|
) -> Tuple[torch.Tensor, torch.Tensor]:
|
|
"""Compute drift and diffusion coefficients for this SDE.
|
|
|
|
Args:
|
|
state: current state of the process, shape (B, C, D, T)
|
|
time: current time of the process, shape (B,)
|
|
prior_mean: mean of the prior distribution
|
|
state_length: length of the valid time steps for each example in the batch
|
|
|
|
Returns:
|
|
Drift and diffusion coefficients.
|
|
"""
|
|
# Drift coefficient
|
|
drift_coefficient = self.stiffness * (prior_mean - state)
|
|
|
|
# Diffusion coefficient
|
|
diffusion_coefficient = self.std_min * torch.pow(self.std_ratio, time) * np.sqrt(2 * self.log_std_ratio)
|
|
# View in the same shape as the state
|
|
diffusion_coefficient = diffusion_coefficient.view(-1, *([1] * (state.dim() - 1)))
|
|
|
|
if state_length is not None:
|
|
drift_coefficient = mask_sequence_tensor(drift_coefficient, state_length)
|
|
diffusion_coefficient = mask_sequence_tensor(diffusion_coefficient, state_length)
|
|
|
|
return drift_coefficient, diffusion_coefficient
|
|
|
|
def prior_sampling(self, prior_mean: torch.Tensor) -> torch.Tensor:
|
|
"""Generate a sample from the prior distribution p_T.
|
|
|
|
Args:
|
|
prior_mean: Mean of the prior distribution
|
|
"""
|
|
# Final time step for all samples in the batch
|
|
time = self.time_max * torch.ones(prior_mean.shape[0], device=prior_mean.device)
|
|
|
|
# Compute the std of the prior distribution
|
|
std = self.perturb_kernel_std(time=time)
|
|
|
|
# view as [B, C, D, T]
|
|
std = std.view(-1, 1, 1, 1)
|
|
|
|
# Generate a sample from a normal distribution centered at prior_mean
|
|
sample = prior_mean + torch.randn_like(prior_mean) * std
|
|
|
|
return sample
|
|
|
|
def copy(self):
|
|
return OrnsteinUhlenbeckVarianceExplodingSDE(
|
|
stiffness=self.stiffness,
|
|
std_min=self.std_min,
|
|
std_max=self.std_max,
|
|
num_steps=self.num_steps,
|
|
time_min=self.time_min,
|
|
time_max=self.time_max,
|
|
eps=self.eps,
|
|
)
|
|
|
|
def __repr__(self):
|
|
desc = f'{self.__class__.__name__}(stiffness={self.stiffness}, std_min={self.std_min}, std_max={self.std_max}, num_steps={self.num_steps}, time_min={self.time_min}, time_max={self.time_max}, eps={self.eps})'
|
|
desc += f'\n\tdt: {self.dt}'
|
|
desc += f'\n\ttime_delta: {self.time_delta}'
|
|
desc += f'\n\tstd_ratio: {self.std_ratio}'
|
|
desc += f'\n\tlog_std_ratio: {self.log_std_ratio}'
|
|
|
|
return desc
|
|
|
|
|
|
class ReverseStochasticDifferentialEquation(StochasticDifferentialEquation):
|
|
def __init__(self, *, sde: Type[StochasticDifferentialEquation], score_estimator: Type[NeuralModule]):
|
|
"""Use the forward SDE and a score estimator to define the reverse SDE.
|
|
|
|
Args:
|
|
sde: forward SDE
|
|
score_estimator: neural score estimator
|
|
"""
|
|
super().__init__(time_min=sde.time_min, time_max=sde.time_max, num_steps=sde.num_steps)
|
|
self.score_estimator = score_estimator
|
|
self.forward_sde = sde
|
|
|
|
logging.debug('Initialized %s', self.__class__.__name__)
|
|
|
|
def coefficients(
|
|
self,
|
|
state: torch.Tensor,
|
|
time: torch.Tensor,
|
|
score_condition: Optional[torch.Tensor] = None,
|
|
state_length: Optional[torch.Tensor] = None,
|
|
**kwargs,
|
|
) -> Tuple[torch.Tensor, torch.Tensor]:
|
|
"""Compute drift and diffusion coefficients for the reverse SDE.
|
|
|
|
Args:
|
|
state: current state of the process, shape (B, C, D, T)
|
|
time: current time of the process, shape (B,)
|
|
"""
|
|
raise NotImplementedError('Coefficients not necessary for the reverse SDE.')
|
|
|
|
def prior_sampling(self, shape: torch.Size, device: torch.device) -> torch.Tensor:
|
|
"""Prior sampling is not necessary for the reverse SDE."""
|
|
raise NotImplementedError('Prior sampling not necessary for the reverse SDE.')
|
|
|
|
def discretize(
|
|
self,
|
|
*,
|
|
state: torch.Tensor,
|
|
time: torch.Tensor,
|
|
score_condition: Optional[torch.Tensor] = None,
|
|
state_length: Optional[torch.Tensor] = None,
|
|
**kwargs,
|
|
) -> Tuple[torch.Tensor, torch.Tensor]:
|
|
"""Discretize the reverse SDE.
|
|
|
|
Args:
|
|
state: current state of the process, shape (B, C, D, T)
|
|
time: current time of the process, shape (B,)
|
|
score_condition: condition for the score estimator
|
|
state_length: length of the valid time steps for each example in the batch
|
|
**kwargs: other parameters for discretization of the forward SDE
|
|
"""
|
|
# Drift and diffusion from the forward SDE
|
|
forward_drift, forward_diffusion = self.forward_sde.discretize(state=state, time=time, **kwargs)
|
|
|
|
# For input for the score estimator:
|
|
# - if no condition is provided, use the state
|
|
# - if a condition is provided, concatenate the state and the condition along the channel dimension
|
|
score_input = state if score_condition is None else torch.cat([state, score_condition], dim=1)
|
|
|
|
# Estimate score
|
|
score, _ = self.score_estimator(input=score_input, input_length=state_length, condition=time)
|
|
|
|
# Adjust drift
|
|
drift = forward_drift - forward_diffusion.pow(2) * score
|
|
|
|
# Adjust diffusion
|
|
diffusion = forward_diffusion
|
|
|
|
if state_length is not None:
|
|
drift = mask_sequence_tensor(drift, state_length)
|
|
diffusion = mask_sequence_tensor(diffusion, state_length)
|
|
|
|
return drift, diffusion
|
|
|
|
def copy(self):
|
|
return ReverseStochasticDifferentialEquation(sde=self.forward_sde.copy(), score_estimator=self.score_estimator)
|
|
|
|
def __repr__(self):
|
|
desc = f'{self.__class__.__name__}(sde={self.forward_sde}, score_estimator={self.score_estimator})'
|
|
return desc
|
|
|
|
|
|
class PredictorCorrectorSampler(NeuralModule):
|
|
"""Predictor-Corrector sampler for the reverse SDE.
|
|
|
|
Args:
|
|
sde: forward SDE
|
|
score_estimator: neural score estimator
|
|
predictor: predictor for the reverse process
|
|
corrector: corrector for the reverse process
|
|
num_steps: number of time steps for the reverse process
|
|
num_corrector_steps: number of corrector steps
|
|
time_max: maximum time
|
|
time_min: minimum time
|
|
snr: SNR for Annealed Langevin Dynamics
|
|
output_type: type of the output ('state' for the final state, or 'mean' for the mean of the final state)
|
|
|
|
References:
|
|
- Song et al., Score-based generative modeling through stochastic differential equations, 2021
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
sde,
|
|
score_estimator,
|
|
predictor: str = 'reverse_diffusion',
|
|
corrector: str = 'annealed_langevin_dynamics',
|
|
num_steps: int = 50,
|
|
num_corrector_steps: int = 1,
|
|
time_max: Optional[float] = None,
|
|
time_min: Optional[float] = None,
|
|
snr: float = 0.5,
|
|
output_type: str = 'mean',
|
|
):
|
|
super().__init__()
|
|
# Create a copy of SDE
|
|
self.sde = sde.copy()
|
|
|
|
# Update SDE parameters for sampling
|
|
if time_max is not None:
|
|
self.sde.time_max = time_max
|
|
logging.info('sde.time_max set to: %s', self.sde.time_max)
|
|
|
|
if time_min is not None:
|
|
self.sde.time_min = time_min
|
|
logging.info('sde.time_min set to: %s', self.sde.time_min)
|
|
|
|
self.sde.num_steps = num_steps
|
|
logging.info('sde.num_steps set to: %s', self.sde.num_steps)
|
|
|
|
# Update local values
|
|
self.time_max = self.sde.time_max
|
|
self.time_min = self.sde.time_min
|
|
self.num_steps = self.sde.num_steps
|
|
|
|
# Predictor setup
|
|
if predictor == 'reverse_diffusion':
|
|
self.predictor = ReverseDiffusionPredictor(sde=self.sde, score_estimator=score_estimator)
|
|
else:
|
|
raise RuntimeError(f'Unexpected predictor: {predictor}')
|
|
|
|
# Corrector setup
|
|
if corrector == 'annealed_langevin_dynamics':
|
|
self.corrector = AnnealedLangevinDynamics(
|
|
sde=self.sde, score_estimator=score_estimator, snr=snr, num_steps=num_corrector_steps
|
|
)
|
|
else:
|
|
raise RuntimeError(f'Unexpected corrector: {corrector}')
|
|
|
|
if output_type not in ['mean', 'state']:
|
|
raise ValueError(f'Unexpected output type: {output_type}')
|
|
self.output_type = output_type
|
|
|
|
logging.debug('Initialized %s with', self.__class__.__name__)
|
|
logging.debug('\tpredictor: %s', predictor)
|
|
logging.debug('\tcorrector: %s', corrector)
|
|
logging.debug('\tnum_steps: %s', self.num_steps)
|
|
logging.debug('\ttime_min: %s', self.time_min)
|
|
logging.debug('\ttime_max: %s', self.time_max)
|
|
logging.debug('\tnum_corrector_steps: %s', num_corrector_steps)
|
|
logging.debug('\tsnr: %s', snr)
|
|
logging.debug('\toutput_type: %s', self.output_type)
|
|
|
|
@typecheck(
|
|
input_types={
|
|
"prior_mean": NeuralType(('B', 'C', 'D', 'T'), SpectrogramType()),
|
|
"score_condition": NeuralType(('B', 'C', 'D', 'T'), SpectrogramType(), optional=True),
|
|
"state_length": NeuralType(tuple('B'), LengthsType(), optional=True),
|
|
},
|
|
output_types={
|
|
"sample": NeuralType(('B', 'C', 'D', 'T'), SpectrogramType()),
|
|
"state_length": NeuralType(tuple('B'), LengthsType(), optional=True),
|
|
},
|
|
)
|
|
@torch.inference_mode()
|
|
def forward(
|
|
self, prior_mean: torch.Tensor, score_condition: torch.Tensor, state_length: Optional[torch.Tensor] = None
|
|
) -> torch.Tensor:
|
|
"""Takes prior (noisy) mean and generates a sample by solving the reverse SDE.
|
|
|
|
Args:
|
|
prior_mean: mean for the prior distribution, e.g., noisy observation
|
|
score_condition: conditioning for the score estimator
|
|
state_length: length of the valid time steps for each example in the batch
|
|
|
|
Returns:
|
|
Generated `sample` and the corresponding `sample_length`.
|
|
"""
|
|
# Sample from the prior distribution
|
|
state = self.sde.prior_sampling(prior_mean=prior_mean)
|
|
|
|
if state_length is not None:
|
|
state = mask_sequence_tensor(state, state_length)
|
|
|
|
# Time steps for evaluation
|
|
time_steps = torch.linspace(self.time_max, self.time_min, self.num_steps, device=state.device)
|
|
|
|
# Sampling
|
|
for t in time_steps:
|
|
# time steps for the whole batch
|
|
time = t * torch.ones(state.shape[0], device=state.device)
|
|
|
|
# corrector step
|
|
state, _ = self.corrector(
|
|
state=state, time=time, score_condition=score_condition, state_length=state_length
|
|
)
|
|
|
|
# predictor step
|
|
state, state_mean = self.predictor(
|
|
state=state,
|
|
time=time,
|
|
score_condition=score_condition,
|
|
prior_mean=prior_mean,
|
|
state_length=state_length,
|
|
)
|
|
|
|
# Final output
|
|
if self.output_type == 'state':
|
|
sample = state
|
|
elif self.output_type == 'mean':
|
|
sample = state_mean
|
|
else:
|
|
raise RuntimeError(f'Unexpected output type: {self.output_type}')
|
|
|
|
if state_length is not None:
|
|
sample = mask_sequence_tensor(sample, state_length)
|
|
|
|
return sample, state_length
|
|
|
|
|
|
class Predictor(torch.nn.Module, ABC):
|
|
"""Predictor for the reverse process.
|
|
|
|
Args:
|
|
sde: forward SDE
|
|
score_estimator: neural score estimator
|
|
"""
|
|
|
|
def __init__(self, sde, score_estimator):
|
|
super().__init__()
|
|
self.reverse_sde = ReverseStochasticDifferentialEquation(sde=sde, score_estimator=score_estimator)
|
|
|
|
@abstractmethod
|
|
@torch.inference_mode()
|
|
def forward(
|
|
self,
|
|
*,
|
|
state: torch.Tensor,
|
|
time: torch.Tensor,
|
|
score_condition: Optional[torch.Tensor] = None,
|
|
state_length: Optional[torch.Tensor] = None,
|
|
**kwargs,
|
|
):
|
|
"""Predict the next state of the reverse process.
|
|
|
|
Args:
|
|
state: current state of the process, shape (B, C, D, T)
|
|
time: current time of the process, shape (B,)
|
|
score_condition: conditioning for the score estimator
|
|
state_length: length of the valid time steps for each example in the batch
|
|
|
|
Returns:
|
|
New state and mean.
|
|
"""
|
|
pass
|
|
|
|
|
|
class ReverseDiffusionPredictor(Predictor):
|
|
"""Predict the next state of the reverse process using the reverse diffusion process.
|
|
|
|
Args:
|
|
sde: forward SDE
|
|
score_estimator: neural score estimator
|
|
"""
|
|
|
|
def __init__(self, sde, score_estimator):
|
|
super().__init__(sde=sde, score_estimator=score_estimator)
|
|
|
|
@torch.inference_mode()
|
|
def forward(self, *, state, time, score_condition=None, state_length=None, **kwargs):
|
|
"""Predict the next state of the reverse process using the reverse diffusion process.
|
|
|
|
Args:
|
|
state: current state of the process, shape (B, C, D, T)
|
|
time: current time of the process, shape (B,)
|
|
score_condition: conditioning for the score estimator
|
|
state_length: length of the valid time steps for each example in the batch
|
|
|
|
Returns:
|
|
New state and mean of the diffusion process.
|
|
"""
|
|
drift, diffusion = self.reverse_sde.discretize(
|
|
state=state, time=time, score_condition=score_condition, state_length=state_length, **kwargs
|
|
)
|
|
|
|
# Generate a random sample from a standard normal distribution
|
|
z_norm = torch.randn_like(state)
|
|
|
|
# Compute the mean of the next state
|
|
mean = state - drift
|
|
|
|
# Compute new state by sampling
|
|
new_state = mean + diffusion * z_norm
|
|
|
|
if state_length is not None:
|
|
new_state = mask_sequence_tensor(new_state, state_length)
|
|
mean = mask_sequence_tensor(mean, state_length)
|
|
|
|
return new_state, mean
|
|
|
|
|
|
class Corrector(NeuralModule, ABC):
|
|
"""Corrector for the reverse process.
|
|
|
|
Args:
|
|
sde: forward SDE
|
|
score_estimator: neural score estimator
|
|
snr: SNR for Annealed Langevin Dynamics
|
|
num_steps: number of steps for the corrector
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
sde: Type[StochasticDifferentialEquation],
|
|
score_estimator: Type[NeuralModule],
|
|
snr: float,
|
|
num_steps: int,
|
|
):
|
|
super().__init__()
|
|
self.sde = sde
|
|
self.score_estimator = score_estimator
|
|
self.snr = snr
|
|
self.num_steps = num_steps
|
|
|
|
logging.debug('Initialized %s with', self.__class__.__name__)
|
|
logging.debug('\tsnr: %s', snr)
|
|
logging.debug('\tnum_steps: %s', num_steps)
|
|
|
|
@abstractmethod
|
|
@typecheck(
|
|
input_types={
|
|
"state": NeuralType(('B', 'C', 'D', 'T'), VoidType()),
|
|
"time": NeuralType(tuple('B'), FloatType()),
|
|
"score_condition": NeuralType(('B', 'C', 'D', 'T'), VoidType(), optional=True),
|
|
"state_length": NeuralType(tuple('B'), LengthsType(), optional=True),
|
|
},
|
|
output_types={
|
|
"state": NeuralType(('B', 'C', 'D', 'T'), VoidType()),
|
|
},
|
|
)
|
|
@torch.inference_mode()
|
|
def forward(self, state, time, score_condition=None, state_length=None):
|
|
"""
|
|
Args:
|
|
state: current state of the process, shape (B, C, D, T)
|
|
time: current time of the process, shape (B,)
|
|
score_condition: conditioning for the score estimator
|
|
state_length: length of the valid time steps for each example in the batch
|
|
|
|
Returns:
|
|
New state and mean.
|
|
"""
|
|
pass
|
|
|
|
|
|
class AnnealedLangevinDynamics(Corrector):
|
|
"""Annealed Langevin Dynamics for the reverse process.
|
|
|
|
References:
|
|
- Song et al., Score-based generative modeling through stochastic differential equations, 2021
|
|
"""
|
|
|
|
def __init__(self, sde, **kwargs):
|
|
if not isinstance(sde, OrnsteinUhlenbeckVarianceExplodingSDE):
|
|
raise ValueError(f'Expected an instance of OrnsteinUhlenbeckVarianceExplodingSDE, got {type(sde)}')
|
|
super().__init__(sde=sde, **kwargs)
|
|
|
|
@torch.inference_mode()
|
|
def forward(self, state, time, score_condition=None, state_length=None):
|
|
"""Correct the state using Annealed Langevin Dynamics.
|
|
|
|
Args:
|
|
state: current state of the process, shape (B, C, D, T)
|
|
time: current time of the process, shape (B,)
|
|
score_condition: conditioning for the score estimator
|
|
state_length: length of the valid time steps for each example in the batch
|
|
|
|
Returns:
|
|
New state and mean of the diffusion process.
|
|
|
|
References:
|
|
Alg. 4 in http://arxiv.org/abs/2011.13456
|
|
"""
|
|
# Compute the standard deviation of the diffusion process
|
|
std = self.sde.perturb_kernel_std(time=time)
|
|
# View as [B, 1, 1, 1]
|
|
std = std.view(-1, *([1] * (state.dim() - 1)))
|
|
|
|
for i in range(self.num_steps):
|
|
# prepare input for the score estimator, concatenate conditioning along the channel dimension
|
|
score_input = state if score_condition is None else torch.cat([state, score_condition], dim=1)
|
|
|
|
# calculate the score
|
|
score, _ = self.score_estimator(input=score_input, input_length=state_length, condition=time)
|
|
|
|
# generate a sample from a standard normal distribution
|
|
z_norm = torch.randn_like(state)
|
|
|
|
# compute the step size
|
|
# note: this is slightly different than in the paper, where std = ||z_norm||_2 / ||score||_2
|
|
step_size = 2 * (self.snr * std).pow(2)
|
|
|
|
# update the mean
|
|
mean = state + step_size * score
|
|
|
|
# update the state
|
|
state = mean + z_norm * torch.sqrt(step_size * 2)
|
|
|
|
if state_length is not None:
|
|
state = mask_sequence_tensor(state, state_length)
|
|
mean = mask_sequence_tensor(mean, state_length)
|
|
|
|
return state, mean
|