(self, block=BasicBlock, layers=[3, 4, 6, 6, 3], strides=[(2,2),(1,2),(2,2),(1,2),(2,2)])
| 43 | |
| 44 | class PSPEncoder(nn.Module): |
| 45 | def __init__(self, block=BasicBlock, layers=[3, 4, 6, 6, 3], strides=[(2,2),(1,2),(2,2),(1,2),(2,2)]): |
| 46 | self.inplanes = 32 |
| 47 | super(PSPEncoder, self).__init__() |
| 48 | self.conv1 = nn.Conv2d(3, 32, kernel_size=3, stride=1, padding=1, |
| 49 | bias=False) |
| 50 | self.relu = nn.LeakyReLU(0.2) |
| 51 | |
| 52 | feature_out_dim = 256 |
| 53 | self.layer1 = self._make_layer(block, 32, layers[0], stride=strides[0]) |
| 54 | self.layer2 = self._make_layer(block, 64, layers[1], stride=strides[1]) |
| 55 | self.layer3 = self._make_layer(block, 128, layers[2], stride=strides[2]) |
| 56 | self.layer4 = self._make_layer(block, 256, layers[3], stride=strides[3]) |
| 57 | self.layer5 = self._make_layer(block, 512, layers[4], stride=strides[4]) |
| 58 | |
| 59 | self.layer512_to_outdim = nn.Sequential( |
| 60 | nn.Conv2d(512, feature_out_dim, kernel_size=1, stride=1, bias=False), |
| 61 | nn.LeakyReLU(0.2) |
| 62 | ) |
| 63 | self.layer256_to_512 = nn.Sequential( |
| 64 | nn.Conv2d(256, 512, kernel_size=1, stride=1, bias=False), |
| 65 | nn.LeakyReLU(0.2) |
| 66 | ) |
| 67 | |
| 68 | |
| 69 | self.down_h = 1 |
| 70 | for stride in strides: |
| 71 | self.down_h *= stride[0] |
| 72 | self.size_h = 32 // self.down_h * 2 |
| 73 | |
| 74 | |
| 75 | self.feature2w = nn.Sequential( |
| 76 | PixelNorm(), |
| 77 | EqualLinear(self.size_h*self.size_h*feature_out_dim, 512, bias=True, bias_init_val=0, lr_mul=1, activation='fused_lrelu'), |
| 78 | EqualLinear(512, 512, bias=True, bias_init_val=0, lr_mul=1, activation='fused_lrelu'), |
| 79 | ) |
| 80 | |
| 81 | for m in self.modules(): |
| 82 | if isinstance(m, nn.Conv2d): |
| 83 | n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels |
| 84 | m.weight.data.normal_(0, math.sqrt(2. / n)) |
| 85 | |
| 86 | |
| 87 | def _make_layer(self, block, planes, blocks, stride=1): |
nothing calls this directly
no test coverage detected