(self, X, prev_reward=None, prev_action=None)
| 290 | self.prev_hidden = None |
| 291 | |
| 292 | def forward(self, X, prev_reward=None, prev_action=None): |
| 293 | |
| 294 | # 0. Setup. |
| 295 | # --------- |
| 296 | batch_size = X.size()[0] |
| 297 | if prev_reward is None: |
| 298 | # (n, 1, 1) |
| 299 | prev_reward = torch.stack([torch.zeros(1, 1)] * batch_size).to(X.device) |
| 300 | else: |
| 301 | prev_reward = prev_reward.unsqueeze(1).unsqueeze(1) |
| 302 | if prev_action is None: |
| 303 | # (n, 1, 1) |
| 304 | prev_action = torch.stack([torch.zeros(1, 1)] * batch_size).to(X.device) |
| 305 | else: |
| 306 | prev_action = prev_action.unsqueeze(1).unsqueeze(1) |
| 307 | # 1 (a). Vision. |
| 308 | # -------------- |
| 309 | |
| 310 | # (n, h, w, c_k + c_v) |
| 311 | O = self.vision(X) |
| 312 | # (n, h, w, c_k), (n, h, w, c_v) |
| 313 | K, V = O.split([self.c_k, self.c_v], dim=3) |
| 314 | # (n, h, w, c_k + c_s), (n, h, w, c_v + c_s) |
| 315 | K, V = self.spatial(K), self.spatial(V) |
| 316 | |
| 317 | # 1 (b). Queries. |
| 318 | # -------------- |
| 319 | if self.prev_output is None: |
| 320 | hidden_state = torch.zeros( |
| 321 | batch_size, self.hidden_size, requires_grad=True |
| 322 | ).to(X.device) |
| 323 | self.prev_output = hidden_state |
| 324 | # (n, h, w, num_queries, c_k + c_s) |
| 325 | Q = self.query(self.prev_output) |
| 326 | |
| 327 | # 2. Answer. |
| 328 | # ---------- |
| 329 | # (n, h, w, num_queries) |
| 330 | A = torch.matmul(K, Q.transpose(2, 1).unsqueeze(1)) |
| 331 | # (n, h, w, num_queries) |
| 332 | A = spatial_softmax(A) |
| 333 | # (n, 1, 1, num_queries) |
| 334 | a = apply_alpha(A, V) |
| 335 | |
| 336 | # (n, (c_v + c_s) * num_queries + (c_k + c_s) * num_queries + 1 + 1) |
| 337 | answer = torch.cat( |
| 338 | torch.chunk(a, 4, dim=1) |
| 339 | + torch.chunk(Q, 4, dim=1) |
| 340 | + (prev_reward.float(), prev_action.float()), |
| 341 | dim=2, |
| 342 | ).squeeze(1) |
| 343 | # (n, hidden_size) |
| 344 | answer = self.answer_processor(answer) |
| 345 | |
| 346 | # 3. Policy. |
| 347 | # ---------- |
| 348 | if self.prev_hidden is None: |
| 349 | h, c = self.policy_core(answer) |
nothing calls this directly
no test coverage detected