Train graph neural nets for millions of proteins on Amazon SageMaker and Amazon DocumentDB (with MongoDB compatibility)

There are over 180,000 unique proteins with 3D structures determined, with tens of thousands new structures resolved every year. This is only a small fraction of the 200 million known proteins with distinctive sequences. Recent deep learning algorithms such as AlphaFold 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 for training GNNs. However, training GNN on millions of protein structures is challenging. Data serialization isn’t scalable to millions of protein structures because it requires loading the entire terabyte-scale dataset into memory.

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) using Amazon SageMaker.

For illustrative purposes, we use publicly available experimentally determined protein structures from the Protein Data Bank and computationally predicted protein structures from the AlphaFold Protein Structure Database. 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.

Overview of solution

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’ atomic coordinates, properties, and identifiers. Storing a JSON record for a protein’s 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 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.

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.

Finally, we analyze the trained GNN model to gain some insights into the predictions.

We walk through the following steps for this tutorial:

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

The code and notebooks used in this post are available in the GitHub repo.

Prerequisites

For this walkthrough, you should have the following prerequisites:

• An AWS account
• Permissions to create AWS resources including SageMaker, Amazon DocumentDB, Amazon Virtual Private Cloud (Amazon VPC), and AWS Secrets Manager
• Knowledge in ML model training and validation

Running this tutorial for an hour should cost no more than $2.00.

Create resources

We provide a CloudFormation template 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. For instructions on creating a CloudFormation stack, see the video Simplify your Infrastructure Management using AWS CloudFormation.

The CloudFormation stack provisions the following:

• A VPC with three private subnets for Amazon DocumentDB and two public subnets intended for the SageMaker notebook instance and ML training containers, respectively.
• An Amazon DocumentDB cluster with three nodes, one in each private subnet.
• A Secrets Manager secret to store login credentials for Amazon DocumentDB. This allows us to avoid storing plaintext credentials in our SageMaker instance.
• A SageMaker notebook instance to prepare data, orchestrate training jobs, and run interactive analyses.

When creating the CloudFormation stack, you need to specify the following:

• Name for your CloudFormation stack
• Amazon DocumentDB user name and password (to be stored in Secrets Manager)
• Amazon DocumentDB instance type (default db.r5.large)
• SageMaker instance type (default ml.t3.xlarge)

It should take about 15 minutes to create the CloudFormation stack. The following diagram shows the resource architecture.

Prepare protein structures and properties and ingest the data into Amazon DocumentDB

All the subsequent code in this section is in the Jupyter notebook Prepare_data.ipynb in the SageMaker instance created in your CloudFormation stack.

This notebook handles the procedures required for preparing and ingesting protein structure data into Amazon DocumentDB.

1. We first download predicted protein structures from AlphaFold DB in PDB format and the matching experimental structures from the Protein Data Bank.

For demonstration purposes, we only use proteins from the thermophilic archaean Methanocaldococcus jannaschii, which has the smallest proteome of 1,773 proteins for us to work with. You are welcome to try using proteins from other species.

2. We connect to an Amazon DocumentDB cluster by retrieving the credentials stored in Secrets Manager:

def get_secret(stack_name):

    # Create a Secrets Manager client
    session = boto3.session.Session()
    client = session.client(
        service_name="secretsmanager",
        region_name=session.region_name
    )
    
    secret_name = f"{stack_name}-DocDBSecret"
    get_secret_value_response = client.get_secret_value(SecretId=secret_name)
    secret = get_secret_value_response["SecretString"]
    
return json.loads(secret)
	
	secrets = get_secret("gnn-proteins")
	
# connect to DocDB
	uri = "mongodb://{}:{}@{}:{}/?tls=true&tlsCAFile=rds-combined-ca-bundle.pem&replicaSet=rs0&readPreference=secondaryPreferred&retryWrites=false"\
    		.format(secrets["username"], secrets["password"], secrets["host"], secrets["port"])
	
client = MongoClient(uri)

db = client["proteins"] # create a database
collection = db["proteins"] # create a collection

3. After we set up the connection to Amazon DocumentDB, we parse the PDB files into JSON records to ingest into the database.

We provide utility functions required for parsing PDB files in pdb_parse.py. The parse_pdb_file_to_json_record function does the heavy lifting of extracting atomic coordinates from one or multiple peptide chains in a PDB file and returns one or a list of JSON documents, which can be directly ingested into the Amazon DocumentDB collection as a document. See the following code:

recs = parse_pdb_file_to_json_record(pdb_parser, pdb_file, pdb_id)
collection.insert_many(recs)

After we ingest the parsed protein data into Amazon DocumentDB, we can update the contents of the protein documents. For instance, it makes our model training logistics easier if we add a field indicating whether a protein structure should be used in the training, validation, or test sets.

