Amazon SageMaker および Amazon DocumentDB (MongoDB 互換性あり) で数百万のタンパク質のグラフ ニューラル ネットをトレーニングする

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

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

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
)

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)

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}

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 )

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

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

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()

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"])

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))

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()