| 49 | |
| 50 | |
| 51 | class PointNetClassifier(FeatureExtractor): |
| 52 | def __init__( |
| 53 | self, |
| 54 | devices: List[Union[str, torch.device]], |
| 55 | device_batch_size: int = 64, |
| 56 | cache_dir: Optional[str] = None, |
| 57 | ): |
| 58 | state_dict = load_checkpoint("pointnet", device=torch.device("cpu"), cache_dir=cache_dir)[ |
| 59 | "model_state_dict" |
| 60 | ] |
| 61 | |
| 62 | self.device_batch_size = device_batch_size |
| 63 | self.devices = devices |
| 64 | self.models = [] |
| 65 | for device in devices: |
| 66 | model = get_model(num_class=40, normal_channel=False, width_mult=2) |
| 67 | model.load_state_dict(state_dict) |
| 68 | model.to(device) |
| 69 | model.eval() |
| 70 | self.models.append(model) |
| 71 | |
| 72 | @property |
| 73 | def supports_predictions(self) -> bool: |
| 74 | return True |
| 75 | |
| 76 | @property |
| 77 | def feature_dim(self) -> int: |
| 78 | return 256 |
| 79 | |
| 80 | @property |
| 81 | def num_classes(self) -> int: |
| 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()) |