def test_mstgcn(): """ Testing MSTGCN block with changing edge index over time. """ 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 = MSTGCN(nb_block, node_features, K, nb_chev_filter, nb_time_filter, nb_time_strides, num_for_predict, len_input) 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_mstgcn(): """ Testing MSTGCN 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 = MSTGCN(DEVICE, nb_block, node_features, K, nb_chev_filter, nb_time_filter, nb_time_strides, adj_mx.toarray(), num_for_predict, len_input) 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)
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 = MSTGCN(DEVICE, nb_block, node_features, K, nb_chev_filter, nb_time_filter, nb_time_strides, adj_mx.toarray(),num_for_predict, len_input) 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 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)