File size: 2,037 Bytes
f7f604d |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 |
import torch
import torch.nn as nn
import torch.nn.functional as F
from .layers import *
class PAA_kernel(nn.Module):
def __init__(self, in_channel, out_channel, receptive_size, stage_size=None):
super(PAA_kernel, self).__init__()
self.conv0 = Conv2d(in_channel, out_channel, 1)
self.conv1 = Conv2d(out_channel, out_channel, kernel_size=(1, receptive_size))
self.conv2 = Conv2d(out_channel, out_channel, kernel_size=(receptive_size, 1))
self.conv3 = Conv2d(out_channel, out_channel, 3, dilation=receptive_size)
self.Hattn = SelfAttention(out_channel, 'h', stage_size[0] if stage_size is not None else None)
self.Wattn = SelfAttention(out_channel, 'w', stage_size[1] if stage_size is not None else None)
def forward(self, x):
x = self.conv0(x)
x = self.conv1(x)
x = self.conv2(x)
Hx = self.Hattn(x)
Wx = self.Wattn(x)
x = self.conv3(Hx + Wx)
return x
class PAA_e(nn.Module):
def __init__(self, in_channel, out_channel, base_size=None, stage=None):
super(PAA_e, self).__init__()
self.relu = nn.ReLU(True)
if base_size is not None and stage is not None:
self.stage_size = (base_size[0] // (2 ** stage), base_size[1] // (2 ** stage))
else:
self.stage_size = None
self.branch0 = Conv2d(in_channel, out_channel, 1)
self.branch1 = PAA_kernel(in_channel, out_channel, 3, self.stage_size)
self.branch2 = PAA_kernel(in_channel, out_channel, 5, self.stage_size)
self.branch3 = PAA_kernel(in_channel, out_channel, 7, self.stage_size)
self.conv_cat = Conv2d(4 * out_channel, out_channel, 3)
self.conv_res = Conv2d(in_channel, out_channel, 1)
def forward(self, x):
x0 = self.branch0(x)
x1 = self.branch1(x)
x2 = self.branch2(x)
x3 = self.branch3(x)
x_cat = self.conv_cat(torch.cat((x0, x1, x2, x3), 1))
x = self.relu(x_cat + self.conv_res(x))
return x
|