(self, streamer: NpzStreamer)
| 82 | return 40 |
| 83 | |
| 84 | def features_and_preds(self, streamer: NpzStreamer) -> Tuple[np.ndarray, np.ndarray]: |
| 85 | batch_size = self.device_batch_size * len(self.devices) |
| 86 | point_clouds = (x["arr_0"] for x in streamer.stream(batch_size, ["arr_0"])) |
| 87 | |
| 88 | output_features = [] |
| 89 | output_predictions = [] |
| 90 | |
| 91 | with ThreadPool(len(self.devices)) as pool: |
| 92 | for batch in point_clouds: |
| 93 | batch = normalize_point_clouds(batch) |
| 94 | batches = [] |
| 95 | for i, device in zip(range(0, len(batch), self.device_batch_size), self.devices): |
| 96 | batches.append( |
| 97 | torch.from_numpy(batch[i : i + self.device_batch_size]) |
| 98 | .permute(0, 2, 1) |
| 99 | .to(dtype=torch.float32, device=device) |
| 100 | ) |
| 101 | |
| 102 | def compute_features(i_batch): |
| 103 | i, batch = i_batch |
| 104 | with torch.no_grad(): |
| 105 | return self.models[i](batch, features=True) |
| 106 | |
| 107 | for logits, _, features in pool.imap(compute_features, enumerate(batches)): |
| 108 | output_features.append(features.cpu().numpy()) |
| 109 | output_predictions.append(logits.exp().cpu().numpy()) |
| 110 | |
| 111 | return np.concatenate(output_features, axis=0), np.concatenate(output_predictions, axis=0) |
| 112 | |
| 113 | |
| 114 | def normalize_point_clouds(pc: np.ndarray) -> np.ndarray: |
no test coverage detected