|
|
|
|
|
|
|
|
import torch |
|
|
import torch.nn as nn |
|
|
from torch.nn import functional as F |
|
|
from torchvision import models |
|
|
|
|
|
class PositionAttentionModule(nn.Module): |
|
|
def __init__(self, in_dim): |
|
|
super(PositionAttentionModule, self).__init__(); self.chanel_in = in_dim |
|
|
self.query_conv = nn.Conv2d(in_channels=in_dim, out_channels=in_dim // 8, kernel_size=1) |
|
|
self.key_conv = nn.Conv2d(in_channels=in_dim, out_channels=in_dim // 8, kernel_size=1) |
|
|
self.value_conv = nn.Conv2d(in_channels=in_dim, out_channels=in_dim, kernel_size=1) |
|
|
self.gamma = nn.Parameter(torch.zeros(1)); self.softmax = nn.Softmax(dim=-1) |
|
|
def forward(self, x): |
|
|
B, C, H, W = x.size() |
|
|
proj_query = self.query_conv(x).view(B, -1, H * W).permute(0, 2, 1) |
|
|
proj_key = self.key_conv(x).view(B, -1, H * W) |
|
|
energy = torch.bmm(proj_query, proj_key) |
|
|
attention_map = self.softmax(energy) |
|
|
proj_value = self.value_conv(x).view(B, -1, H * W) |
|
|
out = torch.bmm(proj_value, attention_map.permute(0, 2, 1)) |
|
|
out = out.view(B, C, H, W); out = self.gamma * out + x |
|
|
return out |
|
|
|
|
|
class ChannelAttentionModule(nn.Module): |
|
|
def __init__(self): |
|
|
super(ChannelAttentionModule, self).__init__() |
|
|
self.beta = nn.Parameter(torch.zeros(1)); self.softmax = nn.Softmax(dim=-1) |
|
|
def forward(self, x): |
|
|
B, C, H, W = x.size() |
|
|
proj_query = x.view(B, C, -1); proj_key = x.view(B, C, -1).permute(0, 2, 1) |
|
|
energy = torch.bmm(proj_query, proj_key); attention_map = self.softmax(energy) |
|
|
proj_value = x.view(B, C, -1); out = torch.bmm(attention_map, proj_value) |
|
|
out = out.view(B, C, H, W); out = self.beta * out + x |
|
|
return out |
|
|
|
|
|
class DANetHead(nn.Module): |
|
|
def __init__(self, in_channels, out_channels): |
|
|
super(DANetHead, self).__init__() |
|
|
inter_channels = in_channels // 4 |
|
|
self.conv5a = nn.Sequential(nn.Conv2d(in_channels, inter_channels, 3, padding=1, bias=False), nn.BatchNorm2d(inter_channels), nn.ReLU()) |
|
|
self.conv5c = nn.Sequential(nn.Conv2d(in_channels, inter_channels, 3, padding=1, bias=False), nn.BatchNorm2d(inter_channels), nn.ReLU()) |
|
|
self.sa = PositionAttentionModule(inter_channels); self.sc = ChannelAttentionModule() |
|
|
self.conv51 = nn.Sequential(nn.Conv2d(inter_channels, inter_channels, 3, padding=1, bias=False), nn.BatchNorm2d(inter_channels), nn.ReLU()) |
|
|
self.conv52 = nn.Sequential(nn.Conv2d(inter_channels, inter_channels, 3, padding=1, bias=False), nn.BatchNorm2d(inter_channels), nn.ReLU()) |
|
|
self.conv8 = nn.Sequential(nn.Dropout2d(0.1, False), nn.Conv2d(inter_channels, out_channels, 1)) |
|
|
def forward(self, x): |
|
|
feat_sa = self.conv5a(x); sa_feat = self.sa(feat_sa); sa_conv = self.conv51(sa_feat) |
|
|
feat_sc = self.conv5c(x); sc_feat = self.sc(feat_sc); sc_conv = self.conv52(sc_feat) |
|
|
feat_sum = sa_conv + sc_conv; s_out = self.conv8(feat_sum) |
|
|
return s_out |
|
|
|
|
|
class DANet(nn.Module): |
|
|
def __init__(self, num_classes=2, backbone='resnet50', pretrained_base=False, aux=False): |
|
|
super(DANet, self).__init__(); self.aux = aux |
|
|
weights = None |
|
|
resnet = models.resnet50(weights=weights); backbone_out_channels = 2048 |
|
|
resnet.layer3[0].conv2.stride = (1, 1); resnet.layer3[0].downsample[0].stride = (1, 1) |
|
|
resnet.layer4[0].conv2.stride = (1, 1); resnet.layer4[0].downsample[0].stride = (1, 1) |
|
|
for i in range(len(resnet.layer4)): resnet.layer4[i].conv2.dilation = (2, 2); resnet.layer4[i].conv2.padding = (2, 2) |
|
|
self.conv1=resnet.conv1; self.bn1=resnet.bn1; self.relu=resnet.relu; self.maxpool=resnet.maxpool |
|
|
self.layer1=resnet.layer1; self.layer2=resnet.layer2; self.layer3=resnet.layer3; self.layer4=resnet.layer4 |
|
|
self.head = DANetHead(backbone_out_channels, num_classes) |
|
|
def forward(self, x): |
|
|
imsize = x.size()[2:]; x = self.conv1(x); x = self.bn1(x); x = self.relu(x); x = self.maxpool(x) |
|
|
x = self.layer1(x); x = self.layer2(x); c3 = self.layer3(x); c4 = self.layer4(c3) |
|
|
main_out = self.head(c4); main_out = F.interpolate(main_out, size=imsize, mode='bilinear', align_corners=True) |
|
|
return main_out |
|
|
|