4. We first retrieve the all the documents with the field is_AF to stratify documents using an aggregation pipeline:

match = {"is_AF": {"$exists": True}}
project = {"y": "$is_AF"}

pipeline = [
    {"$match": match},
    {"$project": project},
]
# aggregation pipeline
cur = collection.aggregate(pipeline)
# retrieve documents from the DB cursor
docs = [doc for doc in cur]
# convert to a data frame:
df = pd.DataFrame(docs)
# stratified split: full -> train/test
df_train, df_test = train_test_split(
    df, 
    test_size=0.2,
    stratify=df["y"], 
    random_state=42
)
# stratified split: train -> train/valid
df_train, df_valid = train_test_split(
    df_train, 
    test_size=0.2,
    stratify=df_train["y"], 
    random_state=42
)

5. Next, we use the update_many function to store the split information back to Amazon DocumentDB:

for split, df_split in zip(
    ["train", "valid", "test"],
    [df_train, df_valid, df_test]
):
    result = collection.update_many(
        {"_id": {"$in": df_split["_id"].tolist()}}, 
        {"$set": {"split": split}}
)
print("Number of documents modified:", result.modified_count)

Train a GNN on the protein structures using SageMaker

All the subsequent code in this section is in the Train_and_eval.ipynb notebook in the SageMaker instance created in your CloudFormation stack.

This notebook trains a GNN model on the protein structure datasets stored in the Amazon DocumentDB.

We first need to implement a PyTorch dataset class for our protein dataset capable of retrieving mini-batches of protein documents from Amazon DocumentDB. It’s more efficient to retrieve batches documents by the built-in primary id (_id).

1. We use the iterable-style dataset by extending the IterableDataset, which pre-fetches the _id and labels of the documents at initialization:

class ProteinDataset(data.IterableDataset):
    """
    An iterable-style dataset for proteins in DocumentDB
    Args:
        pipeline: an aggregation pipeline to retrieve data from DocumentDB
        db_uri: URI of the DocumentDB
        db_name: name of the database
        collection_name: name of the collection
        k: k used for kNN when creating a graph from atomic coordinates
    """

    def __init__(
        self, pipeline, db_uri="", db_name="", collection_name="", k=3
    ):

        self.db_uri = db_uri
        self.db_name = db_name
        self.collection_name = collection_name
        self.k = k

        client = MongoClient(self.db_uri, connect=False)
        collection = client[self.db_name][self.collection_name]
        # pre-fetch the metadata as docs from DocumentDB
        self.docs = [doc for doc in collection.aggregate(pipeline)]
        # mapping document '_id' to label
        self.labels = {doc["_id"]: doc["y"] for doc in self.docs}

2. The ProteinDataset performs a database read operation in the __iter__ method. It tries to evenly split the workload if there are multiple workers:

def __iter__(self):
        worker_info = torch.utils.data.get_worker_info()
        if worker_info is None:
            # single-process data loading, return the full iterator
            protein_ids = [doc["_id"] for doc in self.docs]

        else:  # in a worker process
            # split workload
            start = 0
            end = len(self.docs)
            per_worker = int(
                math.ceil((end - start) / float(worker_info.num_workers))
            )
            worker_id = worker_info.id
            iter_start = start + worker_id * per_worker
            iter_end = min(iter_start + per_worker, end)

            protein_ids = [
                doc["_id"] for doc in self.docs[iter_start:iter_end]
            ]

        # retrieve a list of proteins by _id from DocDB
        with MongoClient(self.db_uri) as client:
            collection = client[self.db_name][self.collection_name]
            cur = collection.find(
                {"_id": {"$in": protein_ids}},
                projection={"coords": True, "seq": True},
            )
            return (
                (
                    convert_to_graph(protein, k=self.k),
                    self.labels[protein["_id"]],
                )
                for protein in cur
            )

3. The preceding __iter__ method also converts the atomic coordinates of proteins into DGLGraph objects after they’re loaded from Amazon DocumentDB via the convert_to_graph function. This function constructs a k-nearest neighbor (kNN) graph for the amino acid residues using the 3D coordinates of the C-alpha atoms and adds one-hot encoded node features to represent residue identities:

def convert_to_graph(protein, k=3):
    """
    Convert a protein (dict) to a dgl graph using kNN.
    """
    coords = torch.tensor(protein["coords"])
    X_ca = coords[:, 1]
    # construct knn graph from C-alpha coordinates
    g = dgl.knn_graph(X_ca, k=k)
    seq = protein["seq"]
    node_features = torch.tensor([d1_to_index[residue] for residue in seq])
    node_features = F.one_hot(node_features, num_classes=len(d1_to_index)).to(
        dtype=torch.float
    )

    # add node features
    g.ndata["h"] = node_features
    return g

