Spaces:
Building
Building
File size: 3,705 Bytes
f6b56a2 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 |
import torch
from gyraudio.audio_separation.architecture.model import SeparationModel
from gyraudio.audio_separation.architecture.building_block import FilterBank
from typing import Optional
class TransformerModel(SeparationModel):
"""Transformer base model
=========================
- Embed signal with a filter bank
- No positional encoding (Potential =add/concatenate positional encoding)
- `nlayers` * transformer blocks
"""
def __init__(self,
nhead: int = 8, # H
nlayers: int = 4, # L
k_size=5,
embedding_dim: int = 64, # D
ch_in: int = 1,
ch_out: int = 1,
dropout: float = 0., # dr
positional_encoding: str = None
) -> None:
"""Transformer base model
Args:
nhead (int): number of heads in each of the MHA models
embedding_dim (int): D number of channels in the audio embeddings
= output of the filter bank
assume `embedding_dim` = `h_dim`
h_dim is the hidden dimension of the model.
nlayers (int): number of nn.TransformerEncoderLayer in nn.TransformerEncoder
dropout (float, optional): dropout value. Defaults to 0.
"""
super().__init__()
self.model_type = "Transformer"
h_dim = embedding_dim # use the same embedding & hidden dimensions
self.encoder = FilterBank(ch_in, embedding_dim, k_size=k_size)
if positional_encoding is None:
self.pos_encoder = torch.nn.Identity()
else:
raise NotImplementedError(
f"Unknown positional encoding {positional_encoding} - should be add/concat in future")
# self.pos_encoder = PositionalEncoding(h_dim, dropout=dropout)
encoder_layers = torch.nn.TransformerEncoderLayer(
d_model=h_dim, # input dimension to the transformer encoder layer
nhead=nhead, # number of heads for MHA (Multi-head attention)
dim_feedforward=h_dim, # output dimension of the MLP on top of the transformer.
dropout=dropout,
batch_first=True
) # we assume h_dim = d_model = dim_feedforward
self.transformer_encoder = torch.nn.TransformerEncoder(
encoder_layers,
num_layers=nlayers
)
self.h_dim = h_dim
self.target_modality_conv = torch.nn.Conv1d(h_dim, ch_out, 1) # conv1x1 channel mixer
# Note: we could finish with a few residual conv blocks... this is pure signal processing
def forward(
self, src: torch.LongTensor,
src_mask: Optional[torch.FloatTensor] = None
) -> torch.FloatTensor:
"""Embdeddings, positional encoders, go trough `nlayers` of residual {multi (`nhead`) attention heads + MLP}.
Args:
src (torch.LongTensor): [N, 1, T] audio signal
Returns:
torch.FloatTensor: separated signal [N, 1, T]
"""
src = self.encoder(src) # [N, 1, T] -> [N, D, T]
src = src.transpose(-1, -2) # [N, D, T] -> [N, T, D] # Transformer expects (batch N, seq "T", features "D")
src = self.pos_encoder(src) # -> [N, T, D] - add positional encoding
output = self.transformer_encoder(src, mask=src_mask) # -> [N, T, D]
output = output.transpose(-1, -2) # -> [N, D, T]
output = self.target_modality_conv(output) # -> [N, 1, T]
return output, None
if __name__ == "__main__":
model = TransformerModel()
inp = torch.rand(2, 1, 2048)
out = model(inp)
print(model)
print(out[0].shape)
|