if __name__ == '__main__':
start_time = time.time()
print('cuda available:', torch.cuda.is_available())
print("device:", args.device)
print("model:", args.model)
print("dataset:", args.dataset)
print("gcn type:", args.gcn_type)
torch.manual_seed(7)
if args.dataset == "metr":
A, X, means, stds = load_metr_la_data()
else:
A, X, means, stds = load_nyc_sharing_bike_data()
print('(num_nodes, num_features, num_time_steps) is ', X.shape)
split_line1 = int(X.shape[2] * 0.6)
split_line2 = int(X.shape[2] * 0.8)
t1 = 100
train_original_data = X[:, :, :split_line1]
val_original_data = X[:, :, split_line1 - t1:split_line2]
test_original_data = X[:, :, split_line2 - t1:]
if args.denosing == "deno":
training_input, training_target = get_denosing_dataset(
train_original_data, num_timesteps_input, num_timesteps_output)
val_input, val_target = get_denosing_dataset(val_original_data,
def init_data(self, data_dir=None):
if data_dir is None:
data_dir = self.config.data_dir
if self.config.dataset == "metr":
A, X, means, stds = load_metr_la_data(data_dir)
elif self.config.dataset == 'nyc-bike':
A, X, means, stds = load_nyc_sharing_bike_data(data_dir)
elif self.config.dataset == 'sf-example':
edge_index, edge_weight, X, shenzhenmask = load_sf_sample_data(
data_dir)
self.node_mask = shenzhenmask
X = X.astype(np.float32)
split_line1 = int(X.shape[2] * 0.2)
split_line2 = int(X.shape[2] * 0.4)
train_original_data = X[:, :, :split_line1]
val_original_data = X[:, :, split_line1:split_line2]
test_original_data = X[:, :, split_line2:]
self.training_input, self.training_target = generate_dataset(
train_original_data,
num_timesteps_input=self.config.num_timesteps_input,
num_timesteps_output=self.config.num_timesteps_output)
self.val_input, self.val_target = generate_dataset(
val_original_data,
num_timesteps_input=self.config.num_timesteps_input,
num_timesteps_output=self.config.num_timesteps_output)
self.test_input, self.test_target = generate_dataset(
test_original_data,
num_timesteps_input=self.config.num_timesteps_input,
num_timesteps_output=self.config.num_timesteps_output)
if self.config.dataset in ['metr', 'nyc-bike']:
self.A = torch.from_numpy(A)
self.sparse_A = self.A.to_sparse()
self.edge_index = self.sparse_A._indices()
self.edge_weight = self.sparse_A._values()
# set config attributes for model initialization
self.config.num_nodes = self.A.shape[0]
self.config.num_edges = self.edge_weight.shape[0]
self.config.num_edge_features = 1
self.config.num_features = self.training_input.shape[3]
self.log('Total nodes: {}'.format(self.config.num_nodes))
self.log('Average degree: {:.3f}'.format(self.config.num_edges /
self.config.num_nodes))
contains_self_loops = torch_geometric.utils.contains_self_loops(
self.edge_index)
self.log('Contains self loops: {}, but we add them.'.format(
contains_self_loops))
if not contains_self_loops:
self.edge_index, self.edge_weight = torch_geometric.utils.add_self_loops(
self.edge_index,
self.edge_weight,
num_nodes=self.config.num_nodes)
elif self.config.dataset in ['sf-example']:
self.edge_index = torch.from_numpy(edge_index)
self.edge_weight = torch.from_numpy(edge_weight)
# set config attributes for model initialization
self.config.num_nodes = len(shenzhenmask)
self.config.num_edges = self.edge_weight.shape[0]
self.config.num_edge_features = self.edge_weight.shape[1]
self.config.num_features = self.training_input.shape[3]
self.log('Total nodes: {}'.format(self.config.num_nodes))
self.log('Average degree: {:.3f}'.format(self.config.num_edges /
self.config.num_nodes))