| 84 | return x |
| 85 | |
| 86 | class DProjector(nn.Module): |
| 87 | def __init__(self, text_dim=512, in_dim=512, kernel_size=1): |
| 88 | super().__init__() |
| 89 | self.in_dim = in_dim |
| 90 | self.kernel_size = kernel_size |
| 91 | # visual projector |
| 92 | |
| 93 | self.vis = nn.Sequential( # os16 -> os4 |
| 94 | nn.Upsample(scale_factor=2, mode='bilinear'), |
| 95 | conv_layer(in_dim, in_dim, 3, padding=1), |
| 96 | nn.Upsample(scale_factor=2, mode='bilinear'), |
| 97 | conv_layer(in_dim, in_dim, 3, padding=1), |
| 98 | nn.Conv2d(in_dim, in_dim, 1)) |
| 99 | |
| 100 | # textual projector |
| 101 | out_dim = 1 * in_dim * kernel_size * kernel_size + 1 |
| 102 | self.txt = nn.Linear(text_dim, out_dim) |
| 103 | |
| 104 | def forward(self, x, text): |
| 105 | ''' |
| 106 | x: b, 512, 104, 104 |
| 107 | text: b, 512 |
| 108 | ''' |
| 109 | x = self.vis(x) # Eq. 8 |
| 110 | |
| 111 | B, C, H, W = x.size() |
| 112 | # 1, b*256, 104, 104 |
| 113 | x = x.reshape(1, B * C, H, W) |
| 114 | # txt: b, 1, (256*3*3 + 1) -> b, 1, 256, 3, 3 / b |
| 115 | text = self.txt(text) # Eq. 8 |
| 116 | |
| 117 | weight, bias = text[:, :-1], text[:, -1] |
| 118 | weight = weight.reshape(B, C, self.kernel_size, self.kernel_size) |
| 119 | # Conv2d - 1, b*256, 104, 104 -> 1, b, 104, 104 |
| 120 | out = F.conv2d(x, |
| 121 | weight, |
| 122 | padding=1, |
| 123 | groups=B, |
| 124 | bias=bias) |
| 125 | |
| 126 | # b, 1, 104, 104 |
| 127 | out = out.transpose(0,1) |
| 128 | return out |
| 129 | |
| 130 | |
| 131 | class CrossAttn(nn.Module): |