| import torch |
| from torch import nn |
|
|
| class PatchEmbed(nn.Module): |
| """Convert a band spectrogram to patch embeddings.""" |
| def __init__(self, band_width, shift_size, in_chans=1, embed_dim=768): |
| super().__init__() |
| |
| self.band_width = band_width |
| self.shift_size = shift_size |
| self.in_chans = in_chans |
| self.embed_dim = embed_dim |
| |
| |
| self.proj = nn.Conv2d(in_chans, embed_dim, kernel_size=(band_width, shift_size), stride=(band_width, shift_size)) |
| |
| def forward(self, x): |
| """x: (B, band_width, T) or (B, 1, band_width, T).""" |
| if x.dim() == 3: |
| x = x.unsqueeze(1) |
|
|
| assert x.dim() == 4, f"Input shape must be (B, band_width, T), but got {x.shape}" |
| B, C, H, W = x.shape |
| |
| assert H == self.band_width, f"Input height ({H}) doesn't match band_width ({self.band_width})" |
|
|
| num_patches = W // self.shift_size |
|
|
| if W % self.shift_size != 0: |
| padding_width = self.shift_size - (W % self.shift_size) |
| x = torch.nn.functional.pad(x, (0, padding_width, 0, 0), mode='constant', value=0) |
| num_patches += 1 |
|
|
| patches = self.proj(x) |
|
|
| patches = patches.squeeze(2).transpose(1, 2) |
| |
| return patches |
|
|
