EGNN - Pytorch

Install

$ pip install egnn-pytorch

Usage

importtorchfromegnn_pytorchimportEGNNlayer1=EGNN(dim=512)layer2=EGNN(dim=512)feats=torch.randn(1,16,512)coors=torch.randn(1,16,3)feats,coors=layer1(feats,coors)feats,coors=layer2(feats,coors)# (1, 16, 512), (1, 16, 3)

importtorchfromegnn_pytorchimportEGNNlayer1=EGNN(dim=512,edge_dim=4)layer2=EGNN(dim=512,edge_dim=4)feats=torch.randn(1,16,512)coors=torch.randn(1,16,3)edges=torch.randn(1,16,16,4)feats,coors=layer1(feats,coors,edges)feats,coors=layer2(feats,coors,edges)# (1, 16, 512), (1, 16, 3)

importtorchfromegnn_pytorchimportEGNN_Networknet=EGNN_Network(num_tokens=21,num_positions=1024,# unless what you are passing in is an unordered set, set this to the maximum sequence lengthdim=32,depth=3,num_nearest_neighbors=8,coor_weights_clamp_value=2.# absolute clamped value for the coordinate weights, needed if you increase the num neareest neighbors)feats=torch.randint(0,21, (1,1024))# (1, 1024)coors=torch.randn(1,1024,3)# (1, 1024, 3)mask=torch.ones_like(feats).bool()# (1, 1024)feats_out,coors_out=net(feats,coors,mask=mask)# (1, 1024, 32), (1, 1024, 3)

importtorchfromegnn_pytorchimportEGNN_Networknet=EGNN_Network(num_tokens=21,dim=32,depth=3,only_sparse_neighbors=True)feats=torch.randint(0,21, (1,1024))coors=torch.randn(1,1024,3)mask=torch.ones_like(feats).bool()# naive adjacency matrix# assuming the sequence is connected as a chain, with at most 2 neighbors - (1024, 1024)i=torch.arange(1024)adj_mat= (i[:,None]>= (i[None, :]-1))& (i[:,None]<= (i[None, :]+1))feats_out,coors_out=net(feats,coors,mask=mask,adj_mat=adj_mat)# (1, 1024, 32), (1, 1024, 3)

importtorchfromegnn_pytorchimportEGNN_Networknet=EGNN_Network(num_tokens=21,dim=32,depth=3,num_adj_degrees=3,# fetch up to 3rd degree neighborsadj_dim=8,# pass an adjacency degree embedding to the EGNN layer, to be used in the edge MLPonly_sparse_neighbors=True)feats=torch.randint(0,21, (1,1024))coors=torch.randn(1,1024,3)mask=torch.ones_like(feats).bool()# naive adjacency matrix# assuming the sequence is connected as a chain, with at most 2 neighbors - (1024, 1024)i=torch.arange(1024)adj_mat= (i[:,None]>= (i[None, :]-1))& (i[:,None]<= (i[None, :]+1))feats_out,coors_out=net(feats,coors,mask=mask,adj_mat=adj_mat)# (1, 1024, 32), (1, 1024, 3)

Edges

importtorchfromegnn_pytorchimportEGNN_Networknet=EGNN_Network(num_tokens=21,dim=32,depth=3,edge_dim=4,num_nearest_neighbors=3)feats=torch.randint(0,21, (1,1024))coors=torch.randn(1,1024,3)mask=torch.ones_like(feats).bool()continuous_edges=torch.randn(1,1024,1024,4)# naive adjacency matrix# assuming the sequence is connected as a chain, with at most 2 neighbors - (1024, 1024)i=torch.arange(1024)adj_mat= (i[:,None]>= (i[None, :]-1))& (i[:,None]<= (i[None, :]+1))feats_out,coors_out=net(feats,coors,edges=continuous_edges,mask=mask,adj_mat=adj_mat)# (1, 1024, 32), (1, 1024, 3)

Stability

importtorchfromegnn_pytorchimportEGNN_Networknet=EGNN_Network(num_tokens=21,dim=32,depth=3,num_nearest_neighbors=32,norm_coors=True,# normalize the relative coordinatescoor_weights_clamp_value=2.# absolute clamped value for the coordinate weights, needed if you increase the num neareest neighbors)feats=torch.randint(0,21, (1,1024))# (1, 1024)coors=torch.randn(1,1024,3)# (1, 1024, 3)mask=torch.ones_like(feats).bool()# (1, 1024)feats_out,coors_out=net(feats,coors,mask=mask)# (1, 1024, 32), (1, 1024, 3)

All parameters

importtorchfromegnn_pytorchimportEGNNmodel=EGNN(dim=dim,# input dimensionedge_dim=0,# dimension of the edges, if exists, should be > 0m_dim=16,# hidden model dimensionfourier_features=0,# number of fourier features for encoding of relative distance - defaults to none as in papernum_nearest_neighbors=0,# cap the number of neighbors doing message passing by relative distancedropout=0.0,# dropoutnorm_feats=False,# whether to layernorm the featuresnorm_coors=False,# whether to normalize the coordinates, using a strategy from the SE(3) Transformers paperupdate_feats=True,# whether to update features - you can build a layer that only updates one or the otherupdate_coors=True,# whether ot update coordinatesonly_sparse_neighbors=False,# using this would only allow message passing along adjacent neighbors, using the adjacency matrix passed invalid_radius=float('inf'),# the valid radius each node considers for message passingm_pool_method='sum',# whether to mean or sum pool for output node representationsoft_edges=False,# extra GLU on the edges, purportedly helps stabilize the network in updated version of the papercoor_weights_clamp_value=None# clamping of the coordinate updates, again, for stabilization purposes)

Examples

$ pip install sidechainnet

$ python denoise_sparse.py

Tests

$ python setup.pytest

Citations

@misc{satorras2021en,title ={E(n) Equivariant Graph Neural Networks},author ={Victor Garcia Satorras and Emiel Hoogeboom and Max Welling},year ={2021},eprint ={2102.09844},archivePrefix ={arXiv},primaryClass ={cs.LG}}