| 227 | return x, h_c |
| 228 | |
| 229 | def _forward_spatial_gru(self, x, t_h, masks, s_h): |
| 230 | if x.size(0) == t_h.size(0): |
| 231 | x, t_h, s_h = self.gru(x.unsqueeze(0), (t_h * masks).unsqueeze(0), s_h.unsqueeze(0)) |
| 232 | x = x.squeeze(0) |
| 233 | t_h = t_h.squeeze(0) |
| 234 | s_h = s_h.squeeze(0) |
| 235 | else: |
| 236 | # x is a (T, N, -1) tensor that has been flatten to (T * N, -1) |
| 237 | N = t_h.size(0) |
| 238 | T = int(x.size(0) / N) |
| 239 | |
| 240 | # unflatten |
| 241 | x = x.view(T, N, x.size(-3), x.size(-2), x.size(-1)) |
| 242 | |
| 243 | # Same deal with masks |
| 244 | masks = masks.view(T, N) |
| 245 | |
| 246 | # Let's figure out which steps in the sequence have a zero for any agent |
| 247 | # We will always assume t=0 has a zero in it as that makes the logic cleaner |
| 248 | has_zeros = ((masks[1:] == 0.0).any(dim=-1).nonzero().squeeze().cpu()) |
| 249 | |
| 250 | # +1 to correct the masks[1:] |
| 251 | if has_zeros.dim() == 0: |
| 252 | # Deal with scalar |
| 253 | has_zeros = [has_zeros.item() + 1] |
| 254 | else: |
| 255 | has_zeros = (has_zeros + 1).numpy().tolist() |
| 256 | |
| 257 | # add t=0 and t=T to the list |
| 258 | has_zeros = [0] + has_zeros + [T] |
| 259 | |
| 260 | t_h = t_h.unsqueeze(0) |
| 261 | s_h = s_h.unsqueeze(0) |
| 262 | outputs = [] |
| 263 | for i in range(len(has_zeros) - 1): |
| 264 | # We can now process steps that don't have any zeros in masks together! |
| 265 | # This is much faster |
| 266 | start_idx = has_zeros[i] |
| 267 | end_idx = has_zeros[i + 1] |
| 268 | |
| 269 | rnn_scores, t_h, s_h = self.gru( |
| 270 | x[start_idx:end_idx], |
| 271 | t_h * (masks[start_idx].unsqueeze(0).unsqueeze(-1)), |
| 272 | s_h) |
| 273 | |
| 274 | outputs.append(rnn_scores) |
| 275 | |
| 276 | # assert len(outputs) == T |
| 277 | # x is a (T, N, -1) tensor |
| 278 | x = torch.cat(outputs, dim=0) |
| 279 | # flatten |
| 280 | x = x.view(T * N, -1) |
| 281 | t_h = t_h.squeeze(0) |
| 282 | s_h = s_h.squeeze(0) |
| 283 | |
| 284 | return x, t_h, s_h |
| 285 | |
| 286 | def forward(self, inputs, temporal_hidden_state=None, masks=None, spatial_hidden_state=None): |