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))