| 61 | return ret |
| 62 | |
| 63 | class Fusion(nn.Module): |
| 64 | def __init__(self, in_dim_1, in_dim_2, out_dim, bias=False) -> None: |
| 65 | super().__init__() |
| 66 | |
| 67 | self.fusion = nn.Sequential( |
| 68 | nn.Conv2d(in_dim_1+in_dim_2, out_dim, 3, padding=1, bias=bias), |
| 69 | nn.BatchNorm2d(out_dim), |
| 70 | nn.ReLU(), |
| 71 | nn.Conv2d(out_dim, out_dim, 3, padding=1, bias=bias), |
| 72 | nn.BatchNorm2d(out_dim), |
| 73 | nn.ReLU(), |
| 74 | ) |
| 75 | |
| 76 | def forward(self, in_1, in_2): |
| 77 | if in_1.shape[-1] < in_2.shape[-1]: |
| 78 | in_1 = F.interpolate(in_1, size=in_2.shape[-2:], mode='bilinear', align_corners=True) |
| 79 | elif in_1.shape[-1] > in_2.shape[-1]: |
| 80 | in_2 = F.interpolate(in_2, size=in_1.shape[-2:], mode='bilinear', align_corners=True) |
| 81 | |
| 82 | x = torch.cat((in_1, in_2), dim=1) |
| 83 | x = self.fusion(x) |
| 84 | return x |
| 85 | |
| 86 | class DProjector(nn.Module): |
| 87 | def __init__(self, text_dim=512, in_dim=512, kernel_size=1): |