Skip to content
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

Jump to bottom

Issue gatv2 conv #3678

Merged
merged 6 commits into from Dec 11, 2021
Merged

Issue gatv2 conv #3678

Changes from all commits
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
cdd50a8
add the lazy initialization features so the class can use with to_hetero
richcmwang Dec 10, 2021
f085a6f
fixed length over 80
richcmwang Dec 10, 2021
aed252b
add lazy initialization
richcmwang Dec 10, 2021
8643055
add lazy clean file
richcmwang Dec 11, 2021
647d850
add lazy clean
richcmwang Dec 11, 2021
6000e57
update doc
rusty1s Dec 11, 2021
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Jump to
Jump to file
Failed to load files.
Diff view
Diff view
31 changes: 21 additions & 10 deletions torch_geometric/nn/conv/gatv2_conv.py
View file
Open in desktop
Expand Up @@ -53,8 +53,10 @@ class GATv2Conv(MessagePassing):
\right)\right)}.

Args:
in_channels (int): Size of each input sample, or :obj:`-1` to derive
the size from the first input(s) to the forward method.
in_channels (int or tuple): Size of each input sample, or :obj:`-1` to
derive the size from the first input(s) to the forward method.
A tuple corresponds to the sizes of source and target
dimensionalities.
out_channels (int): Size of each output sample.
heads (int, optional): Number of multi-head-attentions.
(default: :obj:`1`)
Expand Down Expand Up @@ -90,7 +92,7 @@ class GATv2Conv(MessagePassing):

def __init__(
self,
in_channels: int,
in_channels: Union[int, Tuple[int, int]],
out_channels: int,
heads: int = 1,
concat: bool = True,
Expand All @@ -103,7 +105,7 @@ def __init__(
share_weights: bool = False,
**kwargs,
):
super(GATv2Conv, self).__init__(node_dim=0, **kwargs)
super().__init__(node_dim=0, **kwargs)

self.in_channels = in_channels
self.out_channels = out_channels
Expand All @@ -116,13 +118,22 @@ def __init__(
self.fill_value = fill_value
self.share_weights = share_weights

self.lin_l = Linear(in_channels, heads * out_channels, bias=bias,
weight_initializer='glorot')
if share_weights:
self.lin_r = self.lin_l
else:
self.lin_r = Linear(in_channels, heads * out_channels, bias=bias,
if isinstance(in_channels, int):
self.lin_l = Linear(in_channels, heads * out_channels, bias=bias,
weight_initializer='glorot')
if share_weights:
self.lin_r = self.lin_l
else:
self.lin_r = Linear(in_channels, heads * out_channels,
bias=bias, weight_initializer='glorot')
else:
self.lin_l = Linear(in_channels[0], heads * out_channels,
bias=bias, weight_initializer='glorot')
if share_weights:
self.lin_r = self.lin_l
else:
self.lin_r = Linear(in_channels[1], heads * out_channels,
bias=bias, weight_initializer='glorot')

self.att = Parameter(torch.Tensor(1, heads, out_channels))

Expand Down