| import torch.nn as nn |
| import torch.nn.functional as F |
| import torchvision.models as models |
|
|
| class ResNetUNet(nn.Module): |
| |
| def __init__(self, num_classes=2): |
| super(ResNetUNet, self).__init__() |
| resnet = models.resnet50(pretrained=True) |
|
|
| resnet.conv1 = nn.Conv2d(1, 64, kernel_size=7, stride=2, padding=3, bias=False) |
| nn.init.kaiming_normal_(resnet.conv1.weight, mode='fan_out', nonlinearity='relu') |
|
|
| self.input_block = nn.Sequential( |
| resnet.conv1, |
| resnet.bn1, |
| resnet.relu |
| ) |
| self.maxpool = resnet.maxpool |
|
|
| self.encoder1 = resnet.layer1 |
| self.encoder2 = resnet.layer2 |
| self.encoder3 = resnet.layer3 |
| self.bottleneck = resnet.layer4 |
|
|
| self.up1 = nn.ConvTranspose2d(2048, 1024, kernel_size=2, stride=2) |
| self.up2 = nn.ConvTranspose2d(1024, 512, kernel_size=2, stride=2) |
| self.up3 = nn.ConvTranspose2d(512, 256, kernel_size=2, stride=2) |
| self.up4 = nn.ConvTranspose2d(256, 64, kernel_size=2, stride=2) |
| self.up5 = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True) |
|
|
| self.conv1 = nn.Sequential( |
| nn.Conv2d(1024, 1024, kernel_size=3, padding=1), |
| nn.BatchNorm2d(1024), |
| nn.ReLU() |
| ) |
| self.conv2 = nn.Sequential( |
| nn.Conv2d(512, 512, kernel_size=3, padding=1), |
| nn.BatchNorm2d(512), |
| nn.ReLU() |
| ) |
| self.conv3 = nn.Sequential( |
| nn.Conv2d(256, 256, kernel_size=3, padding=1), |
| nn.BatchNorm2d(256), |
| nn.ReLU() |
| ) |
| self.conv4 = nn.Sequential( |
| nn.Conv2d(64, 64, kernel_size=3, padding=1), |
| nn.BatchNorm2d(64), |
| nn.ReLU() |
| ) |
|
|
| self.out_conv = nn.Conv2d(64, num_classes, kernel_size=1) |
|
|
| def forward(self, x): |
| x0 = self.input_block(x) |
| x1 = self.maxpool(x0) |
| x2 = self.encoder1(x1) |
| x3 = self.encoder2(x2) |
| x4 = self.encoder3(x3) |
| x5 = self.bottleneck(x4) |
|
|
| d1 = F.relu(self.up1(x5) + x4) |
| d1 = self.conv1(d1) |
|
|
| d2 = F.relu(self.up2(d1) + x3) |
| d2 = self.conv2(d2) |
|
|
| d3 = F.relu(self.up3(d2) + x2) |
| d3 = self.conv3(d3) |
|
|
| d4 = F.relu(self.up4(d3) + x0) |
| d4 = self.conv4(d4) |
|
|
| d5 = self.up5(d4) |
| out = self.out_conv(d5) |
| return out |
|
|
|
|
