mirror of
https://github.com/w-okada/voice-changer.git
synced 2025-01-24 14:05:00 +03:00
380 lines
14 KiB
Python
380 lines
14 KiB
Python
import torch
|
|
|
|
from torch import nn
|
|
import math
|
|
from functools import partial
|
|
from einops import rearrange, repeat
|
|
|
|
from local_attention import LocalAttention
|
|
import torch.nn.functional as F
|
|
#import fast_transformers.causal_product.causal_product_cuda
|
|
|
|
def softmax_kernel(data, *, projection_matrix, is_query, normalize_data=True, eps=1e-4, device = None):
|
|
b, h, *_ = data.shape
|
|
# (batch size, head, length, model_dim)
|
|
|
|
# normalize model dim
|
|
data_normalizer = (data.shape[-1] ** -0.25) if normalize_data else 1.
|
|
|
|
# what is ration?, projection_matrix.shape[0] --> 266
|
|
|
|
ratio = (projection_matrix.shape[0] ** -0.5)
|
|
|
|
projection = repeat(projection_matrix, 'j d -> b h j d', b = b, h = h)
|
|
projection = projection.type_as(data)
|
|
|
|
#data_dash = w^T x
|
|
data_dash = torch.einsum('...id,...jd->...ij', (data_normalizer * data), projection)
|
|
|
|
|
|
# diag_data = D**2
|
|
diag_data = data ** 2
|
|
diag_data = torch.sum(diag_data, dim=-1)
|
|
diag_data = (diag_data / 2.0) * (data_normalizer ** 2)
|
|
diag_data = diag_data.unsqueeze(dim=-1)
|
|
|
|
#print ()
|
|
if is_query:
|
|
data_dash = ratio * (
|
|
torch.exp(data_dash - diag_data -
|
|
torch.max(data_dash, dim=-1, keepdim=True).values) + eps)
|
|
else:
|
|
data_dash = ratio * (
|
|
torch.exp(data_dash - diag_data + eps))#- torch.max(data_dash)) + eps)
|
|
|
|
return data_dash.type_as(data)
|
|
|
|
def orthogonal_matrix_chunk(cols, qr_uniform_q = False, device = None):
|
|
unstructured_block = torch.randn((cols, cols), device = device)
|
|
q, r = torch.linalg.qr(unstructured_block.cpu(), mode='reduced')
|
|
q, r = map(lambda t: t.to(device), (q, r))
|
|
|
|
# proposed by @Parskatt
|
|
# to make sure Q is uniform https://arxiv.org/pdf/math-ph/0609050.pdf
|
|
if qr_uniform_q:
|
|
d = torch.diag(r, 0)
|
|
q *= d.sign()
|
|
return q.t()
|
|
def exists(val):
|
|
return val is not None
|
|
|
|
def empty(tensor):
|
|
return tensor.numel() == 0
|
|
|
|
def default(val, d):
|
|
return val if exists(val) else d
|
|
|
|
def cast_tuple(val):
|
|
return (val,) if not isinstance(val, tuple) else val
|
|
|
|
class PCmer(nn.Module):
|
|
"""The encoder that is used in the Transformer model."""
|
|
|
|
def __init__(self,
|
|
num_layers,
|
|
num_heads,
|
|
dim_model,
|
|
dim_keys,
|
|
dim_values,
|
|
residual_dropout,
|
|
attention_dropout):
|
|
super().__init__()
|
|
self.num_layers = num_layers
|
|
self.num_heads = num_heads
|
|
self.dim_model = dim_model
|
|
self.dim_values = dim_values
|
|
self.dim_keys = dim_keys
|
|
self.residual_dropout = residual_dropout
|
|
self.attention_dropout = attention_dropout
|
|
|
|
self._layers = nn.ModuleList([_EncoderLayer(self) for _ in range(num_layers)])
|
|
|
|
# METHODS ########################################################################################################
|
|
|
|
def forward(self, phone, mask=None):
|
|
|
|
# apply all layers to the input
|
|
for (i, layer) in enumerate(self._layers):
|
|
phone = layer(phone, mask)
|
|
# provide the final sequence
|
|
return phone
|
|
|
|
|
|
# ==================================================================================================================== #
|
|
# CLASS _ E N C O D E R L A Y E R #
|
|
# ==================================================================================================================== #
|
|
|
|
|
|
class _EncoderLayer(nn.Module):
|
|
"""One layer of the encoder.
|
|
|
|
Attributes:
|
|
attn: (:class:`mha.MultiHeadAttention`): The attention mechanism that is used to read the input sequence.
|
|
feed_forward (:class:`ffl.FeedForwardLayer`): The feed-forward layer on top of the attention mechanism.
|
|
"""
|
|
|
|
def __init__(self, parent: PCmer):
|
|
"""Creates a new instance of ``_EncoderLayer``.
|
|
|
|
Args:
|
|
parent (Encoder): The encoder that the layers is created for.
|
|
"""
|
|
super().__init__()
|
|
|
|
|
|
self.conformer = ConformerConvModule(parent.dim_model)
|
|
self.norm = nn.LayerNorm(parent.dim_model)
|
|
self.dropout = nn.Dropout(parent.residual_dropout)
|
|
|
|
# selfatt -> fastatt: performer!
|
|
self.attn = SelfAttention(dim = parent.dim_model,
|
|
heads = parent.num_heads,
|
|
causal = False)
|
|
|
|
# METHODS ########################################################################################################
|
|
|
|
def forward(self, phone, mask=None):
|
|
|
|
# compute attention sub-layer
|
|
phone = phone + (self.attn(self.norm(phone), mask=mask))
|
|
|
|
phone = phone + (self.conformer(phone))
|
|
|
|
return phone
|
|
|
|
def calc_same_padding(kernel_size):
|
|
pad = kernel_size // 2
|
|
return (pad, pad - (kernel_size + 1) % 2)
|
|
|
|
# helper classes
|
|
|
|
class Swish(nn.Module):
|
|
def forward(self, x):
|
|
return x * x.sigmoid()
|
|
|
|
class Transpose(nn.Module):
|
|
def __init__(self, dims):
|
|
super().__init__()
|
|
assert len(dims) == 2, 'dims must be a tuple of two dimensions'
|
|
self.dims = dims
|
|
|
|
def forward(self, x):
|
|
return x.transpose(*self.dims)
|
|
|
|
class GLU(nn.Module):
|
|
def __init__(self, dim):
|
|
super().__init__()
|
|
self.dim = dim
|
|
|
|
def forward(self, x):
|
|
out, gate = x.chunk(2, dim=self.dim)
|
|
return out * gate.sigmoid()
|
|
|
|
class DepthWiseConv1d(nn.Module):
|
|
def __init__(self, chan_in, chan_out, kernel_size, padding):
|
|
super().__init__()
|
|
self.padding = padding
|
|
self.conv = nn.Conv1d(chan_in, chan_out, kernel_size, groups = chan_in)
|
|
|
|
def forward(self, x):
|
|
x = F.pad(x, self.padding)
|
|
return self.conv(x)
|
|
|
|
class ConformerConvModule(nn.Module):
|
|
def __init__(
|
|
self,
|
|
dim,
|
|
causal = False,
|
|
expansion_factor = 2,
|
|
kernel_size = 31,
|
|
dropout = 0.):
|
|
super().__init__()
|
|
|
|
inner_dim = dim * expansion_factor
|
|
padding = calc_same_padding(kernel_size) if not causal else (kernel_size - 1, 0)
|
|
|
|
self.net = nn.Sequential(
|
|
nn.LayerNorm(dim),
|
|
Transpose((1, 2)),
|
|
nn.Conv1d(dim, inner_dim * 2, 1),
|
|
GLU(dim=1),
|
|
DepthWiseConv1d(inner_dim, inner_dim, kernel_size = kernel_size, padding = padding),
|
|
#nn.BatchNorm1d(inner_dim) if not causal else nn.Identity(),
|
|
Swish(),
|
|
nn.Conv1d(inner_dim, dim, 1),
|
|
Transpose((1, 2)),
|
|
nn.Dropout(dropout)
|
|
)
|
|
|
|
def forward(self, x):
|
|
return self.net(x)
|
|
|
|
def linear_attention(q, k, v):
|
|
if v is None:
|
|
#print (k.size(), q.size())
|
|
out = torch.einsum('...ed,...nd->...ne', k, q)
|
|
return out
|
|
|
|
else:
|
|
k_cumsum = k.sum(dim = -2)
|
|
#k_cumsum = k.sum(dim = -2)
|
|
D_inv = 1. / (torch.einsum('...nd,...d->...n', q, k_cumsum.type_as(q)) + 1e-8)
|
|
|
|
context = torch.einsum('...nd,...ne->...de', k, v)
|
|
#print ("TRUEEE: ", context.size(), q.size(), D_inv.size())
|
|
out = torch.einsum('...de,...nd,...n->...ne', context, q, D_inv)
|
|
return out
|
|
|
|
def gaussian_orthogonal_random_matrix(nb_rows, nb_columns, scaling = 0, qr_uniform_q = False, device = None):
|
|
nb_full_blocks = int(nb_rows / nb_columns)
|
|
#print (nb_full_blocks)
|
|
block_list = []
|
|
|
|
for _ in range(nb_full_blocks):
|
|
q = orthogonal_matrix_chunk(nb_columns, qr_uniform_q = qr_uniform_q, device = device)
|
|
block_list.append(q)
|
|
# block_list[n] is a orthogonal matrix ... (model_dim * model_dim)
|
|
#print (block_list[0].size(), torch.einsum('...nd,...nd->...n', block_list[0], torch.roll(block_list[0],1,1)))
|
|
#print (nb_rows, nb_full_blocks, nb_columns)
|
|
remaining_rows = nb_rows - nb_full_blocks * nb_columns
|
|
#print (remaining_rows)
|
|
if remaining_rows > 0:
|
|
q = orthogonal_matrix_chunk(nb_columns, qr_uniform_q = qr_uniform_q, device = device)
|
|
#print (q[:remaining_rows].size())
|
|
block_list.append(q[:remaining_rows])
|
|
|
|
final_matrix = torch.cat(block_list)
|
|
|
|
if scaling == 0:
|
|
multiplier = torch.randn((nb_rows, nb_columns), device = device).norm(dim = 1)
|
|
elif scaling == 1:
|
|
multiplier = math.sqrt((float(nb_columns))) * torch.ones((nb_rows,), device = device)
|
|
else:
|
|
raise ValueError(f'Invalid scaling {scaling}')
|
|
|
|
return torch.diag(multiplier) @ final_matrix
|
|
|
|
class FastAttention(nn.Module):
|
|
def __init__(self, dim_heads, nb_features = None, ortho_scaling = 0, causal = False, generalized_attention = False, kernel_fn = nn.ReLU(), qr_uniform_q = False, no_projection = False):
|
|
super().__init__()
|
|
nb_features = default(nb_features, int(dim_heads * math.log(dim_heads)))
|
|
|
|
self.dim_heads = dim_heads
|
|
self.nb_features = nb_features
|
|
self.ortho_scaling = ortho_scaling
|
|
|
|
self.create_projection = partial(gaussian_orthogonal_random_matrix, nb_rows = self.nb_features, nb_columns = dim_heads, scaling = ortho_scaling, qr_uniform_q = qr_uniform_q)
|
|
projection_matrix = self.create_projection()
|
|
self.register_buffer('projection_matrix', projection_matrix)
|
|
|
|
self.generalized_attention = generalized_attention
|
|
self.kernel_fn = kernel_fn
|
|
|
|
# if this is turned on, no projection will be used
|
|
# queries and keys will be softmax-ed as in the original efficient attention paper
|
|
self.no_projection = no_projection
|
|
|
|
self.causal = causal
|
|
if causal:
|
|
try:
|
|
import fast_transformers.causal_product.causal_product_cuda
|
|
self.causal_linear_fn = partial(causal_linear_attention)
|
|
except ImportError:
|
|
print('unable to import cuda code for auto-regressive Performer. will default to the memory inefficient non-cuda version')
|
|
self.causal_linear_fn = causal_linear_attention_noncuda
|
|
@torch.no_grad()
|
|
def redraw_projection_matrix(self):
|
|
projections = self.create_projection()
|
|
self.projection_matrix.copy_(projections)
|
|
del projections
|
|
|
|
def forward(self, q, k, v):
|
|
device = q.device
|
|
|
|
if self.no_projection:
|
|
q = q.softmax(dim = -1)
|
|
k = torch.exp(k) if self.causal else k.softmax(dim = -2)
|
|
|
|
elif self.generalized_attention:
|
|
create_kernel = partial(generalized_kernel, kernel_fn = self.kernel_fn, projection_matrix = self.projection_matrix, device = device)
|
|
q, k = map(create_kernel, (q, k))
|
|
|
|
else:
|
|
create_kernel = partial(softmax_kernel, projection_matrix = self.projection_matrix, device = device)
|
|
|
|
q = create_kernel(q, is_query = True)
|
|
k = create_kernel(k, is_query = False)
|
|
|
|
attn_fn = linear_attention if not self.causal else self.causal_linear_fn
|
|
if v is None:
|
|
out = attn_fn(q, k, None)
|
|
return out
|
|
else:
|
|
out = attn_fn(q, k, v)
|
|
return out
|
|
class SelfAttention(nn.Module):
|
|
def __init__(self, dim, causal = False, heads = 8, dim_head = 64, local_heads = 0, local_window_size = 256, nb_features = None, feature_redraw_interval = 1000, generalized_attention = False, kernel_fn = nn.ReLU(), qr_uniform_q = False, dropout = 0., no_projection = False):
|
|
super().__init__()
|
|
assert dim % heads == 0, 'dimension must be divisible by number of heads'
|
|
dim_head = default(dim_head, dim // heads)
|
|
inner_dim = dim_head * heads
|
|
self.fast_attention = FastAttention(dim_head, nb_features, causal = causal, generalized_attention = generalized_attention, kernel_fn = kernel_fn, qr_uniform_q = qr_uniform_q, no_projection = no_projection)
|
|
|
|
self.heads = heads
|
|
self.global_heads = heads - local_heads
|
|
self.local_attn = LocalAttention(window_size = local_window_size, causal = causal, autopad = True, dropout = dropout, look_forward = int(not causal), rel_pos_emb_config = (dim_head, local_heads)) if local_heads > 0 else None
|
|
|
|
#print (heads, nb_features, dim_head)
|
|
#name_embedding = torch.zeros(110, heads, dim_head, dim_head)
|
|
#self.name_embedding = nn.Parameter(name_embedding, requires_grad=True)
|
|
|
|
|
|
self.to_q = nn.Linear(dim, inner_dim)
|
|
self.to_k = nn.Linear(dim, inner_dim)
|
|
self.to_v = nn.Linear(dim, inner_dim)
|
|
self.to_out = nn.Linear(inner_dim, dim)
|
|
self.dropout = nn.Dropout(dropout)
|
|
|
|
@torch.no_grad()
|
|
def redraw_projection_matrix(self):
|
|
self.fast_attention.redraw_projection_matrix()
|
|
#torch.nn.init.zeros_(self.name_embedding)
|
|
#print (torch.sum(self.name_embedding))
|
|
def forward(self, x, context = None, mask = None, context_mask = None, name=None, inference=False, **kwargs):
|
|
b, n, _, h, gh = *x.shape, self.heads, self.global_heads
|
|
|
|
cross_attend = exists(context)
|
|
|
|
context = default(context, x)
|
|
context_mask = default(context_mask, mask) if not cross_attend else context_mask
|
|
#print (torch.sum(self.name_embedding))
|
|
q, k, v = self.to_q(x), self.to_k(context), self.to_v(context)
|
|
|
|
q, k, v = map(lambda t: rearrange(t, 'b n (h d) -> b h n d', h = h), (q, k, v))
|
|
(q, lq), (k, lk), (v, lv) = map(lambda t: (t[:, :gh], t[:, gh:]), (q, k, v))
|
|
|
|
attn_outs = []
|
|
#print (name)
|
|
#print (self.name_embedding[name].size())
|
|
if not empty(q):
|
|
if exists(context_mask):
|
|
global_mask = context_mask[:, None, :, None]
|
|
v.masked_fill_(~global_mask, 0.)
|
|
if cross_attend:
|
|
pass
|
|
#print (torch.sum(self.name_embedding))
|
|
#out = self.fast_attention(q,self.name_embedding[name],None)
|
|
#print (torch.sum(self.name_embedding[...,-1:]))
|
|
else:
|
|
out = self.fast_attention(q, k, v)
|
|
attn_outs.append(out)
|
|
|
|
if not empty(lq):
|
|
assert not cross_attend, 'local attention is not compatible with cross attention'
|
|
out = self.local_attn(lq, lk, lv, input_mask = mask)
|
|
attn_outs.append(out)
|
|
|
|
out = torch.cat(attn_outs, dim = 1)
|
|
out = rearrange(out, 'b h n d -> b n (h d)')
|
|
out = self.to_out(out)
|
|
return self.dropout(out) |