in_channels: real+imag out_channels: real+imag kernel_size : input [B,C,D,T] kernel size in [D,T] padding : input [B,C,D,T] padding in [D,T] causal: if causal, will padding time dimension's left side, otherwise b
(
self,
in_channels,
out_channels,
kernel_size=(1,1),
stride=(1,1),
padding=(0,0),
dilation=1,
groups = 1,
causal=True,
complex_axis=1,
)
| 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): |