| 243 | |
| 244 | |
| 245 | class UnetLayer(nn.Module): |
| 246 | def __init__(self, in_channels, middle_channels, out_channels): |
| 247 | super(UnetLayer, self).__init__() |
| 248 | self.up = nn.ConvTranspose2d(in_channels, out_channels, kernel_size=2, stride=2) |
| 249 | self.conv_relu = nn.Sequential( |
| 250 | nn.Conv2d(middle_channels, out_channels, kernel_size=3, padding=1), |
| 251 | nn.ReLU(inplace=True), |
| 252 | ) |
| 253 | self.fp16_enabled = False |
| 254 | |
| 255 | @force_fp32() |
| 256 | def forward(self, x1, x2): |
| 257 | x1 = self.up(x1) |
| 258 | x1 = torch.cat((x1, x2), dim=1) |
| 259 | x1 = self.conv_relu(x1) |
| 260 | return x1 |
| 261 | |
| 262 | class UNet(nn.Module): |
| 263 | def __init__(self, n_class, fpn_in_channels): |