def test_astgcn():
"""
Testing ASTGCN block with changing edge index over time or not
"""
node_count = 307
num_classes = 10
edge_per_node = 15
num_for_predict = 12
len_input = 12
nb_time_strides = 1
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
node_features = 2
nb_block = 2
K = 3
nb_chev_filter = 64
nb_time_filter = 64
batch_size = 32
x, edge_index = create_mock_data(node_count, edge_per_node, node_features)
model = ASTGCN(nb_block, node_features, K, nb_chev_filter, nb_time_filter,
nb_time_strides, num_for_predict, len_input,
node_count).to(device)
T = len_input
x_seq = torch.zeros([batch_size, node_count, node_features, T]).to(device)
target_seq = torch.zeros([batch_size, node_count, T]).to(device)
edge_index_seq = []
for b in range(batch_size):
for t in range(T):
x, edge_index = create_mock_data(node_count, edge_per_node,
node_features)
x_seq[b, :, :, t] = x
if b == 0:
edge_index_seq.append(edge_index)
target = create_mock_target(node_count, num_classes)
target_seq[b, :, t] = target
shuffle = True
train_dataset = torch.utils.data.TensorDataset(x_seq, target_seq)
train_loader = torch.utils.data.DataLoader(train_dataset,
batch_size=batch_size,
shuffle=shuffle)
for batch_data in train_loader:
encoder_inputs, _ = batch_data
outputs1 = model(encoder_inputs, edge_index_seq)
outputs2 = model(encoder_inputs, edge_index_seq[0])
assert outputs1.shape == (batch_size, node_count, num_for_predict)
assert outputs2.shape == (batch_size, node_count, num_for_predict)
def test_astgcn():
"""
Testing ASTGCN block
"""
node_count = 307
num_classes = 10
edge_per_node = 15
num_of_vertices = node_count # 307
num_for_predict = 12
len_input = 12
nb_time_strides = 1
DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
node_features = 1
nb_block = 2
K = 3
nb_chev_filter = 64
nb_time_filter = 64
batch_size = 32
x, edge_index = create_mock_data(node_count, edge_per_node, node_features)
adj_mx = to_scipy_sparse_matrix(edge_index)
model = ASTGCN(DEVICE, nb_block, node_features, K,
nb_chev_filter, nb_time_filter, nb_time_strides,
adj_mx.toarray(), num_for_predict, len_input, node_count)
T = len_input
x_seq = torch.zeros([batch_size, node_count, node_features, T]).to(DEVICE)
target_seq = torch.zeros([batch_size, node_count, T]).to(DEVICE)
for b in range(batch_size):
for t in range(T):
x, edge_index = create_mock_data(node_count, edge_per_node,
node_features)
x_seq[b, :, :, t] = x
target = create_mock_target(node_count, num_classes)
target_seq[b, :, t] = target
shuffle = True
train_dataset = torch.utils.data.TensorDataset(x_seq, target_seq)
train_loader = torch.utils.data.DataLoader(train_dataset,
batch_size=batch_size,
shuffle=shuffle)
criterion = torch.nn.MSELoss().to(DEVICE)
for batch_data in train_loader:
encoder_inputs, labels = batch_data
outputs = model(encoder_inputs)
assert outputs.shape == (batch_size, node_count, num_for_predict)
def test_astgcn():
"""
Testing ASTGCN block and its component ChebConvAttention with changing edge index over time or not
"""
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
in_channels, out_channels = (16, 32)
batch_size = 3
edge_index = torch.tensor([[0, 0, 0, 1, 2, 3], [1, 2, 3, 0, 0,
0]]).to(device)
num_nodes = edge_index.max().item() + 1
edge_weight = torch.rand(edge_index.size(1)).to(device)
x = torch.randn((batch_size, num_nodes, in_channels)).to(device)
attention = torch.nn.functional.softmax(torch.rand(
(batch_size, num_nodes, num_nodes)),
dim=1).to(device)
conv = ChebConvAttention(in_channels,
out_channels,
K=3,
normalization='sym').to(device)
assert conv.__repr__(
) == 'ChebConvAttention(16, 32, K=3, normalization=sym)'
out1 = conv(x, edge_index, attention)
assert out1.size() == (batch_size, num_nodes, out_channels)
out2 = conv(x, edge_index, attention, edge_weight)
assert out2.size() == (batch_size, num_nodes, out_channels)
out3 = conv(x, edge_index, attention, edge_weight, lambda_max=3.0)
assert out3.size() == (batch_size, num_nodes, out_channels)
batch = torch.tensor([0, 0, 1, 1]).to(device)
edge_index = torch.tensor([[0, 1, 2, 3], [1, 0, 3, 2]]).to(device)
num_nodes = edge_index.max().item() + 1
edge_weight = torch.rand(edge_index.size(1)).to(device)
x = torch.randn((batch_size, num_nodes, in_channels)).to(device)
lambda_max = torch.tensor([2.0, 3.0]).to(device)
attention = torch.nn.functional.softmax(torch.rand(
(batch_size, num_nodes, num_nodes)),
dim=1).to(device)
out4 = conv(x, edge_index, attention, edge_weight, batch)
assert out4.size() == (batch_size, num_nodes, out_channels)
out5 = conv(x, edge_index, attention, edge_weight, batch, lambda_max)
assert out5.size() == (batch_size, num_nodes, out_channels)
node_count = 307
num_classes = 10
edge_per_node = 15
num_for_predict = 12
len_input = 12
nb_time_strides = 1
node_features = 2
nb_block = 2
K = 3
nb_chev_filter = 64
nb_time_filter = 64
batch_size = 32
normalization = None
bias = True
model = ASTGCN(nb_block, node_features, K, nb_chev_filter, nb_time_filter,
nb_time_strides, num_for_predict, len_input, node_count,
normalization, bias).to(device)
model2 = ASTGCN(nb_block, node_features, K, nb_chev_filter, nb_time_filter,
nb_time_strides, num_for_predict, len_input, node_count,
'sym', False).to(device)
model3 = ASTGCN(nb_block, node_features, K, nb_chev_filter, nb_time_filter,
nb_time_strides, num_for_predict, len_input, node_count,
'rw', bias).to(device)
T = len_input
x_seq = torch.zeros([batch_size, node_count, node_features, T]).to(device)
target_seq = torch.zeros([batch_size, node_count, T]).to(device)
edge_index_seq = []
for b in range(batch_size):
for t in range(T):
x, edge_index = create_mock_data(node_count, edge_per_node,
node_features)
x_seq[b, :, :, t] = x.to(device)
if b == 0:
edge_index_seq.append(edge_index.to(device))
target = create_mock_target(node_count, num_classes).to(device)
target_seq[b, :, t] = target
shuffle = True
train_dataset = torch.utils.data.TensorDataset(x_seq, target_seq)
train_loader = torch.utils.data.DataLoader(train_dataset,
batch_size=batch_size,
shuffle=shuffle)
for batch_data in train_loader:
encoder_inputs, _ = batch_data
outputs0 = model(encoder_inputs, edge_index_seq)
outputs1 = model(encoder_inputs, edge_index_seq[0])
outputs2 = model2(encoder_inputs, edge_index_seq[0])
outputs3 = model2(encoder_inputs, edge_index_seq)
outputs4 = model3(encoder_inputs, edge_index_seq[0])
outputs5 = model3(encoder_inputs, edge_index_seq)
assert outputs0.shape == (batch_size, node_count, num_for_predict)
assert outputs1.shape == (batch_size, node_count, num_for_predict)
assert outputs2.shape == (batch_size, node_count, num_for_predict)
assert outputs3.shape == (batch_size, node_count, num_for_predict)
assert outputs4.shape == (batch_size, node_count, num_for_predict)
assert outputs5.shape == (batch_size, node_count, num_for_predict)
num_for_predict = 12
len_input = 12
nb_time_strides = 1
DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
node_features = 1
nb_block = 2
K = 3
nb_chev_filter = 64
nb_time_filter = 64
batch_size = 32
x, edge_index = create_mock_data(node_count, edge_per_node, node_features)
adj_mx = to_scipy_sparse_matrix(edge_index)
model = ASTGCN(DEVICE, nb_block, node_features,
K, nb_chev_filter, nb_time_filter, nb_time_strides,
adj_mx.toarray(), num_for_predict, len_input, node_count)
optimizer = torch.optim.Adam(model.parameters(),
lr=learning_rate,
weight_decay=weight_decay)
model.train()
T = len_input
x_seq = torch.zeros([batch_size, node_count, node_features, T]).to(DEVICE)
target_seq = torch.zeros([batch_size, node_count, T]).to(DEVICE)
for b in range(batch_size):
for t in range(T):
x, edge_index = create_mock_data(node_count, edge_per_node,
node_features)
x_seq[b, :, :, t] = x