| 208 | |
| 209 | |
| 210 | class ClassificationPredSaver(object): |
| 211 | |
| 212 | def __init__(self, length, save_path=None): |
| 213 | |
| 214 | if save_path is not None: |
| 215 | |
| 216 | # Remove filetype from save_path |
| 217 | save_path = save_path.split('.')[0] |
| 218 | self.save_path = save_path |
| 219 | |
| 220 | self.length = length |
| 221 | |
| 222 | self.all_preds = None |
| 223 | self.all_labels = None |
| 224 | |
| 225 | self.running_start_idx = 0 |
| 226 | |
| 227 | def update(self, preds, labels=None): |
| 228 | |
| 229 | # Expect preds in shape B x C |
| 230 | |
| 231 | if torch.is_tensor(preds): |
| 232 | preds = preds.detach().cpu().numpy() |
| 233 | |
| 234 | b, c = preds.shape |
| 235 | |
| 236 | if self.all_preds is None: |
| 237 | self.all_preds = np.zeros((self.length, c)) |
| 238 | |
| 239 | self.all_preds[self.running_start_idx: self.running_start_idx + b] = preds |
| 240 | |
| 241 | if labels is not None: |
| 242 | if torch.is_tensor(labels): |
| 243 | labels = labels.detach().cpu().numpy() |
| 244 | |
| 245 | if self.all_labels is None: |
| 246 | self.all_labels = np.zeros((self.length,)) |
| 247 | |
| 248 | self.all_labels[self.running_start_idx: self.running_start_idx + b] = labels |
| 249 | |
| 250 | # Maintain running index on dataset being evaluated |
| 251 | self.running_start_idx += b |
| 252 | |
| 253 | def save(self): |
| 254 | |
| 255 | # Softmax over preds |
| 256 | preds = torch.from_numpy(self.all_preds) |
| 257 | preds = torch.nn.Softmax(dim=-1)(preds) |
| 258 | self.all_preds = preds.numpy() |
| 259 | |
| 260 | pred_path = self.save_path + '.pth' |
| 261 | print(f'Saving all predictions to {pred_path}') |
| 262 | |
| 263 | torch.save(self.all_preds, pred_path) |
| 264 | |
| 265 | if self.all_labels is not None: |
| 266 | |
| 267 | # Evaluate |
nothing calls this directly
no outgoing calls
no test coverage detected