(A)
| 224 | |
| 225 | |
| 226 | def spatial_softmax(A): |
| 227 | # A: batch_size x h x w x d |
| 228 | b, h, w, d = A.size() |
| 229 | # Flatten A s.t. softmax is applied to each grid (not over queries) |
| 230 | A = A.reshape(b, h * w, d) |
| 231 | A = F.softmax(A, dim=1) |
| 232 | # Reshape A to original shape. |
| 233 | A = A.reshape(b, h, w, d) |
| 234 | return A |
| 235 | |
| 236 | |
| 237 | def apply_alpha(A, V): |