| 131 | # trading off latency for throughput. |
| 132 | @serve.batch(max_batch_size=128, batch_wait_timeout_s=0.1) |
| 133 | async def predict_batch(self, images: list[np.ndarray]) -> list[dict[str, Any]]: |
| 134 | # Stack all images into a single tensor. |
| 135 | batch_tensor = torch.cat([ |
| 136 | self.transform(img).unsqueeze(0) |
| 137 | for img in images |
| 138 | ]).to(self.device).float() |
| 139 | |
| 140 | # Single forward pass on the entire batch at once. |
| 141 | with torch.no_grad(): |
| 142 | logits = self.model(batch_tensor) |
| 143 | predictions = torch.argmax(logits, dim=1).cpu().numpy() |
| 144 | |
| 145 | # Unbatch the results and preserve their original order. |
| 146 | return [ |
| 147 | { |
| 148 | "predicted_label": int(pred), |
| 149 | "logits": logit.cpu().numpy().tolist() |
| 150 | } |
| 151 | for pred, logit in zip(predictions, logits) |
| 152 | ] |
| 153 | |
| 154 | @app.post("/") |
| 155 | async def handle_request(self, request: ImageRequest): |