speechbrain.nnet.diffusion 模块

去噪扩散的实现

https://arxiv.org/pdf/2006.11239.pdf

某些部分借鉴自 / 灵感来源于 denoising-diffusion-pytorch https://github.com/lucidrains/denoising-diffusion-pytorch

Authors

阿尔乔姆·普洛日尼科夫 2022

摘要

类：

`DenoisingDiffusion`	经典去噪扩散概率模型（DDPM）的实现
`Diffuser`	一个基础的扩散实现
`DiffusionTrainSample`
`GaussianNoise`	添加普通高斯噪声
`LatentDiffusion`	一个潜在扩散包装器。
`LatentDiffusionTrainSample`
`LengthMaskedGaussianNoise`	应用于填充样本的高斯噪声。

函数：

sample_timesteps

返回一个随机的时间步样本作为一维张量（仅一个维度）

参考

class speechbrain.nnet.diffusion.Diffuser(model, timesteps, noise=None)[source]

基础：Module

一个基础的扩散实现

Parameters:

model (nn.Module) – 基础模型
timesteps (int) – 时间步数
noise (callable|str) –
使用的噪声函数/模块

提供了以下预定义的噪声类型 “gaussian”: 高斯噪声，应用于整个样本 “length_masked_gaussian”: 仅应用于样本中非填充部分的高斯噪声

应用于样本中非填充部分

distort(x, timesteps=None)[source]

向一批数据添加噪声

Parameters:

x (torch.Tensor) – 原始数据样本
timesteps (torch.Tensor) – 一个长度为x中批次数的1-D整数张量，其中每个条目对应于批次的timestep编号。如果省略，timesteps将被随机采样

train_sample(x, timesteps=None, condition=None, **kwargs)[source]

为训练循环创建一个带有相应目标的样本

Parameters:

x (torch.Tensor) – 原始数据样本
timesteps (torch.Tensor) – 一个长度为x中批次数的1-D整数张量，其中每个条目对应于批次的timestep编号。如果省略，timesteps将被随机采样
condition (torch.Tensor) – 用于条件生成的condition 在无条件生成时应省略
**kwargs (dict) – 传递给底层模型的参数。

Returns:

pred (torch.Tensor) – 模型输出的预测噪声
noise (torch.Tensor) – 被应用的噪声
noisy_sample (torch.Tensor) – 应用了噪声的样本

sample(shape, **kwargs)[source]

生成由count参数指示的样本数量

Parameters:

shape (可枚举) – 要生成的样本的形状
**kwargs (dict) – 传递给底层模型的参数。

forward(x, timesteps=None)[source]: 计算前向传播，调用distort()

class speechbrain.nnet.diffusion.DenoisingDiffusion(model, timesteps=None, noise=None, beta_start=None, beta_end=None, sample_min=None, sample_max=None, show_progress=False)[source]

基础类: Diffuser

经典去噪扩散概率模型（DDPM）的实现

Parameters:

model (nn.Module) – 基础模型
timesteps (int) – 时间步数
噪声 (str|nn.Module) – 使用的噪声类型 “gaussian” 将产生标准高斯噪声
beta_start (float) – 过程开始时“beta”参数的值（参见论文）
beta_end (float) – 过程结束时“beta”参数的值
sample_min (float)
sample_max (float) – 用于裁剪输出。
show_progress (bool) – 是否在推理过程中显示进度

Example

>>> from speechbrain.nnet.unet import UNetModel
>>> unet = UNetModel(
...     in_channels=1,
...     model_channels=16,
...     norm_num_groups=4,
...     out_channels=1,
...     num_res_blocks=1,
...     attention_resolutions=[]
... )
>>> diff = DenoisingDiffusion(
...     model=unet,
...     timesteps=5
... )
>>> x = torch.randn(4, 1, 64, 64)
>>> pred, noise, noisy_sample = diff.train_sample(x)
>>> pred.shape
torch.Size([4, 1, 64, 64])
>>> noise.shape
torch.Size([4, 1, 64, 64])
>>> noisy_sample.shape
torch.Size([4, 1, 64, 64])
>>> sample = diff.sample((2, 1, 64, 64))
>>> sample.shape
torch.Size([2, 1, 64, 64])

compute_coefficients()[source]: 计算扩散系数（alphas 和 betas）

distort(x, noise=None, timesteps=None, **kwargs)[source]

在前向扩散过程中向样本添加噪声，

Parameters:

x (torch.Tensor) – 一个2维或更高维的数据样本，其中第一个维度表示批次
噪声 (torch.Tensor) – 要添加的噪声
timesteps (torch.Tensor) – 一个长度为x中批次数的1-D整数张量，其中每个条目对应于批次的timestep编号。如果省略，timesteps将被随机采样
**kwargs (dict) – 传递给底层模型的参数。

Returns:

result – 一个与x维度相同的张量

Return type:

torch.Tensor

sample(shape, **kwargs)[source]

生成由count参数指示的样本数量

Parameters:

shape (可枚举) – 要生成的样本的形状
**kwargs (dict) – 传递给底层模型的参数。

Returns:

result – 生成的样本

Return type:

torch.Tensor

sample_step(sample, timestep, **kwargs)[source]

处理采样过程中的单个时间步长

Parameters:

样本 (torch.Tensor) – 下一个时间步的样本
timestep (int) – 时间步数
**kwargs (dict) – 传递给底层模型的参数。

Returns:

predicted_sample – 预测样本（通过一步去噪）

Return type:

torch.Tensor

class speechbrain.nnet.diffusion.LatentDiffusion(autoencoder, diffusion, latent_downsample_factor=None, latent_pad_dim=1)[source]

