| 38 | |
| 39 | |
| 40 | class DownSampleModule(BaseModule): |
| 41 | def __init__(self, nc_in, nf, use_bias, norm, conv_by, conv_type): |
| 42 | super().__init__(conv_type) |
| 43 | self.conv1 = self.ConvBlock( |
| 44 | nc_in, nf * 1, kernel_size=(3, 5, 5), stride=1, |
| 45 | padding=1, bias=use_bias, norm=norm, conv_by=conv_by) |
| 46 | |
| 47 | # Downsample 1 |
| 48 | self.conv2 = self.ConvBlock( |
| 49 | nf * 1, nf * 2, kernel_size=(3, 4, 4), stride=(1, 2, 2), |
| 50 | padding=(1, 2, 2), bias=use_bias, norm=norm, conv_by=conv_by) |
| 51 | self.conv3 = self.ConvBlock( |
| 52 | nf * 2, nf * 2, kernel_size=(3, 3, 3), stride=(1, 1, 1), |
| 53 | padding=1, bias=use_bias, norm=norm, conv_by=conv_by) |
| 54 | # Downsample 2 |
| 55 | self.conv4 = self.ConvBlock( |
| 56 | nf * 2, nf * 4, kernel_size=(3, 4, 4), stride=(1, 2, 2), |
| 57 | padding=1, bias=use_bias, norm=norm, conv_by=conv_by) |
| 58 | self.conv5 = self.ConvBlock( |
| 59 | nf * 4, nf * 4, kernel_size=(3, 3, 3), stride=(1, 1, 1), |
| 60 | padding=1, bias=use_bias, norm=norm, conv_by=conv_by) |
| 61 | self.conv6 = self.ConvBlock( |
| 62 | nf * 4, nf * 4, kernel_size=(3, 3, 3), stride=(1, 1, 1), |
| 63 | padding=1, bias=use_bias, norm=norm, conv_by=conv_by) |
| 64 | |
| 65 | # Dilated Convolutions |
| 66 | self.dilated_conv1 = self.ConvBlock( |
| 67 | nf * 4, nf * 4, kernel_size=(3, 3, 3), stride=(1, 1, 1), |
| 68 | padding=-1, bias=use_bias, norm=norm, conv_by=conv_by, dilation=(1, 2, 2)) |
| 69 | self.dilated_conv2 = self.ConvBlock( |
| 70 | nf * 4, nf * 4, kernel_size=(3, 3, 3), stride=(1, 1, 1), |
| 71 | padding=-1, bias=use_bias, norm=norm, conv_by=conv_by, dilation=(1, 4, 4)) |
| 72 | self.dilated_conv3 = self.ConvBlock( |
| 73 | nf * 4, nf * 4, kernel_size=(3, 3, 3), stride=(1, 1, 1), |
| 74 | padding=-1, bias=use_bias, norm=norm, conv_by=conv_by, dilation=(1, 8, 8)) |
| 75 | self.dilated_conv4 = self.ConvBlock( |
| 76 | nf * 4, nf * 4, kernel_size=(3, 3, 3), stride=(1, 1, 1), |
| 77 | padding=-1, bias=use_bias, norm=norm, conv_by=conv_by, dilation=(1, 16, 16)) |
| 78 | self.conv7 = self.ConvBlock( |
| 79 | nf * 4, nf * 4, kernel_size=(3, 3, 3), stride=(1, 1, 1), padding=1, |
| 80 | bias=use_bias, norm=norm, conv_by=conv_by) |
| 81 | self.conv8 = self.ConvBlock( |
| 82 | nf * 4, nf * 4, kernel_size=(3, 3, 3), stride=(1, 1, 1), padding=1, |
| 83 | bias=use_bias, norm=norm, conv_by=conv_by) |
| 84 | |
| 85 | def forward(self, inp): |
| 86 | c1 = self.conv1(inp) |
| 87 | c2 = self.conv2(c1) |
| 88 | c3 = self.conv3(c2) |
| 89 | c4 = self.conv4(c3) |
| 90 | c5 = self.conv5(c4) |
| 91 | c6 = self.conv6(c5) |
| 92 | |
| 93 | a1 = self.dilated_conv1(c6) |
| 94 | a2 = self.dilated_conv2(a1) |
| 95 | a3 = self.dilated_conv3(a2) |
| 96 | a4 = self.dilated_conv4(a3) |
| 97 | |