{"id":13869885,"url":"https://github.com/lucidrains/egnn-pytorch","last_synced_at":"2025-05-14T22:08:46.641Z","repository":{"id":37557192,"uuid":"342683465","full_name":"lucidrains/egnn-pytorch","owner":"lucidrains","description":"Implementation of E(n)-Equivariant Graph Neural Networks, in Pytorch","archived":false,"fork":false,"pushed_at":"2024-12-06T14:23:04.000Z","size":832,"stargazers_count":464,"open_issues_count":12,"forks_count":77,"subscribers_count":8,"default_branch":"main","last_synced_at":"2025-04-14T22:11:38.390Z","etag":null,"topics":["artificial-intelligence","deep-learning","equivariance","graph-neural-network"],"latest_commit_sha":null,"homepage":"","language":"Python","has_issues":true,"has_wiki":null,"has_pages":null,"mirror_url":null,"source_name":null,"license":"mit","status":null,"scm":"git","pull_requests_enabled":true,"icon_url":"https://github.com/lucidrains.png","metadata":{"files":{"readme":"README.md","changelog":null,"contributing":null,"funding":null,"license":"LICENSE","code_of_conduct":null,"threat_model":null,"audit":null,"citation":null,"codeowners":null,"security":null,"support":null,"governance":null,"roadmap":null,"authors":null,"dei":null,"publiccode":null,"codemeta":null}},"created_at":"2021-02-26T19:44:48.000Z","updated_at":"2025-04-14T09:14:45.000Z","dependencies_parsed_at":"2024-01-16T07:24:15.268Z","dependency_job_id":"628861ec-fb6c-4671-8ede-9ad7a6b97247","html_url":"https://github.com/lucidrains/egnn-pytorch","commit_stats":{"total_commits":134,"total_committers":6,"mean_commits":"22.333333333333332","dds":"0.32835820895522383","last_synced_commit":"3e0a2b82741c678d50b89a5e142c5868e45d14f3"},"previous_names":[],"tags_count":67,"template":false,"template_full_name":null,"repository_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/lucidrains%2Fegnn-pytorch","tags_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/lucidrains%2Fegnn-pytorch/tags","releases_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/lucidrains%2Fegnn-pytorch/releases","manifests_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/lucidrains%2Fegnn-pytorch/manifests","owner_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners/lucidrains","download_url":"https://codeload.github.com/lucidrains/egnn-pytorch/tar.gz/refs/heads/main","host":{"name":"GitHub","url":"https://github.com","kind":"github","repositories_count":254235700,"owners_count":22036964,"icon_url":"https://github.com/github.png","version":null,"created_at":"2022-05-30T11:31:42.601Z","updated_at":"2022-07-04T15:15:14.044Z","host_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub","repositories_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories","repository_names_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repository_names","owners_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners"}},"keywords":["artificial-intelligence","deep-learning","equivariance","graph-neural-network"],"created_at":"2024-08-05T20:01:20.793Z","updated_at":"2025-05-14T22:08:41.624Z","avatar_url":"https://github.com/lucidrains.png","language":"Python","funding_links":[],"categories":["Python","分子"],"sub_categories":["网络服务_其他"],"readme":"\u003cimg src=\"./egnn.png\" width=\"600px\"\u003e\u003c/img\u003e\n\n** A bug has been discovered with the neighbor selection in the presence of masking. If you ran any experiments prior to 0.1.12 that had masking, please rerun them. 🙏 **\n\n## EGNN - Pytorch\n\nImplementation of \u003ca href=\"https://arxiv.org/abs/2102.09844v1\"\u003eE(n)-Equivariant Graph Neural Networks\u003c/a\u003e, in Pytorch. May be eventually used for Alphafold2 replication. This technique went for simple invariant features, and ended up beating all previous methods (including SE3 Transformer and Lie Conv) in both accuracy and performance. SOTA in dynamical system models, molecular activity prediction tasks, etc.\n\n## Install\n\n```bash\n$ pip install egnn-pytorch\n```\n\n## Usage\n\n```python\nimport torch\nfrom egnn_pytorch import EGNN\n\nlayer1 = EGNN(dim = 512)\nlayer2 = EGNN(dim = 512)\n\nfeats = torch.randn(1, 16, 512)\ncoors = torch.randn(1, 16, 3)\n\nfeats, coors = layer1(feats, coors)\nfeats, coors = layer2(feats, coors) # (1, 16, 512), (1, 16, 3)\n```\n\nWith edges\n\n```python\nimport torch\nfrom egnn_pytorch import EGNN\n\nlayer1 = EGNN(dim = 512, edge_dim = 4)\nlayer2 = EGNN(dim = 512, edge_dim = 4)\n\nfeats = torch.randn(1, 16, 512)\ncoors = torch.randn(1, 16, 3)\nedges = torch.randn(1, 16, 16, 4)\n\nfeats, coors = layer1(feats, coors, edges)\nfeats, coors = layer2(feats, coors, edges) # (1, 16, 512), (1, 16, 3)\n```\n\nA full EGNN network\n\n```python\nimport torch\nfrom egnn_pytorch import EGNN_Network\n\nnet = EGNN_Network(\n    num_tokens = 21,\n    num_positions = 1024,           # unless what you are passing in is an unordered set, set this to the maximum sequence length\n    dim = 32,\n    depth = 3,\n    num_nearest_neighbors = 8,\n    coor_weights_clamp_value = 2.   # absolute clamped value for the coordinate weights, needed if you increase the num neareest neighbors\n)\n\nfeats = torch.randint(0, 21, (1, 1024)) # (1, 1024)\ncoors = torch.randn(1, 1024, 3)         # (1, 1024, 3)\nmask = torch.ones_like(feats).bool()    # (1, 1024)\n\nfeats_out, coors_out = net(feats, coors, mask = mask) # (1, 1024, 32), (1, 1024, 3)\n```\n\nOnly attend to sparse neighbors, given to the network as an adjacency matrix.\n\n```python\nimport torch\nfrom egnn_pytorch import EGNN_Network\n\nnet = EGNN_Network(\n    num_tokens = 21,\n    dim = 32,\n    depth = 3,\n    only_sparse_neighbors = True\n)\n\nfeats = torch.randint(0, 21, (1, 1024))\ncoors = torch.randn(1, 1024, 3)\nmask = torch.ones_like(feats).bool()\n\n# naive adjacency matrix\n# assuming the sequence is connected as a chain, with at most 2 neighbors - (1024, 1024)\ni = torch.arange(1024)\nadj_mat = (i[:, None] \u003e= (i[None, :] - 1)) \u0026 (i[:, None] \u003c= (i[None, :] + 1))\n\nfeats_out, coors_out = net(feats, coors, mask = mask, adj_mat = adj_mat) # (1, 1024, 32), (1, 1024, 3)\n```\n\nYou can also have the network automatically determine the Nth-order neighbors, and pass in an adjacency embedding (depending on the order) to be used as an edge, with two extra keyword arguments\n\n```python\nimport torch\nfrom egnn_pytorch import EGNN_Network\n\nnet = EGNN_Network(\n    num_tokens = 21,\n    dim = 32,\n    depth = 3,\n    num_adj_degrees = 3,           # fetch up to 3rd degree neighbors\n    adj_dim = 8,                   # pass an adjacency degree embedding to the EGNN layer, to be used in the edge MLP\n    only_sparse_neighbors = True\n)\n\nfeats = torch.randint(0, 21, (1, 1024))\ncoors = torch.randn(1, 1024, 3)\nmask = torch.ones_like(feats).bool()\n\n# naive adjacency matrix\n# assuming the sequence is connected as a chain, with at most 2 neighbors - (1024, 1024)\ni = torch.arange(1024)\nadj_mat = (i[:, None] \u003e= (i[None, :] - 1)) \u0026 (i[:, None] \u003c= (i[None, :] + 1))\n\nfeats_out, coors_out = net(feats, coors, mask = mask, adj_mat = adj_mat) # (1, 1024, 32), (1, 1024, 3)\n```\n\n## Edges\n\nIf you need to pass in continuous edges\n\n```python\nimport torch\nfrom egnn_pytorch import EGNN_Network\n\nnet = EGNN_Network(\n    num_tokens = 21,\n    dim = 32,\n    depth = 3,\n    edge_dim = 4,\n    num_nearest_neighbors = 3\n)\n\nfeats = torch.randint(0, 21, (1, 1024))\ncoors = torch.randn(1, 1024, 3)\nmask = torch.ones_like(feats).bool()\n\ncontinuous_edges = torch.randn(1, 1024, 1024, 4)\n\n# naive adjacency matrix\n# assuming the sequence is connected as a chain, with at most 2 neighbors - (1024, 1024)\ni = torch.arange(1024)\nadj_mat = (i[:, None] \u003e= (i[None, :] - 1)) \u0026 (i[:, None] \u003c= (i[None, :] + 1))\n\nfeats_out, coors_out = net(feats, coors, edges = continuous_edges, mask = mask, adj_mat = adj_mat) # (1, 1024, 32), (1, 1024, 3)\n```\n\n## Stability\n\nThe initial architecture for EGNN suffered from instability when there was high number of neighbors. Thankfully, there seems to be two solutions that largely mitigate this.\n\n```python\nimport torch\nfrom egnn_pytorch import EGNN_Network\n\nnet = EGNN_Network(\n    num_tokens = 21,\n    dim = 32,\n    depth = 3,\n    num_nearest_neighbors = 32,\n    norm_coors = True,              # normalize the relative coordinates\n    coor_weights_clamp_value = 2.   # absolute clamped value for the coordinate weights, needed if you increase the num neareest neighbors\n)\n\nfeats = torch.randint(0, 21, (1, 1024)) # (1, 1024)\ncoors = torch.randn(1, 1024, 3)         # (1, 1024, 3)\nmask = torch.ones_like(feats).bool()    # (1, 1024)\n\nfeats_out, coors_out = net(feats, coors, mask = mask) # (1, 1024, 32), (1, 1024, 3)\n```\n\n## All parameters\n\n```python\nimport torch\nfrom egnn_pytorch import EGNN\n\nmodel = EGNN(\n    dim = dim,                         # input dimension\n    edge_dim = 0,                      # dimension of the edges, if exists, should be \u003e 0\n    m_dim = 16,                        # hidden model dimension\n    fourier_features = 0,              # number of fourier features for encoding of relative distance - defaults to none as in paper\n    num_nearest_neighbors = 0,         # cap the number of neighbors doing message passing by relative distance\n    dropout = 0.0,                     # dropout\n    norm_feats = False,                # whether to layernorm the features\n    norm_coors = False,                # whether to normalize the coordinates, using a strategy from the SE(3) Transformers paper    \n    update_feats = True,               # whether to update features - you can build a layer that only updates one or the other\n    update_coors = True,               # whether ot update coordinates\n    only_sparse_neighbors = False,     # using this would only allow message passing along adjacent neighbors, using the adjacency matrix passed in \n    valid_radius = float('inf'),       # the valid radius each node considers for message passing\n    m_pool_method = 'sum',             # whether to mean or sum pool for output node representation\n    soft_edges = False,                # extra GLU on the edges, purportedly helps stabilize the network in updated version of the paper\n    coor_weights_clamp_value = None    # clamping of the coordinate updates, again, for stabilization purposes\n)\n\n```\n\n## Examples\n\nTo run the protein backbone denoising example, first install `sidechainnet`\n\n```bash\n$ pip install sidechainnet\n```\n\nThen\n\n```bash\n$ python denoise_sparse.py\n```\n\n## Tests\n\nMake sure you have pytorch geometric installed locally\n\n```bash\n$ python setup.py test\n```\n\n## Citations\n\n```bibtex\n@misc{satorras2021en,\n    title \t= {E(n) Equivariant Graph Neural Networks}, \n    author \t= {Victor Garcia Satorras and Emiel Hoogeboom and Max Welling},\n    year \t= {2021},\n    eprint \t= {2102.09844},\n    archivePrefix = {arXiv},\n    primaryClass = {cs.LG}\n}\n```\n","project_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Flucidrains%2Fegnn-pytorch","html_url":"https://awesome.ecosyste.ms/projects/github.com%2Flucidrains%2Fegnn-pytorch","lists_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Flucidrains%2Fegnn-pytorch/lists"}