Upsampling network module.
| 59 | |
| 60 | |
| 61 | class UpsampleNetwork(torch.nn.Module): |
| 62 | """Upsampling network module.""" |
| 63 | |
| 64 | def __init__(self, |
| 65 | upsample_scales, |
| 66 | nonlinear_activation=None, |
| 67 | nonlinear_activation_params={}, |
| 68 | interpolate_mode="nearest", |
| 69 | freq_axis_kernel_size=1, |
| 70 | use_causal_conv=False, |
| 71 | ): |
| 72 | """Initialize upsampling network module. |
| 73 | |
| 74 | Args: |
| 75 | upsample_scales (list): List of upsampling scales. |
| 76 | nonlinear_activation (str): Activation function name. |
| 77 | nonlinear_activation_params (dict): Arguments for specified activation function. |
| 78 | interpolate_mode (str): Interpolation mode. |
| 79 | freq_axis_kernel_size (int): Kernel size in the direction of frequency axis. |
| 80 | |
| 81 | """ |
| 82 | super(UpsampleNetwork, self).__init__() |
| 83 | self.use_causal_conv = use_causal_conv |
| 84 | self.up_layers = torch.nn.ModuleList() |
| 85 | for scale in upsample_scales: |
| 86 | # interpolation layer |
| 87 | stretch = Stretch2d(scale, 1, interpolate_mode) |
| 88 | self.up_layers += [stretch] |
| 89 | |
| 90 | # conv layer |
| 91 | assert (freq_axis_kernel_size - 1) % 2 == 0, "Not support even number freq axis kernel size." |
| 92 | freq_axis_padding = (freq_axis_kernel_size - 1) // 2 |
| 93 | kernel_size = (freq_axis_kernel_size, scale * 2 + 1) |
| 94 | if use_causal_conv: |
| 95 | padding = (freq_axis_padding, scale * 2) |
| 96 | else: |
| 97 | padding = (freq_axis_padding, scale) |
| 98 | conv = Conv2d(1, 1, kernel_size=kernel_size, padding=padding, bias=False) |
| 99 | self.up_layers += [conv] |
| 100 | |
| 101 | # nonlinear |
| 102 | if nonlinear_activation is not None: |
| 103 | nonlinear = getattr(torch.nn, nonlinear_activation)(**nonlinear_activation_params) |
| 104 | self.up_layers += [nonlinear] |
| 105 | |
| 106 | def forward(self, c): |
| 107 | """Calculate forward propagation. |
| 108 | |
| 109 | Args: |
| 110 | c : Input tensor (B, C, T). |
| 111 | |
| 112 | Returns: |
| 113 | Tensor: Upsampled tensor (B, C, T'), where T' = T * prod(upsample_scales). |
| 114 | |
| 115 | """ |
| 116 | c = c.unsqueeze(1) # (B, 1, C, T) |
| 117 | for f in self.up_layers: |
| 118 | if self.use_causal_conv and isinstance(f, Conv2d): |