(self, dim=128)
| 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 |