| 174 | p.requires_grad = False |
| 175 | |
| 176 | def _forward_gru(self, x, h_c, masks): |
| 177 | if x.size(0) == h_c.size(0): |
| 178 | x, h_c = self.gru(x.unsqueeze(0), (h_c * masks).unsqueeze(0)) |
| 179 | x = x.squeeze(0) |
| 180 | h_c = h_c.squeeze(0) |
| 181 | else: |
| 182 | # x is a (T, N, -1) tensor that has been flatten to (T * N, -1) |
| 183 | N = h_c.size(0) |
| 184 | T = int(x.size(0) / N) |
| 185 | |
| 186 | # unflatten |
| 187 | x = x.view(T, N, x.size(1)) |
| 188 | |
| 189 | # Same deal with masks |
| 190 | masks = masks.view(T, N) |
| 191 | |
| 192 | # Let's figure out which steps in the sequence have a zero for any agent |
| 193 | # We will always assume t=0 has a zero in it as that makes the logic cleaner |
| 194 | has_zeros = ((masks[1:] == 0.0).any(dim=-1).nonzero().squeeze().cpu()) |
| 195 | |
| 196 | # +1 to correct the masks[1:] |
| 197 | if has_zeros.dim() == 0: |
| 198 | # Deal with scalar |
| 199 | has_zeros = [has_zeros.item() + 1] |
| 200 | else: |
| 201 | has_zeros = (has_zeros + 1).numpy().tolist() |
| 202 | |
| 203 | # add t=0 and t=T to the list |
| 204 | has_zeros = [0] + has_zeros + [T] |
| 205 | |
| 206 | h_c = h_c.unsqueeze(0) |
| 207 | outputs = [] |
| 208 | for i in range(len(has_zeros) - 1): |
| 209 | # We can now process steps that don't have any zeros in masks together! |
| 210 | # This is much faster |
| 211 | start_idx = has_zeros[i] |
| 212 | end_idx = has_zeros[i + 1] |
| 213 | |
| 214 | rnn_scores, h_c = self.gru( |
| 215 | x[start_idx:end_idx], |
| 216 | h_c * (masks[start_idx].unsqueeze(0).unsqueeze(-1))) |
| 217 | |
| 218 | outputs.append(rnn_scores) |
| 219 | |
| 220 | # assert len(outputs) == T |
| 221 | # x is a (T, N, -1) tensor |
| 222 | x = torch.cat(outputs, dim=0) |
| 223 | # flatten |
| 224 | x = x.view(T * N, -1) |
| 225 | h_c = h_c.squeeze(0) |
| 226 | |
| 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): |