Skip to content

Added optional tanh to coors_mlp #33

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 1 commit into
base: main
Choose a base branch
from
Open

Added optional tanh to coors_mlp #33

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
9 changes: 7 additions & 2 deletions egnn_pytorch/egnn_pytorch.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -164,11 +164,13 @@ def __init__(
valid_radius = float('inf'),
m_pool_method = 'sum',
soft_edges = False,
coor_weights_clamp_value = None
coor_weights_clamp_value = None,
coors_tanh=False # Only to be used alongside norm_coors - highly recommended for stability
):
super().__init__()
assert m_pool_method in {'sum', 'mean'}, 'pool method must be either sum or mean'
assert update_feats or update_coors, 'you must update either features, coordinates, or both'
assert not (coors_tanh and not norm_coors), 'coors_tanh must be used with norm_coors'

self.fourier_features = fourier_features

Expand Down Expand Up @@ -200,11 +202,14 @@ def __init__(
nn.Linear(dim * 2, dim),
) if update_feats else None

# Tanh layer helps with stability but should only be used in conjuction with
# norm_coors
self.coors_mlp = nn.Sequential(
nn.Linear(m_dim, m_dim * 4),
dropout,
SiLU(),
nn.Linear(m_dim * 4, 1)
nn.Linear(m_dim * 4, 1),
nn.Tanh() if coors_tanh else nn.Identity()
) if update_coors else None

self.num_nearest_neighbors = num_nearest_neighbors
Expand Down