| | import torch.nn as nn
|
| |
|
| | """
|
| | This code refers to "Pyramid attention network for semantic segmentation", that is
|
| | "https://github.com/JaveyWang/Pyramid-Attention-Networks-pytorch/blob/f719365c1780f062058dd0c94550c6c4766cd937/networks.py#L41"
|
| | """
|
| |
|
| | class FPM(nn.Module):
|
| | def __init__(self, channels=1024):
|
| | """
|
| | Feature Pyramid Attention
|
| | :type channels: int
|
| | """
|
| | super(FPM, self).__init__()
|
| | channels_mid = int(channels/4)
|
| |
|
| | self.channels_cond = channels
|
| |
|
| | self.conv_master = nn.Conv2d(self.channels_cond, channels, kernel_size=1, bias=False)
|
| | self.bn_master = nn.BatchNorm2d(channels)
|
| |
|
| |
|
| | self.conv7x7_1 = nn.Conv2d(self.channels_cond, channels_mid, kernel_size=(7, 7), stride=2, padding=3, bias=False)
|
| | self.bn1_1 = nn.BatchNorm2d(channels_mid)
|
| | self.conv5x5_1 = nn.Conv2d(channels_mid, channels_mid, kernel_size=(5, 5), stride=2, padding=2, bias=False)
|
| | self.bn2_1 = nn.BatchNorm2d(channels_mid)
|
| | self.conv3x3_1 = nn.Conv2d(channels_mid, channels_mid, kernel_size=(3, 3), stride=2, padding=1, bias=False)
|
| | self.bn3_1 = nn.BatchNorm2d(channels_mid)
|
| |
|
| | self.conv7x7_2 = nn.Conv2d(channels_mid, channels_mid, kernel_size=(7, 7), stride=1, padding=3, bias=False)
|
| | self.bn1_2 = nn.BatchNorm2d(channels_mid)
|
| | self.conv5x5_2 = nn.Conv2d(channels_mid, channels_mid, kernel_size=(5, 5), stride=1, padding=2, bias=False)
|
| | self.bn2_2 = nn.BatchNorm2d(channels_mid)
|
| | self.conv3x3_2 = nn.Conv2d(channels_mid, channels_mid, kernel_size=(3, 3), stride=1, padding=1, bias=False)
|
| | self.bn3_2 = nn.BatchNorm2d(channels_mid)
|
| |
|
| |
|
| | self.conv_upsample_3 = nn.ConvTranspose2d(channels_mid, channels_mid, kernel_size=4, stride=2, padding=1, bias=False)
|
| | self.bn_upsample_3 = nn.BatchNorm2d(channels_mid)
|
| |
|
| | self.conv_upsample_2 = nn.ConvTranspose2d(channels_mid, channels_mid, kernel_size=4, stride=2, padding=1, bias=False)
|
| | self.bn_upsample_2 = nn.BatchNorm2d(channels_mid)
|
| |
|
| | self.conv_upsample_1 = nn.ConvTranspose2d(channels_mid, channels, kernel_size=4, stride=2, padding=1, bias=False)
|
| | self.bn_upsample_1 = nn.BatchNorm2d(channels)
|
| |
|
| | self.relu = nn.ReLU(inplace=False)
|
| |
|
| | def forward(self, x):
|
| | """
|
| | :param x: Shape: [b, 2048, h, w]
|
| | :return: out: Feature maps. Shape: [b, 2048, h, w]
|
| | """
|
| |
|
| | x_master = self.conv_master(x)
|
| | x_master = self.bn_master(x_master)
|
| |
|
| |
|
| | x1_1 = self.conv7x7_1(x)
|
| | x1_1 = self.bn1_1(x1_1)
|
| | x1_1 = self.relu(x1_1)
|
| | x1_2 = self.conv7x7_2(x1_1)
|
| | x1_2 = self.bn1_2(x1_2)
|
| |
|
| |
|
| | x2_1 = self.conv5x5_1(x1_1)
|
| | x2_1 = self.bn2_1(x2_1)
|
| | x2_1 = self.relu(x2_1)
|
| | x2_2 = self.conv5x5_2(x2_1)
|
| | x2_2 = self.bn2_2(x2_2)
|
| |
|
| |
|
| | x3_1 = self.conv3x3_1(x2_1)
|
| | x3_1 = self.bn3_1(x3_1)
|
| | x3_1 = self.relu(x3_1)
|
| | x3_2 = self.conv3x3_2(x3_1)
|
| | x3_2 = self.bn3_2(x3_2)
|
| |
|
| |
|
| | x3_upsample = self.relu(self.bn_upsample_3(self.conv_upsample_3(x3_2)))
|
| | x2_merge = self.relu(x2_2 + x3_upsample)
|
| | x2_upsample = self.relu(self.bn_upsample_2(self.conv_upsample_2(x2_merge)))
|
| | x1_merge = self.relu(x1_2 + x2_upsample)
|
| |
|
| | x_master = x_master * self.relu(self.bn_upsample_1(self.conv_upsample_1(x1_merge)))
|
| |
|
| | out = self.relu(x_master)
|
| |
|
| | return out |
