Pytorch-基于GCN/GAT/Chebnet图神经网络实现的交通流预测(附代码)

代码地址
代码地址：https://github.com/LeronQ/GCN_predict-Pytorch

知乎主页

1:目录结构

基于图神经网络实现的交通流量预测，主要包括：GCN、GAR、ChebNet算法。

2:数据集信息

数据来自美国的加利福尼亚州的洛杉矶市，CSV文件是关于节点的表示情况，一共有307个节点，npz文件是交通流量的文件，每5分钟输出节点数据信息。
数据集信息：

• PEMS-04
• 数据时间范围：2018.1.1—2018.2.28
• 节点信息：共307个检测器，即Nodes数量为307
• 3个特征：流量值，占有率，速度值。

3:数据探索

dataView.py 文件为数据可视化探索分析，可直接运行，查看3个特征的数据分布。
可以看到每个节点有三个特征，但有两个节点基本是平稳不变的，所以只取第一维特征（即流量特征）。

4:数据处理

traffic_dataset.py文件为数据处理模块，主要包含：

• 读取邻接矩阵函数：get_adjacent_matrix
• 读取数据集函数：get_flow_data
• 加载数据类：class LoadData

其中，LoadData 主要包含数据读取，数据切分和数据归一化处理等。
可单独运行该模块查看结果。

class LoadData(Dataset):  # 这个就是把读入的数据处理成模型需要的训练数据和测试数据，一个一个样本能读取出来
    def __init__(self, data_path, num_nodes, divide_days, time_interval, history_length, train_mode):
        """
        :param data_path: list, ["graph file name" , "flow data file name"], path to save the data file names.
        :param num_nodes: int, number of nodes.
        :param divide_days: list, [ days of train data, days of test data], list to divide the original data.
        :param time_interval: int, time interval between two traffic data records (mins).---5 mins
        :param history_length: int, length of history data to be used.
        :param train_mode: list, ["train", "test"].
        """

        self.data_path = data_path
        self.num_nodes = num_nodes
        self.train_mode = train_mode
        self.train_days = divide_days[0]  # 59-14 = 45, train_data
        self.test_days = divide_days[1]  # 7*2 = 14 ,test_data
        self.history_length = history_length  # 30/5 = 6, 历史长度为6
        self.time_interval = time_interval  # 5 min

        self.one_day_length = int(24 * 60 / self.time_interval) # 一整天的数据量

        self.graph = get_adjacent_matrix(distance_file=data_path[0], num_nodes=num_nodes)

        self.flow_norm, self.flow_data = self.pre_process_data(data=get_flow_data(data_path[1]), norm_dim=1) # self.flow_norm为归一化的基

    def __len__(self):  # 表示数据集的长度
        """
        :return: length of dataset (number of samples).
        """
        if self.train_mode == "train":
            return self.train_days * self.one_day_length - self.history_length #　训练的样本数　＝　训练集总长度　－　历史数据长度
        elif self.train_mode == "test":
            return self.test_days * self.one_day_length  #　每个样本都能测试，测试样本数　＝　测试总长度
        else:
            raise ValueError("train mode: [{}] is not defined".format(self.train_mode))

    def __getitem__(self, index):  # 功能是如何取每一个样本 (x, y), index = [0, L1 - 1]这个是根据数据集的长度确定的
        """
        :param index: int, range between [0, length - 1].
        :return:
            graph: torch.tensor, [N, N].
            data_x: torch.tensor, [N, H, D].
            data_y: torch.tensor, [N, 1, D].
        """
        if self.train_mode == "train":
            index = index#训练集的数据是从时间０开始的，这个是每一个流量数据，要和样本（ｘ,y）区别
        elif self.train_mode == "test":
            index += self.train_days * self.one_day_length#有一个偏移量
        else:
            raise ValueError("train mode: [{}] is not defined".format(self.train_mode))

        data_x, data_y = LoadData.slice_data(self.flow_data, self.history_length, index, self.train_mode)#这个就是样本（ｘ,y）

        data_x = LoadData.to_tensor(data_x)  # [N, H, D] # 转换成张量
        data_y = LoadData.to_tensor(data_y).unsqueeze(1)