基础：Module

一个潜在扩散包装器。潜在扩散是将去噪扩散应用于潜在空间而不是原始数据空间。

Parameters:

autoencoder (speechbrain.nnet.autoencoders.Autoencoder) – 一个将原始空间转换为潜在空间的自动编码器
扩散 (speechbrain.nnet.diffusion.Diffuser) – 一个扩散包装器
latent_downsample_factor (int) – 潜在空间维度需要被整除的因子。如果扩散包装器的基础模型基于类似UNet的架构，其中输入逐渐被下采样和上采样，因子为二，这将非常有用。
latent_pad_dim (int|list[int]) – 潜在空间将沿其填充的维度

Example

>>> import torch
>>> from torch import nn
>>> from speechbrain.nnet.CNN import Conv2d
>>> from speechbrain.nnet.autoencoders import NormalizingAutoencoder
>>> from speechbrain.nnet.unet import UNetModel

设置一个简单的自动编码器（一个真正的自动编码器将是一个深度神经网络）

>>> ae_enc = Conv2d(
...     kernel_size=3,
...     stride=4,
...     in_channels=1,
...     out_channels=1,
...     skip_transpose=True,
... )
>>> ae_dec = nn.ConvTranspose2d(
...     kernel_size=3,
...     stride=4,
...     in_channels=1,
...     out_channels=1,
...     output_padding=1
... )
>>> ae = NormalizingAutoencoder(
...     encoder=ae_enc,
...     decoder=ae_dec,
... )

使用UNet架构构建扩散模型

>>> unet = UNetModel(
...     in_channels=1,
...     model_channels=16,
...     norm_num_groups=4,
...     out_channels=1,
...     num_res_blocks=1,
...     attention_resolutions=[]
... )
>>> diff = DenoisingDiffusion(
...     model=unet,
...     timesteps=5
... )
>>> latent_diff = LatentDiffusion(
...     autoencoder=ae,
...     diffusion=diff,
...     latent_downsample_factor=4,
...     latent_pad_dim=2
... )
>>> x = torch.randn(4, 1, 64, 64)
>>> latent_sample = latent_diff.train_sample_latent(x)
>>> diff_sample, ae_sample = latent_sample
>>> pred, noise, noisy_sample = diff_sample
>>> pred.shape
torch.Size([4, 1, 16, 16])
>>> noise.shape
torch.Size([4, 1, 16, 16])
>>> noisy_sample.shape
torch.Size([4, 1, 16, 16])
>>> ae_sample.latent.shape
torch.Size([4, 1, 16, 16])

创建一些样本（给定的形状应该是潜在空间的形状）

>>> sample = latent_diff.sample((2, 1, 16, 16))
>>> sample.shape
torch.Size([2, 1, 64, 64])

train_sample(x, **kwargs)[source]

为训练循环创建一个带有相应目标的样本

Parameters:

x (torch.Tensor) – 原始数据样本
**kwargs (dict) – 传递给底层模型的参数。

Returns:

pred (torch.Tensor) – 模型输出的预测噪声
noise (torch.Tensor) – 被应用的噪声
noisy_sample – 应用了噪声的样本

train_sample_latent(x, **kwargs)[source]

返回一个带有自动编码器输出的训练样本 - 可用于联合训练扩散模型和自动编码器

Parameters:

x (torch.Tensor) – 原始数据样本
**kwargs (dict) – 传递给底层模型的参数。

Returns:

训练样本。

Return type:

LatentDiffusionTrainSample

distort(x)[source]

在前向扩散过程中向样本添加噪声，

Parameters:: x (torch.Tensor) – 一个2维或更高维的数据样本，其中第一个维度表示批次
Returns:: result – 一个与x维度相同的张量
Return type:: torch.Tensor

sample(shape)[source]

从扩散模型中获取样本

Parameters:: 形状 (torch.Tensor)
Returns:: sample – 指定形状的样本
Return type:: torch.Tensor

speechbrain.nnet.diffusion.sample_timesteps(x, num_timesteps)[source]

返回一个随机的时间步样本作为一维张量（仅一个维度）

Parameters:

x (torch.Tensor) – 任意维度的样本张量
num_timesteps (int) – 总的时间步数

Return type:

时间戳的随机样本。

class speechbrain.nnet.diffusion.GaussianNoise(*args, **kwargs)[source]

基础：Module

添加普通高斯噪声

forward(sample, **kwargs)[source]

前向传播

Parameters:

样本 (原始样本)
**kwargs (dict) – 传递给底层模型的参数。

Return type:

样本形状中的噪声。

class speechbrain.nnet.diffusion.LengthMaskedGaussianNoise(length_dim=1)[source]

基础：Module

对填充样本应用高斯噪声。不会对作为填充部分的位置添加噪声。

Parameters:: length_dim (int) – 应用长度的时间维度。

forward(sample, length=None, **kwargs)[source]

创建高斯噪声。如果提供了长度张量，则不会在填充位置添加噪声。

Parameters:

样本 (torch.Tensor) – 一批数据
length (torch.Tensor) – 相对长度
**kwargs (dict) – 传递给底层模型的参数。

Return type:

样本形状的高斯噪声。

class speechbrain.nnet.diffusion.DiffusionTrainSample(pred, noise, noisy_sample)

基础：tuple

noise: 字段编号1的别名

noisy_sample: 字段编号2的别名

pred: 字段编号 0 的别名

class speechbrain.nnet.diffusion.LatentDiffusionTrainSample(diffusion, autoencoder)

基础：tuple

autoencoder: 字段编号1的别名

diffusion: 字段编号 0 的别名