4. With the ProteinDataset implemented, we can initialize instances for train, validation, and test datasets and wrap the training instance with BufferedShuffleDataset to enable shuffling.
5. We further wrap them with torch.utils.data.DataLoader to work with other components of the SageMaker PyTorch Estimator training script.
6. Next, we implement a simple two-layered graph convolution network (GCN) with a global attention pooling layer for ease of interpretation:

class GCN(nn.Module):
    """A two layer Graph Conv net with Global Attention Pooling over the
    nodes.
    Args:
        in_feats: int, dim of input node features
        h_feats: int, dim of hidden layers
        num_classes: int, number of output units
    """

    def __init__(self, in_feats, h_feats, num_classes):
        super(GCN, self).__init__()
        self.conv1 = GraphConv(in_feats, h_feats)
        self.conv2 = GraphConv(h_feats, h_feats)
        # the gate layer that maps node feature to outputs
        self.gate_nn = nn.Linear(h_feats, num_classes)
        self.gap = GlobalAttentionPooling(self.gate_nn)
        # the output layer making predictions
        self.output = nn.Linear(h_feats, num_classes)

    def _conv_forward(self, g):
        """forward pass through the GraphConv layers"""
        in_feat = g.ndata["h"]
        h = self.conv1(g, in_feat)
        h = F.relu(h)
        h = self.conv2(g, h)
        h = F.relu(h)
        return h

    def forward(self, g):
        h = self._conv_forward(g)
        h = self.gap(g, h)
        return self.output(h)

    def attention_scores(self, g):
        """Calculate attention scores"""
        h = self._conv_forward(g)
        with g.local_scope():
            gate = self.gap.gate_nn(h)
            g.ndata["gate"] = gate
            gate = dgl.softmax_nodes(g, "gate")
            g.ndata.pop("gate")
            return gate

7. Afterwards, we can train this GCN on the ProteinDataset instance for a binary classification task of predicting whether a protein structure is predicted by AlphaFold or not. We use binary cross entropy as the objective function and Adam optimizer for stochastic gradient optimization. The full training script can be found in src/main.py.

Next, we set up the SageMaker PyTorch Estimator to handle the training job. To allow the managed Docker container initiated by SageMaker to connect to Amazon DocumentDB, we need to configure the subnet and security group for the Estimator.

8. We retrieve the subnet ID where the Network Address Translation (NAT) gateway resides, as well as the security group ID of our Amazon DocumentDB cluster by name:

ec2 = boto3.client("ec2")
# find the NAT gateway's subnet ID 
resp = ec2.describe_subnets(
    Filters=[{"Name": "tag:Name", "Values": ["{}-NATSubnet".format(stack_name)]}]
)
nat_subnet_id = resp["Subnets"][0]["SubnetId"]
# find security group id of the DocumentDB
resp = ec2.describe_security_groups(
    Filters=[{
        "Name": "tag:Name", 
        "Values": ["{}-SG-DocumentDB".format(stack_name)]
    }])
sg_id = resp["SecurityGroups"][0]["GroupId"]
Finally, we can kick off the training of our GCN model using SageMaker: 
from sagemaker.pytorch import PyTorch

CODE_PATH = "main.py"

params = {
    "patience": 5, 
    "n-epochs": 200,
    "batch-size": 64,
    "db-host": secrets["host"],
    "db-username": secrets["username"], 
    "db-password": secrets["password"], 
    "db-port": secrets["port"],
    "knn": 4,
}

estimator = PyTorch(
    entry_point=CODE_PATH,
    source_dir="src",
    role=role,
    instance_count=1,
    instance_type="ml.p3.2xlarge", # 'ml.c4.2xlarge' for CPU
    framework_version="1.7.1",
    py_version="py3",
    hyperparameters=params,
    sagemaker_session=sess,
    subnets=[nat_subnet_id], 
    security_group_ids=[sg_id],
)
# run the training job:
estimator.fit()

Load and evaluate the trained GNN model

When the training job is complete, we can load the trained GCN model and perform some in-depth evaluation.

The codes for the following steps are also available in the notebook Train_and_eval.ipynb.

SageMaker training jobs save the model artifacts into the default S3 bucket, the URI of which can be accessed from the estimator.model_data attribute. We can also navigate to the Training jobs page on the SageMaker console to find the trained model to evaluate.

1. For research purposes, we can load the model artifact (learned parameters) into a PyTorch state_dict using the following function:

