| 34 | x = self.bottleneck(x) |
| 35 | return x |
| 36 | class AlignedModule(nn.Module): |
| 37 | #SFNet-DFNet |
| 38 | def __init__(self, inplane, outplane): |
| 39 | super(AlignedModule, self).__init__() |
| 40 | self.down_h = nn.Conv2d(inplane, outplane, 1, bias=False) |
| 41 | self.down_l = nn.Conv2d(inplane, outplane, 1, bias=False) |
| 42 | self.flow_make = nn.Conv2d(outplane*2, 2, kernel_size=3, padding=1, bias=False) |
| 43 | |
| 44 | def forward(self, x): |
| 45 | low_feature, h_feature = x |
| 46 | h_feature_orign = h_feature |
| 47 | h, w = low_feature.size()[2:] |
| 48 | size = (h, w) |
| 49 | low_feature = self.down_l(low_feature) |
| 50 | h_feature= self.down_h(h_feature) |
| 51 | h_feature = F.interpolate(h_feature,size=size,mode="bilinear",align_corners=False) |
| 52 | flow = self.flow_make(torch.cat([h_feature, low_feature], 1)) |
| 53 | h_feature = self.flow_warp(h_feature_orign, flow, size=size) |
| 54 | |
| 55 | return h_feature |
| 56 | |
| 57 | def flow_warp(self, input, flow, size): |
| 58 | out_h, out_w = size |
| 59 | n, c, h, w = input.size() |
| 60 | # n, c, h, w |
| 61 | # n, 2, h, w |
| 62 | |
| 63 | norm = torch.tensor([[[[out_w, out_h]]]]).type_as(input).to(input.device) |
| 64 | h = torch.linspace(-1.0, 1.0, out_h).view(-1, 1).repeat(1, out_w) |
| 65 | w = torch.linspace(-1.0, 1.0, out_w).repeat(out_h, 1) |
| 66 | grid = torch.cat((w.unsqueeze(2), h.unsqueeze(2)), 2) |
| 67 | grid = grid.repeat(n, 1, 1, 1).type_as(input).to(input.device) |
| 68 | grid = grid + flow.permute(0, 2, 3, 1) / norm |
| 69 | |
| 70 | output = F.grid_sample(input, grid,align_corners=False) |
| 71 | return output |
| 72 | class FeatureFusionModule(nn.Module): |
| 73 | # BiseNet |
| 74 | def __init__(self, in_chan, out_chan): |