{"id":1437,"date":"2022-01-04T19:41:02","date_gmt":"2022-01-04T19:41:02","guid":{"rendered":"https:\/\/salarydistribution.com\/machine-learning\/2022\/01\/04\/train-graph-neural-nets-for-millions-of-proteins-on-amazon-sagemaker-and-amazon-documentdb-with-mongodb-compatibility\/"},"modified":"2022-01-04T19:41:02","modified_gmt":"2022-01-04T19:41:02","slug":"train-graph-neural-nets-for-millions-of-proteins-on-amazon-sagemaker-and-amazon-documentdb-with-mongodb-compatibility","status":"publish","type":"post","link":"https:\/\/salarydistribution.com\/machine-learning\/2022\/01\/04\/train-graph-neural-nets-for-millions-of-proteins-on-amazon-sagemaker-and-amazon-documentdb-with-mongodb-compatibility\/","title":{"rendered":"Train graph neural nets for millions of proteins on Amazon SageMaker and Amazon DocumentDB (with MongoDB compatibility)"},"content":{"rendered":"
\n

There are over 180,000 unique proteins with 3D structures determined, with tens of thousands new structures resolved every year<\/a>. This is only a small fraction of the 200 million known proteins with distinctive sequences<\/a>. Recent deep learning algorithms such as AlphaFold<\/a> can accurately predict 3D structures of proteins using their sequences, which help scale the protein 3D structure data to the millions. Graph neural network (GNN) has emerged as an effective deep learning approach to extract information from protein structures, which can be represented by graphs of amino acid residues. Individual protein graphs usually contain a few hundred nodes, which is manageable in size. Tens of thousands of protein graphs can be easily stored in serialized data structures such as TFrecord<\/a> for training GNNs. However, training GNN on millions of protein structures is challenging. Data serialization isn\u2019t scalable to millions of protein structures because it requires loading the entire terabyte-scale dataset into memory.<\/p>\n

In this post, we introduce a scalable deep learning solution that allows you to train GNNs on millions of proteins stored in Amazon DocumentDB (with MongoDB compatibility)<\/a> using Amazon SageMaker<\/a>.<\/p>\n

For illustrative purposes, we use publicly available experimentally determined protein structures from the Protein Data Bank<\/a> and computationally predicted protein structures from the AlphaFold Protein Structure Database<\/a>. The machine learning (ML) problem is to develop a discriminator GNN model to distinguish between experimental and predicted structures based on protein graphs constructed from their 3D structures.<\/p>\n

Overview of solution<\/h2>\n

We first parse the protein structures into JSON records with multiple types of data structures, such as an n-dimensional array and nested object, to store the proteins\u2019 atomic coordinates, properties, and identifiers. Storing a JSON record for a protein\u2019s structure takes 45 KB on average; we project storing 100 million proteins would take around 4.2 TB. Amazon DocumentDB storage automatically scales with the data<\/a> in your cluster volume in 10 GB increments, up to 64 TB. Therefore, the support for JSON data structure and scalability makes Amazon DocumentDB a natural choice.<\/p>\n

We next build a GNN model to predict protein properties using graphs of amino acid residues constructed from their structures. The GNN model is trained using SageMaker and configured to efficiently retrieve batches of protein structures from the database.<\/p>\n

Finally, we analyze the trained GNN model to gain some insights into the predictions.<\/p>\n

We walk through the following steps for this tutorial:<\/p>\n

    \n
  1. Create resources using an AWS CloudFormation<\/a> template.<\/li>\n
  2. Prepare protein structures and properties and ingest the data into Amazon DocumentDB.<\/li>\n
  3. Train a GNN on the protein structures using SageMaker.<\/li>\n
  4. Load and evaluate the trained GNN model.<\/li>\n<\/ol>\n

    The code and notebooks used in this post are available in the GitHub repo<\/a>.<\/p>\n

    Prerequisites<\/h2>\n

    For this walkthrough, you should have the following prerequisites:<\/p>\n

    Running this tutorial for an hour should cost no more than $2.00.<\/p>\n

    Create resources<\/h2>\n

    We provide a CloudFormation template<\/a> to create the required AWS resources for this post, with a similar architecture as in the post Analyzing data stored in Amazon DocumentDB (with MongoDB compatibility) using Amazon SageMaker<\/a>. For instructions on creating a CloudFormation stack, see the video Simplify your Infrastructure Management using AWS CloudFormation<\/a>.<\/p>\n

    The CloudFormation stack provisions the following:<\/p>\n