| 147 | return out |
| 148 | |
| 149 | class ComplexConvTranspose2d(nn.Module): |
| 150 | |
| 151 | def __init__( |
| 152 | self, |
| 153 | in_channels, |
| 154 | out_channels, |
| 155 | kernel_size=(1,1), |
| 156 | stride=(1,1), |
| 157 | padding=(0,0), |
| 158 | output_padding=(0,0), |
| 159 | causal=False, |
| 160 | complex_axis=1, |
| 161 | groups=1, |
| 162 | dilation=1, |
| 163 | ): |
| 164 | ''' |
| 165 | in_channels: real+imag |
| 166 | out_channels: real+imag |
| 167 | ''' |
| 168 | super(ComplexConvTranspose2d, self).__init__() |
| 169 | self.in_channels = in_channels//2 |
| 170 | self.out_channels = out_channels//2 |
| 171 | self.kernel_size = kernel_size |
| 172 | self.stride = stride |
| 173 | self.padding = padding |
| 174 | self.output_padding=output_padding |
| 175 | self.groups = groups |
| 176 | self.dilation = dilation |
| 177 | |
| 178 | self.real_conv = nn.ConvTranspose2d(self.in_channels, self.out_channels,kernel_size, self.stride,padding=self.padding,output_padding=output_padding, dilation=self.dilation,groups=self.groups) |
| 179 | self.imag_conv = nn.ConvTranspose2d(self.in_channels, self.out_channels,kernel_size, self.stride,padding=self.padding,output_padding=output_padding, dilation=self.dilation,groups=self.groups) |
| 180 | self.complex_axis=complex_axis |
| 181 | |
| 182 | nn.init.normal_(self.real_conv.weight,std=0.05) |
| 183 | nn.init.normal_(self.imag_conv.weight,std=0.05) |
| 184 | nn.init.constant_(self.real_conv.bias,0.) |
| 185 | nn.init.constant_(self.imag_conv.bias,0.) |
| 186 | |
| 187 | def forward(self,inputs): |
| 188 | |
| 189 | if isinstance(inputs, torch.Tensor): |
| 190 | real,imag = torch.chunk(inputs, 2, self.complex_axis) |
| 191 | elif isinstance(inputs, tuple) or isinstance(inputs, list): |
| 192 | real = inputs[0] |
| 193 | imag = inputs[1] |
| 194 | if self.complex_axis == 0: |
| 195 | real = self.real_conv(inputs) |
| 196 | imag = self.imag_conv(inputs) |
| 197 | real2real,imag2real = torch.chunk(real,2, self.complex_axis) |
| 198 | real2imag,imag2imag = torch.chunk(imag,2, self.complex_axis) |
| 199 | |
| 200 | else: |
| 201 | if isinstance(inputs, torch.Tensor): |
| 202 | real,imag = torch.chunk(inputs, 2, self.complex_axis) |
| 203 | |
| 204 | real2real = self.real_conv(real,) |
| 205 | imag2imag = self.imag_conv(imag,) |
| 206 | |