| 77 | return outputs |
| 78 | |
| 79 | class ComplexConv2d(nn.Module): |
| 80 | |
| 81 | def __init__( |
| 82 | self, |
| 83 | in_channels, |
| 84 | out_channels, |
| 85 | kernel_size=(1,1), |
| 86 | stride=(1,1), |
| 87 | padding=(0,0), |
| 88 | dilation=1, |
| 89 | groups = 1, |
| 90 | causal=True, |
| 91 | complex_axis=1, |
| 92 | ): |
| 93 | ''' |
| 94 | in_channels: real+imag |
| 95 | out_channels: real+imag |
| 96 | kernel_size : input [B,C,D,T] kernel size in [D,T] |
| 97 | padding : input [B,C,D,T] padding in [D,T] |
| 98 | causal: if causal, will padding time dimension's left side, |
| 99 | otherwise both |
| 100 | |
| 101 | ''' |
| 102 | super(ComplexConv2d, self).__init__() |
| 103 | self.in_channels = in_channels//2 |
| 104 | self.out_channels = out_channels//2 |
| 105 | self.kernel_size = kernel_size |
| 106 | self.stride = stride |
| 107 | self.padding = padding |
| 108 | self.causal = causal |
| 109 | self.groups = groups |
| 110 | self.dilation = dilation |
| 111 | self.complex_axis=complex_axis |
| 112 | self.real_conv = nn.Conv2d(self.in_channels, self.out_channels, kernel_size, self.stride,padding=[self.padding[0],0],dilation=self.dilation, groups=self.groups) |
| 113 | self.imag_conv = nn.Conv2d(self.in_channels, self.out_channels, kernel_size, self.stride,padding=[self.padding[0],0],dilation=self.dilation, groups=self.groups) |
| 114 | |
| 115 | nn.init.normal_(self.real_conv.weight.data,std=0.05) |
| 116 | nn.init.normal_(self.imag_conv.weight.data,std=0.05) |
| 117 | nn.init.constant_(self.real_conv.bias,0.) |
| 118 | nn.init.constant_(self.imag_conv.bias,0.) |
| 119 | |
| 120 | |
| 121 | def forward(self,inputs): |
| 122 | if self.padding[1] != 0 and self.causal: |
| 123 | inputs = F.pad(inputs,[self.padding[1], 0,0,0]) |
| 124 | else: |
| 125 | inputs = F.pad(inputs,[self.padding[1], self.padding[1],0,0]) |
| 126 | |
| 127 | if self.complex_axis == 0: |
| 128 | real = self.real_conv(inputs) |
| 129 | imag = self.imag_conv(inputs) |
| 130 | real2real,imag2real = torch.chunk(real,2, self.complex_axis) |
| 131 | real2imag,imag2imag = torch.chunk(imag,2, self.complex_axis) |
| 132 | |
| 133 | else: |
| 134 | if isinstance(inputs, torch.Tensor): |
| 135 | real,imag = torch.chunk(inputs, 2, self.complex_axis) |
| 136 | |