Function predict_batch

python_server/server.py:146–166

(req: BatchPredictRequest)

# 接口 2: 批量预测 (支撑 5x11 并发)
@app.post("/predict_batch", response_model=BatchPredictResponse)
async def predict_batch(req: BatchPredictRequest):
    if model is None: raise HTTPException(500, "Model not ready")
    
    batch_size = len(req.images)
    if batch_size == 0: return BatchPredictResponse(labels=[], logits_batch=[])

    try:
        # List -> Tensor -> GPU
        # 这一步对于大 Batch 可能会有点慢，但比起 HTTP RTT 已经很快了
        input_tensor = torch.tensor(req.images).float().to(DEVICE)
        input_tensor = input_tensor.view(-1, 3, 32, 32)
        
        with torch.no_grad(): # 必加：显存保护
            output = model(input_tensor)
            pred_labels = output.argmax(dim=1).cpu().tolist()
            pred_logits = output.cpu().tolist()
            
            return BatchPredictResponse(labels=pred_labels, logits_batch=pred_logits)
            
    except Exception as e:
        raise HTTPException(500, str(e))

# 本地调试用
if __name__ == "__main__":