Spaces:
Sleeping
Sleeping
File size: 3,236 Bytes
35c1cfd |
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 |
# MIT License
#
# Copyright 2023 ByteDance Inc.
#
# Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the “Software”),
# to deal in the Software without restriction, including without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense,
# and/or sell copies of the Software, and to permit persons to whom the Software is furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED “AS IS”, WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS
# IN THE SOFTWARE.
from torch import nn
from einops import rearrange
class Res2dModule(nn.Module):
def __init__(self, idim, odim, stride=(2, 2)):
super(Res2dModule, self).__init__()
self.conv1 = nn.Conv2d(idim, odim, 3, padding=1, stride=stride)
self.bn1 = nn.BatchNorm2d(odim)
self.conv2 = nn.Conv2d(odim, odim, 3, padding=1)
self.bn2 = nn.BatchNorm2d(odim)
self.relu = nn.ReLU()
# residual
self.diff = False
if (idim != odim) or (stride[0] > 1):
self.conv3 = nn.Conv2d(idim, odim, 3, padding=1, stride=stride)
self.bn3 = nn.BatchNorm2d(odim)
self.diff = True
def forward(self, x):
out = self.bn2(self.conv2(self.relu(self.bn1(self.conv1(x)))))
if self.diff:
x = self.bn3(self.conv3(x))
out = x + out
out = self.relu(out)
return out
class Conv2dSubsampling(nn.Module):
"""Convolutional 2D subsampling (to 1/4 length).
Args:
idim (int): Input dimension.
hdim (int): Hidden dimension.
odim (int): Output dimension.
strides (list): Sizes of strides.
n_bands (int): Number of frequency bands.
"""
def __init__(self, idim, hdim, odim, strides=[2, 2], n_bands=64):
"""Construct an Conv2dSubsampling object."""
super(Conv2dSubsampling, self).__init__()
self.conv = nn.Sequential(
Res2dModule(idim, hdim, (2, strides[0])),
Res2dModule(hdim, hdim, (2, strides[1])),
)
self.linear = nn.Linear(hdim * n_bands // 2 // 2, odim)
def forward(self, x):
"""Subsample x.
Args:
x (torch.Tensor): Input tensor (#batch, idim, time).
Returns:
torch.Tensor: Subsampled tensor (#batch, time', odim),
where time' = time // 4.
"""
if x.dim() == 3:
x = x.unsqueeze(1) # (b, c, f, t)
x = self.conv(x)
x = rearrange(x, "b c f t -> b t (c f)")
x = self.linear(x)
return x
|