| 168 | ## Version 2 |
| 169 | class StreamConvTranspose2d(nn.Module): |
| 170 | def __init__(self, |
| 171 | in_channels: int, |
| 172 | out_channels: int, |
| 173 | kernel_size: Union[int, Tuple[int, int]], |
| 174 | stride: Union[int, Tuple[int, int]] = 1, |
| 175 | padding: Union[str, int, Tuple[int, int]] = 0, |
| 176 | dilation: Union[int, Tuple[int, int]] = 1, |
| 177 | groups: int = 1, |
| 178 | bias: bool = True, |
| 179 | *args, **kargs): |
| 180 | super().__init__(*args, **kargs) |
| 181 | """ |
| 182 | kernel_size = [T_size, F_size] by default |
| 183 | stride = [T_stride, F_stride], and T_stride == 1 |
| 184 | """ |
| 185 | self.in_channels = in_channels |
| 186 | self.out_channels = out_channels |
| 187 | if type(kernel_size) is int: |
| 188 | self.T_size = kernel_size |
| 189 | self.F_size = kernel_size |
| 190 | elif type(kernel_size) in [list, tuple]: |
| 191 | self.T_size, self.F_size = kernel_size |
| 192 | else: |
| 193 | raise ValueError('Invalid kernel size.') |
| 194 | |
| 195 | if type(stride) is int: |
| 196 | self.T_stride = stride |
| 197 | self.F_stride = stride |
| 198 | elif type(stride) in [list, tuple]: |
| 199 | self.T_stride, self.F_stride = stride |
| 200 | else: |
| 201 | raise ValueError('Invalid stride size.') |
| 202 | |
| 203 | assert self.T_stride == 1 |
| 204 | |
| 205 | if type(padding) is int: |
| 206 | self.T_pad = padding |
| 207 | self.F_pad = padding |
| 208 | elif type(padding) in [list, tuple]: |
| 209 | self.T_pad, self.F_pad = padding |
| 210 | else: |
| 211 | raise ValueError('Invalid padding size.') |
| 212 | assert(self.T_pad == 0) |
| 213 | |
| 214 | if type(dilation) is int: |
| 215 | self.T_dilation = dilation |
| 216 | self.F_dilation = dilation |
| 217 | elif type(dilation) in [list, tuple]: |
| 218 | self.T_dilation, self.F_dilation = dilation |
| 219 | else: |
| 220 | raise ValueError('Invalid dilation size.') |
| 221 | |
| 222 | # Implementing ConvTranspose2d using Conv2d with weight-time reversal. |
| 223 | self.ConvTranspose2d = nn.Conv2d(in_channels = in_channels, |
| 224 | out_channels = out_channels, |
| 225 | kernel_size = kernel_size, |
| 226 | stride = (self.T_stride, 1), # An additional upsampling will be used in forward, if F_stride != 1 |
| 227 | padding = (self.T_pad, 0), # An additional padding will be used in forward, if F_pad != 0 |