Simple block, preserve spatial resolution.
| 163 | |
| 164 | |
| 165 | class SameBlock2d(nn.Module): |
| 166 | """ |
| 167 | Simple block, preserve spatial resolution. |
| 168 | """ |
| 169 | |
| 170 | def __init__(self, in_features, out_features, groups=1, kernel_size=3, padding=1, lrelu=False): |
| 171 | super(SameBlock2d, self).__init__() |
| 172 | self.conv = nn.Conv2d(in_channels=in_features, out_channels=out_features, kernel_size=kernel_size, padding=padding, groups=groups) |
| 173 | self.norm = nn.BatchNorm2d(out_features, affine=True) |
| 174 | if lrelu: |
| 175 | self.ac = nn.LeakyReLU() |
| 176 | else: |
| 177 | self.ac = nn.ReLU() |
| 178 | |
| 179 | def forward(self, x): |
| 180 | out = self.conv(x) |
| 181 | out = self.norm(out) |
| 182 | out = self.ac(out) |
| 183 | return out |
| 184 | |
| 185 | |
| 186 | class Encoder(nn.Module): |