| 185 | return preds[:, 1:, :] |
| 186 | |
| 187 | def forward_test(self, feat, holistic_feat, data=None): |
| 188 | bsz = feat.shape[0] |
| 189 | seq_len = self.max_seq_len |
| 190 | holistic_feat = holistic_feat.unsqueeze(1) |
| 191 | tokens = torch.full((bsz, ), |
| 192 | self.start_idx, |
| 193 | device=feat.device, |
| 194 | dtype=torch.long) |
| 195 | outputs = [] |
| 196 | tokens = self.embedding(tokens) |
| 197 | tokens = tokens.unsqueeze(1).expand(-1, seq_len, -1) |
| 198 | tokens = torch.cat((holistic_feat, tokens), dim=1) |
| 199 | for i in range(1, seq_len + 1): |
| 200 | Hidden_state, attn_feat = self._2d_attation(feat, |
| 201 | tokens, |
| 202 | data=data, |
| 203 | training=self.training) |
| 204 | if self.pred_concat: |
| 205 | f_c = holistic_feat.size(-1) |
| 206 | holistic_feat = holistic_feat.expand(bsz, seq_len + 1, f_c) |
| 207 | preds = self.prediction( |
| 208 | torch.cat((Hidden_state, attn_feat, holistic_feat), 2)) |
| 209 | else: |
| 210 | preds = self.prediction(attn_feat) |
| 211 | # bsz * (seq_len + 1) * num_classes |
| 212 | char_output = preds[:, i, :] |
| 213 | char_output = F.softmax(char_output, -1) |
| 214 | outputs.append(char_output) |
| 215 | _, max_idx = torch.max(char_output, dim=1, keepdim=False) |
| 216 | char_embedding = self.embedding(max_idx) |
| 217 | if (i < seq_len): |
| 218 | tokens[:, i + 1, :] = char_embedding |
| 219 | if (tokens == self.end_idx).any(dim=-1).all(): |
| 220 | break |
| 221 | outputs = torch.stack(outputs, 1) |
| 222 | |
| 223 | return outputs |
| 224 | |
| 225 | def forward(self, feat, data=None): |
| 226 | if self.use_lstm: |