| 169 | weight_init(self) |
| 170 | |
| 171 | class MSA_module(nn.Module): |
| 172 | def __init__(self, dim=128): |
| 173 | super(MSA_module, self).__init__() |
| 174 | self.B_TA = MSA_head() |
| 175 | self.F_TA = MSA_head() |
| 176 | self.TA = MSA_head() |
| 177 | self.Fuse = nn.Conv2d(3*dim,dim,kernel_size=3,padding=1) |
| 178 | self.Fuse2 = nn.Sequential(nn.Conv2d(dim, dim, kernel_size=1), nn.Conv2d(dim, dim, kernel_size=3, padding=1), nn.BatchNorm2d(dim), nn.ReLU(inplace=True)) |
| 179 | |
| 180 | def forward(self,x,side_x,mask): |
| 181 | N,C,H,W = x.shape |
| 182 | mask = F.interpolate(mask,size=x.size()[2:],mode='bilinear') |
| 183 | mask_d = mask.detach() |
| 184 | mask_d = torch.sigmoid(mask_d) |
| 185 | xf = self.F_TA(x,mask_d) |
| 186 | xb = self.B_TA(x,1-mask_d) |
| 187 | x = self.TA(x) |
| 188 | x = torch.cat((xb,xf,x),1) |
| 189 | x = x.view(N,3*C,H,W) |
| 190 | x = self.Fuse(x) |
| 191 | D = self.Fuse2(side_x+side_x*x) |
| 192 | return D |
| 193 | |
| 194 | def initialize(self): |
| 195 | weight_init(self) |
| 196 | |
| 197 | class Conv_Block(nn.Module): |
| 198 | def __init__(self, channels): |