Atrial Fibrillation Classification

Technical Details

System Architecture

Our overall architecture relies heavily on Google Colab for initial model development and AWS SageMaker for fine-tuning and deployment. The final model parameters were stored in an S3 bucket to be deployed using a SageMaker Endpoint. The image for the containerized app is stored in AWS ECR, which will pull the model parameters, instantiate S3 buckets for temporary data storage, and create a SageMaker endpoint for model inference. We deployed the Dash app using Fargate. We registered our domain using Route 53 and leveraged a load balancer to manage security and traffic flow while connecting our dashboard to the domain name.

Deployment Architecture

AWS Cloud Amazon SageMaker LSTM Model Training Training Jobs Hyperparameter Tuning Model Artifacts Model Registry Version Control Model Metadata Approval Workflow SageMaker Endpoint Real-time Inference Auto-scaling Model Monitoring S3 Bucket Training Data Model Artifacts API Gateway REST API Authentication Web Application Dash/Flask Frontend

Architecture Overview

Our LSTM model is deployed using AWS SageMaker, providing a scalable and robust inference solution. The architecture includes:

  • Model Training: LSTM models are trained using SageMaker training jobs with hyperparameter optimization.
  • Model Registry: Trained models are versioned and stored in the SageMaker Model Registry.
  • Endpoint Deployment: Approved models are deployed to a SageMaker endpoint for real-time inference.
  • API Gateway: Provides secure RESTful API access to the model endpoint.
  • Web Application: A Dash/Flask application serves as the user interface, making API calls to the model endpoint.

This architecture enables scalable, secure, and reliable access to our LSTM model predictions while maintaining version control and monitoring capabilities.

Key Components

AWS SageMaker

Used to support model deployment.

AWS S3 Buckets

Model artifact and data storage

AWS ECR/ECS

Deployed app to the web using ECS via images hosted in ECR

Technology Stack

SageMaker Tensorflow Keras Google Colab Jupyter

Challenges & Solutions

Data

We had to find a way to process vast amounts of signals data to develop the model in a standardized, replicable, and easy-to-interpret way. We addressed this by developing a scalable and secure deployment pipeline to handle different data type inputs, clean/standardize signals, and generate classifications all in once place.

Accuracy and F1 score

We had to balance accurately identifying the type of heartbeat rhythm with mitigating the risk of a false positive, which may cause unnecessary anxiety, and of a false negative, which could create a dangerous false sense of security. We addressed this and achieved close to SoTA accuracy by combining a K-means smoothing layer with our baseline LSTM model, which led to substantial improvements in out-of-sample accuracy.

Deployment

We wanted to find a way to give users access to their insights in a way that was easy-to-use, secure, and had low(ish) latency. We addressed this by tying model classifications back to timestamped heartbeats in a user-friendly way, providing utility for patients to better understand their health alongside commonly used metrics that cardiologists use for diagnosis.

Future Work

  • Fine-tuning on individual users as they upload more data

  • Reduce latency in various places for better UX

  • User account management for data to persist across sessions

  • Incorporate trend analysis on AFib episodes into the dashboard