Spaces:
Configuration error
Configuration error
| # Copyright (c) 2021 PaddlePaddle Authors. All Rights Reserved. | |
| # | |
| # Licensed under the Apache License, Version 2.0 (the "License"); | |
| # you may not use this file except in compliance with the License. | |
| # You may obtain a copy of the License at | |
| # | |
| # http://www.apache.org/licenses/LICENSE-2.0 | |
| # | |
| # Unless required by applicable law or agreed to in writing, software | |
| # distributed under the License is distributed on an "AS IS" BASIS, | |
| # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. | |
| # See the License for the specific language governing permissions and | |
| # limitations under the License. | |
| import paddle | |
| import paddle.nn as nn | |
| import paddle.nn.functional as F | |
| from paddleseg.models import layers | |
| class NonLocal2D(nn.Layer): | |
| """Basic Non-local module. | |
| This model is the implementation of "Non-local Neural Networks" | |
| (https://arxiv.org/abs/1711.07971) | |
| Args: | |
| in_channels (int): Channels of the input feature map. | |
| reduction (int): Channel reduction ratio. Default: 2. | |
| use_scale (bool): Whether to scale pairwise_weight by `1/sqrt(inter_channels)` when the mode is `embedded_gaussian`. Default: True. | |
| sub_sample (bool): Whether to utilize max pooling after pairwise function. Default: False. | |
| mode (str): Options are `gaussian`, `concatenation`, `embedded_gaussian` and `dot_product`. Default: embedded_gaussian. | |
| """ | |
| def __init__(self, | |
| in_channels, | |
| reduction=2, | |
| use_scale=True, | |
| sub_sample=False, | |
| mode='embedded_gaussian'): | |
| super(NonLocal2D, self).__init__() | |
| self.in_channels = in_channels | |
| self.reduction = reduction | |
| self.use_scale = use_scale | |
| self.sub_sample = sub_sample | |
| self.mode = mode | |
| if mode not in [ | |
| 'gaussian', 'embedded_gaussian', 'dot_product', 'concatenation' | |
| ]: | |
| raise ValueError( | |
| "Mode should be in 'gaussian', 'concatenation','embedded_gaussian' or 'dot_product'." | |
| ) | |
| self.inter_channels = max(in_channels // reduction, 1) | |
| self.g = nn.Conv2D( | |
| in_channels=self.in_channels, | |
| out_channels=self.inter_channels, | |
| kernel_size=1) | |
| self.conv_out = layers.ConvBNReLU( | |
| in_channels=self.inter_channels, | |
| out_channels=self.in_channels, | |
| kernel_size=1, | |
| bias_attr=False) | |
| if self.mode != "gaussian": | |
| self.theta = nn.Conv2D( | |
| in_channels=self.in_channels, | |
| out_channels=self.inter_channels, | |
| kernel_size=1) | |
| self.phi = nn.Conv2D( | |
| in_channels=self.in_channels, | |
| out_channels=self.inter_channels, | |
| kernel_size=1) | |
| if self.mode == "concatenation": | |
| self.concat_project = layers.ConvBNReLU( | |
| in_channels=self.inter_channels * 2, | |
| out_channels=1, | |
| kernel_size=1, | |
| bias_attr=False) | |
| if self.sub_sample: | |
| max_pool_layer = nn.MaxPool2D(kernel_size=(2, 2)) | |
| self.g = nn.Sequential(self.g, max_pool_layer) | |
| if self.mode != 'gaussian': | |
| self.phi = nn.Sequential(self.phi, max_pool_layer) | |
| else: | |
| self.phi = max_pool_layer | |
| def gaussian(self, theta_x, phi_x): | |
| pairwise_weight = paddle.matmul(theta_x, phi_x) | |
| pairwise_weight = F.softmax(pairwise_weight, axis=-1) | |
| return pairwise_weight | |
| def embedded_gaussian(self, theta_x, phi_x): | |
| pairwise_weight = paddle.matmul(theta_x, phi_x) | |
| if self.use_scale: | |
| pairwise_weight /= theta_x.shape[-1]**0.5 | |
| pairwise_weight = F.softmax(pairwise_weight, -1) | |
| return pairwise_weight | |
| def dot_product(self, theta_x, phi_x): | |
| pairwise_weight = paddle.matmul(theta_x, phi_x) | |
| pairwise_weight /= pairwise_weight.shape[-1] | |
| return pairwise_weight | |
| def concatenation(self, theta_x, phi_x): | |
| h = theta_x.shape[2] | |
| w = phi_x.shape[3] | |
| theta_x = paddle.tile(theta_x, [1, 1, 1, w]) | |
| phi_x = paddle.tile(phi_x, [1, 1, h, 1]) | |
| concat_feature = paddle.concat([theta_x, phi_x], axis=1) | |
| pairwise_weight = self.concat_project(concat_feature) | |
| n, _, h, w = pairwise_weight.shape | |
| pairwise_weight = paddle.reshape(pairwise_weight, [n, h, w]) | |
| pairwise_weight /= pairwise_weight.shape[-1] | |
| return pairwise_weight | |
| def forward(self, x): | |
| n, c, h, w = x.shape | |
| g_x = paddle.reshape(self.g(x), [n, self.inter_channels, -1]) | |
| g_x = paddle.transpose(g_x, [0, 2, 1]) | |
| if self.mode == 'gaussian': | |
| theta_x = paddle.reshape(x, [n, self.inter_channels, -1]) | |
| theta_x = paddle.transpose(theta_x, [0, 2, 1]) | |
| if self.sub_sample: | |
| phi_x = paddle.reshape( | |
| self.phi(x), [n, self.inter_channels, -1]) | |
| else: | |
| phi_x = paddle.reshape(x, [n, self.in_channels, -1]) | |
| elif self.mode == 'concatenation': | |
| theta_x = paddle.reshape( | |
| self.theta(x), [n, self.inter_channels, -1, 1]) | |
| phi_x = paddle.reshape(self.phi(x), [n, self.inter_channels, 1, -1]) | |
| else: | |
| theta_x = paddle.reshape( | |
| self.theta(x), [n, self.inter_channels, -1]) | |
| theta_x = paddle.transpose(theta_x, [0, 2, 1]) | |
| phi_x = paddle.reshape(self.phi(x), [n, self.inter_channels, -1]) | |
| pairwise_func = getattr(self, self.mode) | |
| pairwise_weight = pairwise_func(theta_x, phi_x) | |
| y = paddle.matmul(pairwise_weight, g_x) | |
| y = paddle.transpose(y, [0, 2, 1]) | |
| y = paddle.reshape(y, [n, self.inter_channels, h, w]) | |
| output = x + self.conv_out(y) | |
| return output | |