Ejemplo n.º 1
0

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,
Ejemplo n.º 2
0
    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))