def __init__(self, in_dim, z_dim, edge_index, edge_weight, h_dim=256):
"""
Args:
in_dim (int): dimensionality of the input
z_dim (int): dimensionality of the low-dimensional space
edge_index (LongTensor): shape (2, ?), edge indices
edge_weight (FloatTensor): shape (?), edge weights.
h_dim (int): dimensionality of intermediate layers in MLP
"""
super(Encoder, self).__init__()
self.edge_index = edge_index
self.edge_weight = edge_weight
self.fn = nn.Sequential(
nn.Linear(in_dim, h_dim, bias=True),
nn.GELU(),
nn.Linear(h_dim, h_dim, bias=True),
nn.GELU(),
)
self.gc = Sequential(
"x, edge_index, edge_weight",
[
(GCNConv(h_dim, z_dim, cached=False, add_self_loops=True),
"x, edge_index, edge_weight -> x"),
# nn.BatchNorm1d(z_dim),
nn.GELU(),
(GCNConv(z_dim, z_dim, cached=False, add_self_loops=True),
"x, edge_index, edge_weight -> x"),
# nn.BatchNorm1d(z_dim),
nn.GELU(),
nn.Linear(z_dim, z_dim)
])
self.gen = nn.Sequential(nn.Linear(z_dim, z_dim, bias=True))
def test_sequential_tracable():
model = Sequential('x, edge_index', [
(GCNConv(16, 64), 'x, edge_index -> x1'),
ReLU(inplace=True),
(GCNConv(64, 64), 'x1, edge_index -> x2'),
ReLU(inplace=True),
(lambda x1, x2: x1 + x2, 'x1, x2 -> x'),
Linear(64, 7),
])
symbolic_trace(model)
def __init__(self, num_features, hidden_channels_list, num_classes):
super(GAT, self).__init__()
torch.manual_seed(12345)
hns = [num_features] + hidden_channels_list
conv_list = []
for idx in range(len(hidden_channels_list)):
conv_list.append((GATConv(hns[idx], hns[idx+1]), 'x, edge_index -> x'))
conv_list.append(ReLU(inplace=True),)
self.convseq = Sequential('x, edge_index', conv_list)
self.linear = Linear(hidden_channels_list[-1], num_classes)
def test_sequential_jittable():
x = torch.randn(4, 16)
edge_index = torch.tensor([[0, 0, 0, 1, 2, 3], [1, 2, 3, 0, 0, 0]])
adj_t = SparseTensor(row=edge_index[0], col=edge_index[1]).t()
model = Sequential('x: Tensor, edge_index: Tensor', [
(GCNConv(16, 64).jittable(), 'x, edge_index -> x'),
ReLU(inplace=True),
(GCNConv(64, 64).jittable(), 'x, edge_index -> x'),
ReLU(inplace=True),
Linear(64, 7),
])
torch.jit.script(model)(x, edge_index)
model = Sequential('x: Tensor, edge_index: SparseTensor', [
(GCNConv(16, 64).jittable(), 'x, edge_index -> x'),
ReLU(inplace=True),
(GCNConv(64, 64).jittable(), 'x, edge_index -> x'),
ReLU(inplace=True),
Linear(64, 7),
])
torch.jit.script(model)(x, adj_t)
def test_sequential_with_multiple_return_values():
x = torch.randn(4, 16)
edge_index = torch.tensor([[0, 0, 0, 1, 2, 3], [1, 2, 3, 0, 0, 0]])
model = Sequential('x, edge_index', [
(GCNConv(16, 32), 'x, edge_index -> x1'),
(GCNConv(32, 64), 'x1, edge_index -> x2'),
(lambda x1, x2: (x1, x2), 'x1, x2 -> x1, x2'),
])
x1, x2 = model(x, edge_index)
assert x1.size() == (4, 32)
assert x2.size() == (4, 64)
def test_sequential():
x = torch.randn(4, 16)
edge_index = torch.tensor([[0, 0, 0, 1, 2, 3], [1, 2, 3, 0, 0, 0]])
batch = torch.zeros(4, dtype=torch.long)
model = Sequential('x, edge_index', [
(GCNConv(16, 64), 'x, edge_index -> x'),
ReLU(inplace=True),
(GCNConv(64, 64), 'x, edge_index -> x'),
ReLU(inplace=True),
Linear(64, 7),
])
model.reset_parameters()
assert len(model) == 5
assert str(model) == (
'Sequential(\n'
' (0): GCNConv(16, 64)\n'
' (1): ReLU(inplace=True)\n'
' (2): GCNConv(64, 64)\n'
' (3): ReLU(inplace=True)\n'
' (4): Linear(in_features=64, out_features=7, bias=True)\n'
')')
assert isinstance(model[0], GCNConv)
assert isinstance(model[1], ReLU)
assert isinstance(model[2], GCNConv)
assert isinstance(model[3], ReLU)
assert isinstance(model[4], Linear)
out = model(x, edge_index)
assert out.size() == (4, 7)
model = Sequential('x, edge_index, batch', [
(Dropout(p=0.5), 'x -> x'),
(GCNConv(16, 64), 'x, edge_index -> x1'),
ReLU(inplace=True),
(GCNConv(64, 64), 'x1, edge_index -> x2'),
ReLU(inplace=True),
(lambda x1, x2: [x1, x2], 'x1, x2 -> xs'),
(JumpingKnowledge('cat', 64, num_layers=2), 'xs -> x'),
(global_mean_pool, 'x, batch -> x'),
Linear(2 * 64, 7),
])
model.reset_parameters()
out = model(x, edge_index, batch)
assert out.size() == (1, 7)
def test_sequential_with_ordered_dict():
x = torch.randn(4, 16)
edge_index = torch.tensor([[0, 0, 0, 1, 2, 3], [1, 2, 3, 0, 0, 0]])
model = Sequential(
'x, edge_index', modules=OrderedDict([
('conv1', (GCNConv(16, 32), 'x, edge_index -> x')),
('conv2', (GCNConv(32, 64), 'x, edge_index -> x')),
]))
assert isinstance(model.conv1, GCNConv)
assert isinstance(model.conv2, GCNConv)
x = model(x, edge_index)
assert x.size() == (4, 64)
def __init__(
self,
embed_dim: int,
num_layers: int,
heads: int = 8,
normalization: str = "batch",
feed_forward_hidden: int = 512,
pooling_method: str = "mean",
) -> None:
super().__init__()
self.embed_dim = embed_dim
self.num_layers = num_layers
self.heads = heads
assert (self.embed_dim % self.heads) == 0
self.pooling_func = get_pooling_func(pooling_method)
self.norm_class = get_normalization_class(normalization)
self.gnn_layer_list = nn.ModuleList()
self.norm_list = nn.ModuleList()
for i in range(self.num_layers):
gnn_layer = TransformerConv(
in_channels=self.embed_dim,
out_channels=self.embed_dim // self.heads,
heads=self.heads,
edge_dim=self.embed_dim,
)
self.gnn_layer_list.append(gnn_layer)
self.norm_list.append(GraphNorm(in_channels=self.embed_dim))
self.feed_forward = Sequential(
"x, batch",
[
(nn.Linear(self.embed_dim, feed_forward_hidden), "x -> x"),
nn.GELU(),
# (GraphNorm(in_channels=feed_forward_hidden), "x, batch -> x"),
nn.Linear(feed_forward_hidden, self.embed_dim),
],
)
self.ff_norm = GraphNorm(in_channels=self.embed_dim)
**kwargs)
else:
train_loader = HGTLoader(data,
num_samples=[1024] * 4,
shuffle=True,
input_nodes=train_input_nodes,
**kwargs)
val_loader = HGTLoader(data,
num_samples=[1024] * 4,
input_nodes=val_input_nodes,
**kwargs)
model = Sequential('x, edge_index', [
(SAGEConv((-1, -1), 64), 'x, edge_index -> x'),
ReLU(inplace=True),
(SAGEConv((-1, -1), 64), 'x, edge_index -> x'),
ReLU(inplace=True),
(Linear(-1, dataset.num_classes), 'x -> x'),
])
model = to_hetero(model, data.metadata(), aggr='sum').to(device)
@torch.no_grad()
def init_params():
# Initialize lazy parameters via forwarding a single batch to the model:
batch = next(iter(train_loader))
batch = batch.to(device, 'edge_index')
model(batch.x_dict, batch.edge_index_dict)
def train():
def __init__(self, input_feat=2, Conv_outputs=[5], LSTM_output=[5],
K=1, linear_output=3):
super(social_stgcn, self).__init__()
self.input_feat = input_feat
self.Conv_outputs = Conv_outputs
self.LSTM_output = LSTM_output
self.linear_output = linear_output
self.K = K
self.gcn1 = GCNConv(in_channels=self.input_feat,
out_channels=self.Conv_outputs[0],
improved=True)
self.gcn2 = GCNConv(in_channels=self.Conv_outputs[0],
out_channels=self.Conv_outputs[1],
improved=True)
self.gclstm1 = GConvLSTM(in_channels=self.Conv_outputs[1],
out_channels=self.Conv_outputs[1],
K=K, normalization="sym", bias=True)
self.gclstm2 = GConvLSTM(in_channels=self.Conv_outputs[1],
out_channels=self.Conv_outputs[1],
K=K, normalization="sym", bias=True)
self.gclstm3 = GConvLSTM(in_channels=self.Conv_outputs[1],
out_channels=self.LSTM_output[0],
K=K, normalization="sym", bias=True)
self.no_lstm = 3
self.linear = nn.Linear(in_features=self.LSTM_output[0],
out_features=self.linear_output)
self.gcn = Sequential('x, edge_index, h_none', [
(GCNConv(in_channels=self.input_feat,
out_channels=self.Conv_outputs[0],
improved=True), 'x, edge_index -> x'),
(nn.ReLU(), 'x -> x'),
(GCNConv(in_channels=self.Conv_outputs[0],
out_channels=self.Conv_outputs[1],
improved=True), 'x, edge_index -> x'),
(nn.ReLU(), 'x -> x'),
(GConvLSTM(in_channels=self.Conv_outputs[1],
out_channels=self.Conv_outputs[1],
K=K, normalization="sym", bias=True), 'x, edge_index -> h, c'),
(GConvLSTM(in_channels=self.Conv_outputs[1],
out_channels=self.Conv_outputs[1],
K=K, normalization="sym", bias=True), 'h, edge_index, h_none, c -> h, c'),
(GConvLSTM(in_channels=self.Conv_outputs[1],
out_channels=self.LSTM_output[0],
K=K, normalization="sym", bias=True), 'h, edge_index, h_none, c-> h, c'),
# TO BE TESTED LATER
# (GCLSTM(in_channels=self.Conv_outputs[1],
# out_channels=self.LSTM_output[0],
# K=K, normalization="sym", bias=True), 'x, edge_index -> h, c'),
#
# (GCLSTM(in_channels=self.LSTM_output[0],
# out_channels=self.LSTM_output[1],
# K=K, normalization="sym", bias=True), 'x, edge_index, h, c -> h, c'),
#
# (GCLSTM(in_channels=self.LSTM_output[1],
# out_channels=self.LSTM_output[2],
# K=K, normalization="sym", bias=True), 'h, edge_index, c -> h, _'),
(nn.ReLU(), "h -> h"),
(nn.Linear(in_features=self.LSTM_output[0],
out_features=self.linear_output), "h -> x")])