def load_sagemaker_model_artifact(s3_bucket, key):
    """Load a PyTorch model artifact (model.tar.gz) produced by a SageMaker
    Training job.
    Args:
        s3_bucket: str, s3 bucket name (s3://bucket_name)
        key: object key: path to model.tar.gz from within the bucket
    Returns:
        state_dict: dict representing the PyTorch checkpoint
    """
    # load the s3 object
    s3 = boto3.client("s3")
    obj = s3.get_object(Bucket=s3_bucket, Key=key)
    # read into memory
    model_artifact = BytesIO(obj["Body"].read())
    # parse out the state dict from the tar.gz file
    tar = tarfile.open(fileobj=model_artifact)
    for member in tar.getmembers():
        pth = tar.extractfile(member).read()

    state_dict = torch.load(BytesIO(pth), map_location=torch.device("cpu"))
return state_dict

	state_dict = load_sagemaker_model_artifact(
bucket, 
key=estimator.model_data.split(bucket)[1].lstrip("/")
)

# initialize a GCN model
model = GCN(dim_nfeats, 16, n_classes)
# load the learned parameters
model.load_state_dict(state_dict["model_state_dict"])

2. Next, we perform quantitative model evaluation on the full test set by calculating accuracy:

device = torch.device("cuda:0") if torch.cuda.is_available() else torch.device("cpu")
num_correct = 0
num_tests = 0
model.eval()
with torch.no_grad():
    for batched_graph, labels in test_loader:
        batched_graph = batched_graph.to(device)
        labels = labels.to(device)
        logits = model(batched_graph)
        preds = (logits.sigmoid() > 0.5).to(labels.dtype)
        num_correct += (preds == labels).sum().item()
        num_tests += len(labels)

print('Test accuracy: {:.6f}'.format(num_correct / num_tests))

We found our GCN model achieved an accuracy of 74.3%, whereas the dummy baseline model making predictions based on class priors only achieved 56.3%.

We’re also interested in interpretability of our GCN model. Because we implement a global attention pooling layer, we can compute the attention scores across nodes to explain specific predictions made by the model.

3. Next, we compute the attention scores and overlay them on the protein graphs for a pair of structures (AlphaFold predicted and experimental) from the same peptide:

pair = ["AF-Q57887", "1JT8-A"]
cur = collection.find(
    {"id": {"$in": pair}},
)

for doc in cur:
    # convert to dgl.graph object
    graph = convert_to_graph(doc, k=4)
    
    with torch.no_grad():
        # make prediction
        pred = model(graph).sigmoid()
        # calculate attention scores for a protein graph
        attn = model.attention_scores(graph)
    
    pred = pred.item()
    attn = attn.numpy()
    
    # convert to networkx graph for visualization
    graph = graph.to_networkx().to_undirected()
    # calculate graph layout
    pos = nx.spring_layout(graph, iterations=500)
    
    fig, ax = plt.subplots(figsize=(8, 8))
    nx.draw(
        graph, 
        pos, 
        node_color=attn.flatten(),
        cmap="Reds",
        with_labels=True, 
        font_size=8,
        ax=ax
    )
    ax.set(title="{}, p(is_predicted)={:.6f}".format(doc["id"], pred))
plt.show()

The preceding codes produce the following protein graphs overlaid with attention scores on the nodes. We find the model’s global attentive pooling layer can highlight certain residues in the protein graph as being important for making the prediction of whether the protein structure is predicted by AlphaFold. This indicates that these residues may have distinctive graph topologies in predicted and experimental protein structures.

In summary, we showcase a scalable deep learning solution to train GNNs on protein structures stored in Amazon DocumentDB. Although the tutorial only uses thousands of proteins for training, this solution is scalable to millions of proteins. Unlike other approaches such as serializing the entire protein dataset, our approach transfers the memory-heavy workloads to the database, making the memory complexity for the training jobs O(batch_size), which is independent of the total number of proteins to train.

Clean up

To avoid incurring future charges, delete the CloudFormation stack you created. This removes all the resources you provisioned using the CloudFormation template, including the VPC, Amazon DocumentDB cluster, and SageMaker instance. For instructions, see Deleting a stack on the AWS CloudFormation console.

Conclusion

We described a cloud-based deep learning architecture scalable to millions of protein structures by storing them in Amazon DocumentDB and efficiently retrieving mini-batches of data from SageMaker.

To learn more about the use of GNN in protein property predictions, check out our recent publication LM-GVP, A Generalizable Deep Learning Framework for Protein Property Prediction from Sequence and Structure.

About the Authors

Zichen Wang, PhD, is an Applied Scientist in the Amazon Machine Learning Solutions Lab. With several years of research experience in developing ML and statistical methods using biological and medical data, he works with customers across various verticals to solve their ML problems.

Selvan Senthivel is a Senior ML Engineer with the Amazon ML Solutions Lab at AWS, focusing on helping customers on machine learning, deep learning problems, and end-to-end ML solutions. He was a founding engineering lead of Amazon Comprehend Medical and contributed to the design and architecture of multiple AWS AI services.