| 149 | return corr |
| 150 | |
| 151 | class MemoryDecoder(nn.Module): |
| 152 | def __init__(self, cfg): |
| 153 | super(MemoryDecoder, self).__init__() |
| 154 | dim = self.dim = cfg.query_latent_dim |
| 155 | self.cfg = cfg |
| 156 | |
| 157 | self.flow_token_encoder = nn.Sequential( |
| 158 | nn.Conv2d(81*cfg.cost_heads_num, dim, 1, 1), |
| 159 | nn.GELU(), |
| 160 | nn.Conv2d(dim, dim, 1, 1) |
| 161 | ) |
| 162 | self.proj = nn.Conv2d(256, 256, 1) |
| 163 | self.depth = cfg.decoder_depth |
| 164 | self.decoder_layer = MemoryDecoderLayer(dim, cfg) |
| 165 | |
| 166 | if self.cfg.gma: |
| 167 | self.update_block = GMAUpdateBlock(self.cfg, hidden_dim=128) |
| 168 | self.att = Attention(args=self.cfg, dim=128, heads=1, max_pos_size=160, dim_head=128) |
| 169 | else: |
| 170 | self.update_block = BasicUpdateBlock(self.cfg, hidden_dim=128) |
| 171 | |
| 172 | def upsample_flow(self, flow, mask): |
| 173 | """ Upsample flow field [H/8, W/8, 2] -> [H, W, 2] using convex combination """ |
| 174 | N, _, H, W = flow.shape |
| 175 | mask = mask.view(N, 1, 9, 8, 8, H, W) |
| 176 | mask = torch.softmax(mask, dim=2) |
| 177 | |
| 178 | up_flow = F.unfold(8 * flow, [3,3], padding=1) |
| 179 | up_flow = up_flow.view(N, 2, 9, 1, 1, H, W) |
| 180 | |
| 181 | up_flow = torch.sum(mask * up_flow, dim=2) |
| 182 | up_flow = up_flow.permute(0, 1, 4, 2, 5, 3) |
| 183 | return up_flow.reshape(N, 2, 8*H, 8*W) |
| 184 | |
| 185 | def encode_flow_token(self, cost_maps, coords): |
| 186 | """ |
| 187 | cost_maps - B*H1*W1, cost_heads_num, H2, W2 |
| 188 | coords - B, 2, H1, W1 |
| 189 | """ |
| 190 | coords = coords.permute(0, 2, 3, 1) |
| 191 | batch, h1, w1, _ = coords.shape |
| 192 | |
| 193 | r = 4 |
| 194 | dx = torch.linspace(-r, r, 2*r+1) |
| 195 | dy = torch.linspace(-r, r, 2*r+1) |
| 196 | delta = torch.stack(torch.meshgrid(dy, dx), axis=-1).to(coords.device) |
| 197 | |
| 198 | centroid = coords.reshape(batch*h1*w1, 1, 1, 2) |
| 199 | delta = delta.view(1, 2*r+1, 2*r+1, 2) |
| 200 | coords = centroid + delta |
| 201 | corr = bilinear_sampler(cost_maps, coords) |
| 202 | corr = corr.view(batch, h1, w1, -1).permute(0, 3, 1, 2) |
| 203 | return corr |
| 204 | |
| 205 | def forward(self, cost_memory, context, data={}, flow_init=None): |
| 206 | """ |
| 207 | memory: [B*H1*W1, H2'*W2', C] |
| 208 | context: [B, D, H1, W1] |