Deploying a Faster-R-CNN model on Kubernetes
Helping lawyers conduct document reviews
My team works very closely with corporate lawyers. A big part of what we do is to create AI solutions for them to search, sort, automate parts of the document review process. One of them is to detect signatures in documents.
This task, seemingly straightforward, unfolded into a complex, year-long project. It began with the meticulous setup of a labeling process, a foundational step crucial for training our AI models. Imagine teaching a child to recognize various shapes; the clarity and diversity of examples you provide are pivotal. Similarly, to ensure our model could accurately detect and locate five types of signatures in a PDF document, we had to eliminate biases and human errors from our dataset.
Once we had a gold-standard dataset, the development of our object detection models began. This wasn't a sprint; it was a marathon with hurdles. It took us quite a awhile to train a performant model, or one that our business stakeholders are ok with. The model of choice was FasterRCNN, renowned for its efficacy in object detection tasks, which we adapted to our specific need: locating signatures on a page-wise image.
I'm going to briefly describe how we deploy this model on our existing Azure Kubernetes Service (AKS) cluster
Our Kubernetes Cluster Setup
Our infrastructure leverages a comprehensive stack centered around Azure Kubernetes Service (AKS), which provides a managed Kubernetes environment. This is a pretty neat setup i'd say. A managed Kubernetes cluster is generally easier to set up and maintain, reduces operational costs, and provides improved security. This leaves us with additional time to worry about our application deployments, helm chart definitions, MLOps pipelines, etc.
Our stack is designed to support a ML-powered APIs and applications. A brief breakdown of our stack is as follows:
- AKS: Supports a mix of MLOps and SRE tasks
- Centralized EFK Logging Stack: Comprising Elasticsearch, Fluentd, and Kibana, it provides comprehensive logging and monitoring capabilities
- Filebeat for Log Collection: Gathers logs from stdout across all production pods, ensuring no data is missed
- Data orchestration with Apache Airflow: Manages complex workflows, enhancing our ability to orchestrate ML pipelines efficiently
- Advanced Search with OpenSearch: Facilitates efficient data retrieval and management, crucial for powering our in-house search engines
Here, AKS isn't merely a tool but a framework enabling our AI models to serve their purpose across multiple domains at scale.
Signature detection model
Training an object detection model itself isn't rocket science, with plenty of open-source resources available. The real challenge lay in managing business stakeholder expectations. Data scientists often face the uphill battle of explaining the inherent uncertainty and potential biases in these models, which can clash with stakeholders' desire for definitive guarantees. It's been an iterative process of continuous performance improvement, all while navigating the sometimes tricky terrain of expectation management.
Looking back at our project roadmap:
- Scope requirements with business stakeholders
- Collect raw unlabelled data from stakeholders
- Design and implement labelling effort to label all pdfs, i.e. draw bounding boxes around each page in pdfs.
- Build a POC to prove the feasibility of signature detection model, that 1) detects signature type and 2) draw a bounding box around each detected signature
- Get buy-in and official nod of approval from business
- Improve model to fit stakeholder expectations, i.e. hyperparamter tuning, trying different NN architectures, error analysis, etc.
- Scope requirements to integrate model into existing applications, e.g. back-of-envelop calculations, set service level objectives and indicators, etc.
- Local deploy model with Torchserve, conduct load testing
- Deploy model on tech stack
- Validate monitoring, alerting, logging are in-place, incident response SOPs, canary deployments, etc.
Of course, we did it. Our best models' have 91% mAPs across 5 different types of signatures: handwritten, docusign, stamp, typed, acknowledgement. Our best model is a Faster-R-CNN model.
Faster R-CNN is a state-of-the-art deep learning model for object detection in images and videos. It builds upon previous architectures like R-CNN and Fast R-CNN, addressing their limitations and achieving significant improvements in speed and accuracy.
Using Torchserve, FastAPI, Celery, Redis
Deploying our FasterRCNN signature detection model within our Azure Kubernetes Service (AKS) cluster presented a unique set of challenges and opportunities, especially considering the critical role this model plays in streamlining document review processes for corporate lawyers. Our deployment strategy lies a stack composed of TorchServe, FastAPI, Celery, and Redis, each selected for its specific strengths in handling model serving, request handling, task queuing, and caching, respectively.
TorchServe was our go-to for serving the FasterRCNN model, mainly because of its seamless integration with PyTorch models, enabling us to efficiently manage model versions and handle concurrent inference requests. This choice allowed us to leverage TorchServe's built-in functionalities like model snapshotting, logging, and metrics collection, which are crucial for maintaining the model's performance and ensuring its reliability in a production environment.
FastAPI came into play as the interface layer, acting as the bridge between our model and the end-users. Its asynchronous request handling capabilities meant we could serve more requests at lower latency, a critical factor when dealing with high volumes of document scans for signature detection.
For managing background tasks, such as spliting up the PDFs into images and applying inference on them, or handling intensive I/O operations without blocking the main request-response cycle, we employed Celery with Redis as the message broker. This combination allowed us to offload heavy tasks, ensuring that the API remains responsive and that tasks are processed efficiently in the background. Celery's distributed nature was particularly beneficial, allowing us to scale our processing power horizontally by adding more workers as the load increases.
The integration of Redis served multiple purposes: it acted as the backend for Celery, enabling task queuing and result storage, and also provided a caching layer for our application. By caching frequently accessed data, such as precomputed signature detection results for common document templates, we significantly reduced the load on our model and decreased response times, enhancing the overall user experience.
Example output from API
{
"task_id": "e4ceedce-41b9-445e-93b1-dd364d552bf8",
"status": "SUCCESS",
"results": [
{
"page_num": "1",
"label": "DOCUSIGN",
"bbox": [
370.5526123046875,
1026.4556884765625,
532.9699096679688,
1074.72802734375
],
"score": 0.9958266615867615
},
{
"page_num": "1",
"label": "DOCUSIGN",
"bbox": [
196.8037109375,
1025.1981201171875,
349.9497985839844,
1078.6612548828125
],
"score": 0.9925532937049866
},
{
"page_num": "1",
"label": "STAMPSIGNATURE",
"bbox": [
692.742431640625,
954.4534912109375,
742.327392578125,
1015.265625
],
"score": 0.706244945526123
},
{
"page_num": "1",
"label": "DOCUSIGN",
"bbox": [
302.35113525390625,
1030.09326171875,
341.5270690917969,
1074.838134765625
],
"score": 0.6967933773994446
},
{
"page_num": "1",
"label": "DOCUSIGN",
"bbox": [
270.1528015136719,
1023.7876586914062,
339.8675842285156,
1078.1610107421875
],
"score": 0.6473340392112732
}
]
}Deploying this stack on AKS required careful consideration of resource allocation to balance CPU, I/O, and latency requirements. AKS facilitated this by allowing us to specify resource requests and limits for each component, ensuring that TorchServe had enough computational power for model inference, while Celery workers were allocated resources that matched their I/O and CPU-intensive task profiles. This fine-grained control over resources, combined with AKS's auto-scaling capabilities, meant our deployment could adapt to varying loads, maintaining high availability and consistent performance.
Reflecting on this deployment, the synergy between these technologies and AKS highlights the importance of choosing the right tools and platforms that complement each other. This stack not only provided the flexibility and scalability needed to meet our technical and business requirements but also exemplified how modern cloud-native technologies could be leveraged to deploy complex AI models in a production environment efficiently.