limingjuan
/
Prediction-MindSpore
Not watched
1
0
Fork 0
Code
Releases 0 Wiki evaluate Activity
Issues 0 Pull Requests 0 Datasets Model Cloudbrain HPC
You can not select more than 25 topics Topics must start with a chinese character,a letter or number, can include dashes ('-') and can be up to 35 characters long.
Tree: abd9fc84aa
Branches Tags
${ item.name }
Create branch ${ searchTerm }
from 'abd9fc84aa'
${ noResults }
Prediction-MindSpore/model/attention/PolarizedSelfAttention.py

101 lines
4.7 kB
Raw Blame History 

  1. """

  2. MindSpore implementation of 'PolarizedSelfAttention'

  3. Refer to "Polarized Self-Attention: Towards High-quality Pixel-wise Regression"

  4. """

  5. import mindspore as ms

  6. from mindspore import nn

  7. 

  8. 

  9. class ParallelPolarizedSelfAttention(nn.Cell):

  10.     """ Parallel Polarized Self Attention """

  11.     def __init__(self, channel=512):

  12.         super().__init__()

  13.         self.ch_wv = nn.Conv2d(channel, channel // 2, kernel_size=(1, 1))

  14.         self.ch_wq = nn.Conv2d(channel, 1, kernel_size=(1, 1))

  15.         self.softmax_channel = nn.Softmax(1)

  16.         self.softmax_spatial = nn.Softmax(-1)

  17.         self.ch_wz = nn.Conv2d(channel // 2, channel, kernel_size=(1, 1))

  18.         self.ln = nn.LayerNorm((channel,))

  19.         self.sigmoid = nn.Sigmoid()

  20.         self.sp_wv = nn.Conv2d(channel, channel // 2, kernel_size=(1, 1))

  21.         self.sp_wq = nn.Conv2d(channel, channel // 2, kernel_size=(1, 1))

  22.         self.agp = nn.AdaptiveAvgPool2d((1, 1))

  23. 

  24.     def construct(self, x):

  25.         b, c, h, w = x.shape

  26. 

  27.         # Channel-only Self-Attention

  28.         channel_wv = self.ch_wv(x)  # bs,c//2,h,w

  29.         channel_wq = self.ch_wq(x)  # bs,1,h,w

  30.         channel_wv = channel_wv.reshape(b, c // 2, -1)  # bs,c//2,h*w

  31.         channel_wq = channel_wq.reshape(b, -1, 1)  # bs,h*w,1

  32.         channel_wq = self.softmax_channel(channel_wq)

  33.         channel_wz = ms.ops.matmul(channel_wv, channel_wq).unsqueeze(-1)  # bs,c//2,1,1

  34.         channel_weight = self.sigmoid(self.ln(self.ch_wz(channel_wz).reshape(b, c, 1).permute(0, 2, 1))).permute(0, 2,

  35.                                                                                                                  1).reshape(

  36.             b, c, 1, 1)  # bs,c,1,1

  37.         channel_out = channel_weight * x

  38. 

  39.         # Spatial-only Self-Attention

  40.         spatial_wv = self.sp_wv(x)  # bs,c//2,h,w

  41.         spatial_wq = self.sp_wq(x)  # bs,c//2,h,w

  42.         spatial_wq = self.agp(spatial_wq)  # bs,c//2,1,1

  43.         spatial_wv = spatial_wv.reshape(b, c // 2, -1)  # bs,c//2,h*w

  44.         spatial_wq = spatial_wq.permute(0, 2, 3, 1).reshape(b, 1, c // 2)  # bs,1,c//2

  45.         spatial_wq = self.softmax_spatial(spatial_wq)

  46.         spatial_wz = ms.ops.matmul(spatial_wq, spatial_wv)  # bs,1,h*w

  47.         spatial_weight = self.sigmoid(spatial_wz.reshape(b, 1, h, w))  # bs,1,h,w

  48.         spatial_out = spatial_weight * x

  49.         out = spatial_out + channel_out

  50.         return out

  51. 

  52. 

  53. class SequentialPolarizedSelfAttention(nn.Cell):

  54.     """ Sequential Polarized Self Attention """

  55.     def __init__(self, channel=512):

  56.         super().__init__()

  57.         self.ch_wv = nn.Conv2d(channel, channel // 2, kernel_size=(1, 1))

  58.         self.ch_wq = nn.Conv2d(channel, 1, kernel_size=(1, 1))

  59.         self.softmax_channel = nn.Softmax(1)

  60.         self.softmax_spatial = nn.Softmax(-1)

  61.         self.ch_wz = nn.Conv2d(channel // 2, channel, kernel_size=(1, 1))

  62.         self.ln = nn.LayerNorm((channel,))

  63.         self.sigmoid = nn.Sigmoid()

  64.         self.sp_wv = nn.Conv2d(channel, channel // 2, kernel_size=(1, 1))

  65.         self.sp_wq = nn.Conv2d(channel, channel // 2, kernel_size=(1, 1))

  66.         self.agp = nn.AdaptiveAvgPool2d((1, 1))

  67. 

  68.     def construct(self, x):

  69.         b, c, h, w = x.shape

  70. 

  71.         # Channel-only Self-Attention

  72.         channel_wv = self.ch_wv(x)  # bs,c//2,h,w

  73.         channel_wq = self.ch_wq(x)  # bs,1,h,w

  74.         channel_wv = channel_wv.reshape(b, c // 2, -1)  # bs,c//2,h*w

  75.         channel_wq = channel_wq.reshape(b, -1, 1)  # bs,h*w,1

  76.         channel_wq = self.softmax_channel(channel_wq)

  77.         channel_wz = ms.ops.matmul(channel_wv, channel_wq).unsqueeze(-1)  # bs,c//2,1,1

  78.         channel_weight = self.sigmoid(self.ln(self.ch_wz(channel_wz).reshape(b, c, 1).permute(0, 2, 1))).permute(0, 2,

  79.                                                                                                                  1).reshape(

  80.             b, c, 1, 1)  # bs,c,1,1

  81.         channel_out = channel_weight * x

  82. 

  83.         # Spatial-only Self-Attention

  84.         spatial_wv = self.sp_wv(channel_out)  # bs,c//2,h,w

  85.         spatial_wq = self.sp_wq(channel_out)  # bs,c//2,h,w

  86.         spatial_wq = self.agp(spatial_wq)  # bs,c//2,1,1

  87.         spatial_wv = spatial_wv.reshape(b, c // 2, -1)  # bs,c//2,h*w

  88.         spatial_wq = spatial_wq.permute(0, 2, 3, 1).reshape(b, 1, c // 2)  # bs,1,c//2

  89.         spatial_wq = self.softmax_spatial(spatial_wq)

  90.         spatial_wz = ms.ops.matmul(spatial_wq, spatial_wv)  # bs,1,h*w

  91.         spatial_weight = self.sigmoid(spatial_wz.reshape(b, 1, h, w))  # bs,1,h,w

  92.         spatial_out = spatial_weight * channel_out

  93.         return spatial_out

  94. 

  95. 

  96. if __name__ == '__main__':

  97.     dummy_input = ms.ops.randn((1, 512, 7, 7))

  98.     psa = SequentialPolarizedSelfAttention(channel=512)

  99.     output = psa(dummy_input)

  100.     print(output.shape)