{ "cells": [ { "cell_type": "markdown", "id": "8d7675a4", "metadata": {}, "source": [ "### Divide training and testing sets" ] }, { "cell_type": "code", "execution_count": 23, "id": "9cee47ff", "metadata": {}, "outputs": [], "source": [ "import numpy as np" ] }, { "cell_type": "code", "execution_count": 24, "id": "adf0f0be", "metadata": {}, "outputs": [], "source": [ "def preprocess_data(data, time_len, rate, seq_len, pre_len): \n", " train_size = int(time_len * rate)\n", " train_data = data[0:train_size]\n", " test_data = data[train_size:time_len]\n", "\n", " trainX, trainY, testX, testY = [], [], [], []\n", " for i in range(len(train_data) - seq_len - pre_len):\n", " a = train_data[i: i + seq_len + pre_len]\n", " trainX.append(a[0: seq_len])\n", " trainY.append(a[seq_len: seq_len + pre_len])\n", "\n", " for i in range(len(test_data) - seq_len - pre_len):\n", " b = test_data[i: i + seq_len + pre_len]\n", " testX.append(b[0: seq_len])\n", " testY.append(b[seq_len: seq_len + pre_len])\n", "\n", " trainX1 = np.array(trainX)\n", " trainY1 = np.array(trainY)\n", " testX1 = np.array(testX)\n", " testY1 = np.array(testY)\n", " return trainX1, trainY1, testX1, testY1" ] }, { "cell_type": "markdown", "id": "36b6ca60", "metadata": {}, "source": [ "### Define TC-GAT traffic flow prediction models" ] }, { "cell_type": "code", "execution_count": 25, "id": "b678e3a7", "metadata": {}, "outputs": [], "source": [ "import torch\n", "import torch.nn as nn\n", "from torch.nn.utils import weight_norm\n", "import torch.nn.functional as F\n", "import numpy as np\n", "\n", "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n", "from torch_geometric.nn import GATConv" ] }, { "cell_type": "markdown", "id": "be5b2c57", "metadata": {}, "source": [ "#### Attention" ] }, { "cell_type": "code", "execution_count": 26, "id": "578ba35c", "metadata": {}, "outputs": [], "source": [ "class Attention(nn.Module): # Attention模块\n", " \"\"\"\n", " Scaled dot-product attention\n", " \"\"\"\n", "\n", " def __init__(self, d_model, d_k, d_v, h, dropout=.1):\n", " \"\"\"\n", " :param d_model: Output dimensionality of the model\n", " :param d_k: Dimensionality of queries and keys\n", " :param d_v: Dimensionality of values\n", " :param h: Number of heads 多头注意力机制\n", " \"\"\"\n", " super(Attention, self).__init__()\n", " self.fc_q = nn.Linear(d_model, h * d_k)\n", " self.fc_k = nn.Linear(d_model, h * d_k)\n", " self.fc_v = nn.Linear(d_model, h * d_v)\n", " self.fc_o = nn.Linear(h * d_v, d_model)\n", " self.dropout = nn.Dropout(dropout)\n", "\n", " self.d_model = d_model\n", " self.d_k = d_k\n", " self.d_v = d_v\n", " self.h = h\n", "\n", " def forward(self, queries, keys, values, attention_mask=None, attention_weights=None):\n", "\n", " b_s, nq = queries.shape[:2]\n", " nk = keys.shape[1]\n", " q = self.fc_q(queries).view(b_s, nq, self.h, self.d_k).permute(0, 2, 1, 3) # (b_s, h, nq, d_k)\n", " k = self.fc_k(keys).view(b_s, nk, self.h, self.d_k).permute(0, 2, 3, 1) # (b_s, h, d_k, nk)\n", " v = self.fc_v(values).view(b_s, nk, self.h, self.d_v).permute(0, 2, 1, 3) # (b_s, h, nk, d_v)\n", " att = torch.matmul(q, k) / np.sqrt(self.d_k) # (b_s, h, nq, nk) 自注意力机制的实现\n", " if attention_weights is not None:\n", " att = att * attention_weights\n", " if attention_mask is not None:\n", " att = att.masked_fill(attention_mask, -np.inf)\n", " att = torch.softmax(att, -1)\n", " att = self.dropout(att)\n", "\n", " out = torch.matmul(att, v).permute(0, 2, 1, 3).contiguous().view(b_s, nq, self.h * self.d_v) # (b_s, nq, h*d_v)\n", " out = self.fc_o(out) # (b_s, nq, d_model)\n", " return out" ] }, { "cell_type": "markdown", "id": "52139a7c", "metadata": {}, "source": [ "#### Edge construction for GAT" ] }, { "cell_type": "code", "execution_count": 27, "id": "1eede2db", "metadata": {}, "outputs": [], "source": [ "def edge_index_func(matrix_path):\n", " # print(\"In edge index function\")\n", " a, b = [], []\n", " matrix = matrix_path\n", "\n", " for i in range(matrix.shape[0]):\n", " for j in range(matrix.shape[1]):\n", " # if(matrix[i][j] == 1 and i != j):\n", " if matrix[i][j] == 1:\n", " a.append(i)\n", " b.append(j)\n", " edge = [a, b]\n", " edge_index = torch.tensor(edge, dtype=torch.long)\n", " return edge_index" ] }, { "cell_type": "markdown", "id": "dd92c663", "metadata": {}, "source": [ "#### GAT" ] }, { "cell_type": "code", "execution_count": 28, "id": "ab38f23d", "metadata": {}, "outputs": [], "source": [ "class gat(nn.Module):\n", " def __init__(self, in_channels, out_channels, adj, embeddim, heads=0, dropout=0):\n", " super(gat, self).__init__()\n", " self.in_channels = in_channels\n", " self.out_channels = out_channels\n", " self.head = heads\n", " self.drop = dropout\n", " self.edge_index = edge_index_func(adj).to(device)\n", " self.gat1 = GATConv(in_channels * embeddim, out_channels * embeddim, heads=heads, dropout=dropout, concat=False)\n", " # self.gat2 = GATConv(in_channels, out_channels, heads=heads, dropout=dropout,concat = False)\n", " self.conv1 = nn.Conv2d(embeddim, embeddim, 1, 1)\n", " self.conv2 = nn.Conv2d(embeddim, embeddim, 1, 1)\n", "\n", " def forward(self, x):\n", " # B H T N\n", " x = x.permute(0, 3, 2, 1)\n", " x = self.conv1(x)\n", " x = x.permute(0, 3, 1, 2) # 把时间放在最后一个维度用于输入到TCN\n", " B, N, H, T = x.shape\n", " res = x # 做一个残差\n", " x = x.reshape(B, N, T * H)\n", " X_G = torch.Tensor(0, N, T * H).to(device)\n", "\n", " for b in range(B):\n", " temp = self.gat1(x[b], self.edge_index)\n", " X_G = torch.cat((X_G, temp.unsqueeze(0)), dim=0)\n", " x = X_G\n", " x = x.reshape(B, N, H, T)\n", " x = (x + res).permute(0, 2, 3, 1) # 做一个残差\n", " x = self.conv2(x).permute(0, 3, 2, 1)\n", " return x" ] }, { "cell_type": "markdown", "id": "12317d61", "metadata": {}, "source": [ "#### TCN" ] }, { "cell_type": "code", "execution_count": 29, "id": "3b8d813c", "metadata": {}, "outputs": [], "source": [ "# 剪枝,一维卷积后会出现多余的padding\n", "class Chomp1d(nn.Module):\n", " def __init__(self, chomp_size):\n", " super(Chomp1d, self).__init__()\n", " self.chomp_size = chomp_size\n", "\n", " def forward(self, x):\n", " return x[:, :, :-self.chomp_size].contiguous()" ] }, { "cell_type": "code", "execution_count": 30, "id": "ef37c38a", "metadata": {}, "outputs": [], "source": [ "# 时序模块,两层一维卷积，两层Weight_Norm,两层Chomd1d，非线性激活函数为Relu,dropout为0.2。\n", "class TemporalBlock(nn.Module):\n", " def __init__(self, n_inputs, n_outputs, kernel_size, stride, dilation, padding, dropout=0.2):\n", " \"\"\"\n", " 相当于一个Residual block\n", "\n", " :param n_inputs: int, 输入通道数\n", " :param n_outputs: int, 输出通道数\n", " :param kernel_size: int, 卷积核尺寸\n", " :param stride: int, 步长，一般为1\n", " :param dilation: int, 膨胀系数\n", " :param padding: int, 填充系数\n", " :param dropout: float, dropout比率\n", " \"\"\"\n", " super(TemporalBlock, self).__init__()\n", " self.conv1 = weight_norm(nn.Conv1d(n_inputs, n_outputs, kernel_size,\n", " stride=stride, padding=padding, dilation=dilation))\n", " # 经过conv1，输出的size其实是(Batch, input_channel, seq_len + padding)\n", " self.chomp1 = Chomp1d(padding) # 裁剪掉多出来的padding部分，维持输出时间步为seq_len\n", " self.relu1 = nn.ReLU()\n", " self.dropout1 = nn.Dropout(dropout)\n", "\n", " self.conv2 = weight_norm(nn.Conv1d(n_outputs, n_outputs, kernel_size,\n", " stride=stride, padding=padding, dilation=dilation))\n", " self.chomp2 = Chomp1d(padding) # 裁剪掉多出来的padding部分，维持输出时间步为seq_len\n", " self.relu2 = nn.ReLU()\n", " self.dropout2 = nn.Dropout(dropout)\n", "\n", " self.net = nn.Sequential(self.conv1, self.chomp1, self.relu1, self.dropout1,\n", " self.conv2, self.chomp2, self.relu2, self.dropout2)\n", " self.downsample = nn.Conv1d(n_inputs, n_outputs, 1) if n_inputs != n_outputs else None\n", " self.relu = nn.ReLU()\n", " self.init_weights()\n", "\n", " def init_weights(self):\n", " self.conv1.weight.data.normal_(0, 0.01)\n", " self.conv2.weight.data.normal_(0, 0.01)\n", " if self.downsample is not None:\n", " self.downsample.weight.data.normal_(0, 0.01)\n", "\n", " def forward(self, x):\n", " \"\"\"\n", " :param x: size of (Batch, input_channel, seq_len)\n", " :return:\n", " \"\"\"\n", " out = self.net(x)\n", " res = x if self.downsample is None else self.downsample(x)\n", " return self.relu(out + res) # 残差块" ] }, { "cell_type": "code", "execution_count": 31, "id": "2180e87b", "metadata": {}, "outputs": [], "source": [ "# 时序卷积模块,其中*layer不是c中的指针，* 表示迭代器拆分layers为一层层网络。\n", "class TemporalConvNet(nn.Module):\n", " def __init__(self, num_inputs, num_channels, kernel_size=2, dropout=0.2):\n", " \"\"\"\n", " TCN，目前paper给出的TCN结构很好的支持每个时刻为一个数的情况，即sequence结构，\n", " 对于每个时刻为一个向量这种一维结构，勉强可以把向量拆成若干该时刻的输入通道，\n", " 对于每个时刻为一个矩阵或更高维图像的情况，就不太好办。\n", "\n", " :param num_inputs: int， 输入通道数\n", " :param num_channels: list，每层的hidden_channel数，例如[25,25,25,25]表示有4个隐层，每层hidden_channel数为25\n", " :param kernel_size: int, 卷积核尺寸\n", " :param dropout: float, drop_out比率\n", " \"\"\"\n", " super(TemporalConvNet, self).__init__()\n", " layers = []\n", " num_levels = len(num_channels)\n", " for i in range(num_levels):\n", " dilation_size = 2 ** i # 膨胀系数：1，2，4，8……\n", " in_channels = num_inputs if i == 0 else num_channels[i - 1] # 确定每一层的输入通道数,输入层通道为1，隐含层是25。\n", " out_channels = num_channels[i] # 确定每一层的输出通道数\n", " layers += [TemporalBlock(in_channels, out_channels, kernel_size, stride=1, dilation=dilation_size,\n", " padding=(kernel_size - 1) * dilation_size, dropout=dropout)]\n", "\n", " self.network = nn.Sequential(*layers) # *作用是将输入迭代器拆成一个个元素\n", "\n", " def forward(self, x):\n", " \"\"\"\n", " 输入x的结构不同于RNN，一般RNN的size为(Batch, seq_len, channels)或者(seq_len, Batch, channels)，\n", " 这里把seq_len放在channels后面，把所有时间步的数据拼起来，当做Conv1d的输入尺寸，实现卷积跨时间步的操作，\n", " 很巧妙的设计。\n", "\n", " :param x: size of (Batch, input_channel, seq_len)\n", " :return: size of (Batch, output_channel, seq_len)\n", " \"\"\"\n", " return self.network(x)" ] }, { "cell_type": "markdown", "id": "d3b1be35", "metadata": {}, "source": [ "#### TC-GAT" ] }, { "cell_type": "code", "execution_count": 32, "id": "3bbb4cfb", "metadata": {}, "outputs": [], "source": [ "class GATNet(nn.Module):\n", " def __init__(self, num_nodes, In_len, out_len, embedsize, head, adj):\n", " super(GATNet, self).__init__()\n", " self.conv1 = nn.Conv2d(1, embedsize, 1, 1) # 输入维度为1 输出维度大小为64\n", " self.tcn1 = TemporalConvNet(num_inputs=num_nodes, num_channels=[num_nodes, num_nodes], kernel_size=5,\n", " dropout=0.1)\n", " self.gat = gat(In_len, In_len, adj, embeddim=embedsize, heads=head, dropout=0)\n", " self.Att = Attention(d_model=embedsize * In_len, d_k=embedsize * 2, d_v=embedsize * 2, h=head, dropout=0)\n", " self.feed_forward = nn.Sequential(\n", " nn.Linear(embedsize * In_len, 4 * embedsize * In_len),\n", " nn.ReLU(),\n", " nn.Linear(4 * embedsize * In_len, embedsize * In_len),\n", " )\n", " self.conv2 = nn.Sequential(nn.Conv2d(In_len, out_len, 1, 1))\n", " self.conv3 = nn.Sequential(nn.Conv2d(embedsize, embedsize, 1, 1), nn.ReLU(), nn.Conv2d(embedsize, 1, 1, 1))\n", " self.norm1 = nn.LayerNorm(embedsize)\n", " self.norm2 = nn.LayerNorm(embedsize * In_len)\n", " self.norm3 = nn.LayerNorm(embedsize * In_len)\n", "\n", " def forward(self, x):\n", " # x [B,T,N]\n", " # B,N,T,H\n", " x = x.unsqueeze(1) # 达到适应2维卷积的维度 二维卷积的输入为[batch_size, channel_num, H, W]\n", " x = self.conv1(x).permute(0, 3, 2, 1)\n", " res = x\n", " B, N, T, H = x.shape\n", "\n", " ################################GAT\n", " x = self.norm1(self.gat(x) + res)\n", " ################################TCN+ATT\n", "\n", " x = x.reshape(B, N, -1)\n", " res1 = x\n", " x = self.tcn1(x)# 可以多次叠加TCN\n", " x = self.norm2(self.Att(x, x, x) + res1) # 此处参考Transformer用两个全链接层做输出\n", " res2 = x\n", " x = self.norm3(self.feed_forward(x) + res2)\n", " x = x.reshape(B, N, T, H)\n", "\n", " ##################输出层\n", " x = x.permute(0, 2, 1, 3)\n", " x = F.relu(x)\n", " x = self.conv2(x)\n", " x = x.permute(0, 3, 1, 2)\n", " x = F.relu(x)\n", " x = self.conv3(x).squeeze(1) # 恢复维度,1D·FCN\n", " return x\n" ] }, { "cell_type": "markdown", "id": "d1eb782b", "metadata": {}, "source": [ "### Main函数" ] }, { "cell_type": "code", "execution_count": 33, "id": "b65962fe", "metadata": {}, "outputs": [], "source": [ "import warnings\n", "\n", "warnings.filterwarnings(\"ignore\")\n", "import torch.nn as nn\n", "from torch.utils.data import DataLoader, TensorDataset\n", "import pandas as pd\n", "import torch" ] }, { "cell_type": "code", "execution_count": 34, "id": "5db5f848", "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "2976\n", "156\n" ] } ], "source": [ "data = np.array(pd.read_csv(r'data/sz_speed.csv'))\n", "print(data.shape[0])\n", "print(data.shape[1])" ] }, { "cell_type": "code", "execution_count": null, "id": "b0c206f2", "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "(2366, 12, 156) (2366, 2, 156) (582, 12, 156) (582, 2, 156)\n", "rmse,mae: 0 7.171055 5.296642\n", "rmse,mae: 1 5.6994705 4.2154827\n" ] } ], "source": [ "# 评价指标 平均绝对误差(MAE) RMSE(均方根误差)\n", "def metric(pred, label):\n", " with np.errstate(divide='ignore', invalid='ignore'):\n", " mask = np.not_equal(label, 0)\n", " mask = mask.astype(np.float32)\n", " mask /= np.mean(mask)\n", " mae = np.abs(np.subtract(pred, label)).astype(np.float32)\n", " rmse = np.square(mae)\n", " mae = np.nan_to_num(mae * mask)\n", " mae = np.mean(mae)\n", " rmse = np.nan_to_num(rmse * mask)\n", " rmse = np.sqrt(np.mean(rmse))\n", " return mae, rmse\n", "\n", "\n", "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n", "data = np.array(pd.read_csv(r'data/sz_speed.csv'))\n", "adj = np.array(pd.read_csv(r'data/sz_adj.csv', header=None))\n", "time_len = data.shape[0]\n", "num_nodes = data.shape[1]\n", "train_rate = 0.8\n", "seq_len = 12\n", "pre_len = 2\n", "trainX, trainY, testX, testY = preprocess_data(data, time_len, train_rate, seq_len, pre_len)\n", "print(trainX.shape, trainY.shape, testX.shape, testY.shape)\n", "# 归一化\n", "##########################\n", "mean, std = np.mean(trainX), np.std(trainX)\n", "trainX = (trainX - mean) / std\n", "testX = (testX - mean) / std\n", "##########################\n", "\n", "\n", "trainX = torch.tensor(trainX, dtype=torch.float32)\n", "trainY = torch.tensor(trainY, dtype=torch.float32)\n", "testX = torch.tensor(testX, dtype=torch.float32)\n", "testY = torch.tensor(testY, dtype=torch.float32)\n", "adj = torch.tensor(np.array(adj), dtype=torch.float32)\n", "\n", "train_data = TensorDataset(trainX, trainY)\n", "train_dataloader = DataLoader(train_data, batch_size=128, shuffle=True)\n", "\n", "test_data = TensorDataset(testX, testY)\n", "test_dataloader = DataLoader(test_data, batch_size=128)\n", "\n", "model = GATNet(num_nodes=num_nodes, In_len=seq_len, out_len=pre_len, embedsize=64, head=4, adj=adj)\n", "model = model.to(device)\n", "criterion = nn.MSELoss()\n", "\n", "optimizer = torch.optim.Adam(model.parameters(), lr=0.0001)\n", "best_mae = 100\n", "best_rmse = None\n", "best_mape = None\n", "\n", "for epoch in range(100):\n", " # 训练\n", " model.train()\n", " for x, y in train_dataloader:\n", " x = x.to(device)\n", " y = y.to(device)\n", " pre = model(x).reshape(-1, adj.shape[0])\n", " y = y.reshape(-1, adj.shape[0])\n", "\n", " loss = criterion(pre * std + mean, y)\n", " optimizer.zero_grad()\n", " loss.backward()\n", " optimizer.step()\n", "\n", " model.eval()\n", " # 测试\n", " with torch.no_grad():\n", " P = []\n", " L = []\n", " for x, y in test_dataloader:\n", " x = x.to(device)\n", " pre = model(x) * std + mean\n", " P.append(pre.cpu().detach())\n", " L.append(y)\n", " pre = torch.cat(P, 0)\n", " label = torch.cat(L, 0)\n", "\n", " pre = pre.reshape(-1, adj.shape[0])\n", " label = label.reshape(-1, adj.shape[0])\n", " mae, rmse = metric(pre.numpy(), label.numpy())\n", " print(\"rmse,mae:\", epoch, rmse, mae)" ] } ], "metadata": { "kernelspec": { "display_name": "Python 3 (ipykernel)", "language": "python", "name": "python3" }, "language_info": { "codemirror_mode": { "name": "ipython", "version": 3 }, "file_extension": ".py", "mimetype": "text/x-python", "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", "version": "3.9.7" } }, "nbformat": 4, "nbformat_minor": 5 }