diff --git a/.gitignore b/.gitignore index de47f30..efca7c6 100644 --- a/.gitignore +++ b/.gitignore @@ -134,3 +134,6 @@ FETCH_HEAD # vscode # .vscode + +# experiments +experiments/* diff --git a/.vscode/launch.json b/.vscode/launch.json index 93462f8..859a5ab 100644 --- a/.vscode/launch.json +++ b/.vscode/launch.json @@ -70,5 +70,19 @@ "PYTHONPATH": "${workspaceFolder}", } }, + { + "name": "example_CorrSTN_ddeint", + "type": "python", + "request": "launch", + "program": "${workspaceFolder}/example/CorrSTN/train_dde.py", + "console": "integratedTerminal", + "justMyCode": false, + "args": [], + "env": { + "CUDA_VISIBLE_DEVICES": "0", + "PYTHONPATH": "${workspaceFolder}", + } + }, + ] } \ No newline at end of file diff --git a/example/CorrSTN/TrafficFlowData b/example/CorrSTN/TrafficFlowData index e4fd17a..f8a2d29 160000 --- a/example/CorrSTN/TrafficFlowData +++ b/example/CorrSTN/TrafficFlowData @@ -1 +1 @@ -Subproject commit e4fd17ae1b463eab790b71d3a0dbaec698d5747e +Subproject commit f8a2d293a8d24b8a683b129b9bffb7c89be71e6e diff --git a/example/CorrSTN/args.py b/example/CorrSTN/args.py index 76838e8..487c460 100644 --- a/example/CorrSTN/args.py +++ b/example/CorrSTN/args.py @@ -31,7 +31,7 @@ # model config -parser.add_argument("--model_name", type=str, default="WindSTN", help="model name") +parser.add_argument("--model_name", type=str, default="DDECorrSTN", help="model name") parser.add_argument("--his_len", type=int, default=288, help="history data length") parser.add_argument("--tgt_len", type=int, default=12, help="tgt data length") parser.add_argument("--encoder_input_size", type=int, default=1) @@ -39,27 +39,31 @@ parser.add_argument("--decoder_output_size", type=int, default=1) parser.add_argument("--encoder_num_layers", type=int, default=4) parser.add_argument("--decoder_num_layers", type=int, default=4) -parser.add_argument("--d_model", type=int, default=64) +parser.add_argument("--d_model", type=int, default=64, help="d_proj+d_sect*2") +parser.add_argument("--d_proj", type=int, default=32) +parser.add_argument("--d_sect", type=int, default=16) parser.add_argument("--attention", type=str, default="Corr", help="Corr,Vanilla") parser.add_argument("--split_seq", type=bool, default=False, help="split q k v") parser.add_argument("--head", type=int, default=8, help="head") parser.add_argument("--kernel_size", type=int, default=3, help="kernel_size") parser.add_argument("--top_k", type=int, default=5, help="top_k") parser.add_argument("--smooth_layer_num", type=int, default=1) - +parser.add_argument("--no_adj", type=bool, default=False, help="no adj") # train config parser.add_argument("--learning_rate", type=float, default=1e-3) -parser.add_argument("--weight_decay", type=float, default=0.01) +parser.add_argument("--weight_decay", type=float, default=0.0) parser.add_argument("--start_epoch", type=int, default=0, help="start epoch") -parser.add_argument("--train_epochs", type=int, default=75, help="train epochs") +parser.add_argument("--train_epochs", type=int, default=100, help="train epochs") parser.add_argument("--finetune_epochs", type=int, default=50, help="finetune epochs") parser.add_argument("--batch_size", type=int, default=16, help="batch_size") -parser.add_argument("--patience", type=int, default=8, help="early stopping patience") +parser.add_argument("--patience", type=int, default=15, help="early stopping patience") parser.add_argument("--loss", type=str, default="mse", help="loss function") parser.add_argument("--dropout", type=float, default=0.0, help="dropout") parser.add_argument("--continue_training", type=bool, default=False, help="") parser.add_argument("--fp16", type=bool, default=False, help="") +parser.add_argument("--distribute", type=bool, default=False, help="") -args = parser.parse_args("") +# args = parser.parse_args("") +args = parser.parse_args() os.environ["CUDA_VISIBLE_DEVICES"] = args.devices diff --git a/example/CorrSTN/attention.py b/example/CorrSTN/attention.py index a870839..4c25de8 100644 --- a/example/CorrSTN/attention.py +++ b/example/CorrSTN/attention.py @@ -125,6 +125,7 @@ def __init__( args.d_model, (1, args.kernel_size), padding=(0, self.padding_1DConv), + bias_attr=True, ) else: self.query_conv = nn.Conv2D( @@ -132,6 +133,7 @@ def __init__( args.d_model, (1, args.kernel_size), padding=(0, self.padding_causal), + bias_attr=True, ) if key_conv_type == "1DConv": @@ -140,6 +142,7 @@ def __init__( args.d_model, (1, args.kernel_size), padding=(0, self.padding_1DConv), + bias_attr=True, ) else: self.key_conv = nn.Conv2D( @@ -147,6 +150,7 @@ def __init__( args.d_model, (1, args.kernel_size), padding=(0, self.padding_causal), + bias_attr=True, ) self.dropout = nn.Dropout(p=args.dropout) diff --git a/example/CorrSTN/corrstn.ipynb b/example/CorrSTN/corrstn.ipynb deleted file mode 100644 index 34c205b..0000000 --- a/example/CorrSTN/corrstn.ipynb +++ /dev/null @@ -1,150 +0,0 @@ -{ - "cells": [ - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "import numpy as np\n", - "import paddle\n", - "import paddle.nn as nn\n", - "from args import args\n", - "from utils import norm_adj_matrix\n", - "\n", - "corr = np.load(\"data/HZME_OUTFLOW/SCORR_HZME_OUTFLOW_0.6_C16.npy\")[0]\n", - "adj_matrix = paddle.to_tensor(norm_adj_matrix(corr), paddle.get_default_dtype())" - ] - }, - { - "cell_type": "code", - "execution_count": 2, - "metadata": {}, - "outputs": [], - "source": [ - "from dataset import TrafficFlowDataset\n", - "from paddle.io import DataLoader\n", - "\n", - "train_dataset = TrafficFlowDataset(args, \"train\")\n", - "train_dataloader = DataLoader(\n", - " train_dataset, batch_size=args.batch_size, shuffle=False\n", - ")" - ] - }, - { - "cell_type": "code", - "execution_count": 3, - "metadata": {}, - "outputs": [], - "source": [ - "his, tgt = next(iter(train_dataloader))" - ] - }, - { - "cell_type": "code", - "execution_count": 4, - "metadata": {}, - "outputs": [], - "source": [ - "his_index = paddle.randint(\n", - " shape=[args.batch_size * args.num_nodes, 36], high=his.shape[-2], low=0\n", - ")\n", - "axis_bs = paddle.arange(args.batch_size * args.num_nodes)[:, None, None]\n", - "axis_index = his_index[:, :, None]\n", - "axis_dim = paddle.arange(1)[None, None, :]\n", - "his = his.reshape([-1] + his.shape[-2:])\n", - "his_input = his[axis_bs, axis_index, axis_dim]" - ] - }, - { - "cell_type": "code", - "execution_count": 5, - "metadata": {}, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "Attention Type: Vanilla\n", - "Attention Type: Vanilla\n", - "Attention Type: Vanilla\n" - ] - }, - { - "name": "stderr", - "output_type": "stream", - "text": [ - "Traceback (most recent call last):\n", - " File \"_pydevd_bundle/pydevd_cython.pyx\", line 577, in _pydevd_bundle.pydevd_cython.PyDBFrame._handle_exception\n", - " File \"_pydevd_bundle/pydevd_cython.pyx\", line 312, in _pydevd_bundle.pydevd_cython.PyDBFrame.do_wait_suspend\n", - " File \"/Users/drownfish19/anaconda3/envs/paddle-2.5.1/lib/python3.10/site-packages/debugpy/_vendored/pydevd/pydevd.py\", line 2070, in do_wait_suspend\n", - " keep_suspended = self._do_wait_suspend(thread, frame, event, arg, suspend_type, from_this_thread, frames_tracker)\n", - " File \"/Users/drownfish19/anaconda3/envs/paddle-2.5.1/lib/python3.10/site-packages/debugpy/_vendored/pydevd/pydevd.py\", line 2106, in _do_wait_suspend\n", - " time.sleep(0.01)\n", - "KeyboardInterrupt\n" - ] - }, - { - "ename": "AttributeError", - "evalue": "'Tensor' object has no attribute 'permute'", - "output_type": "error", - "traceback": [ - "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m", - "\u001b[0;31mAttributeError\u001b[0m Traceback (most recent call last)", - "Cell \u001b[0;32mIn[5], line 7\u001b[0m\n\u001b[1;32m 3\u001b[0m nn\u001b[39m.\u001b[39minitializer\u001b[39m.\u001b[39mset_global_initializer(\n\u001b[1;32m 4\u001b[0m nn\u001b[39m.\u001b[39minitializer\u001b[39m.\u001b[39mXavierUniform(), nn\u001b[39m.\u001b[39minitializer\u001b[39m.\u001b[39mXavierUniform()\n\u001b[1;32m 5\u001b[0m )\n\u001b[1;32m 6\u001b[0m model \u001b[39m=\u001b[39m CorrSTN(args, adj_matrix)\n\u001b[0;32m----> 7\u001b[0m model(his_input, his_index, tgt)\n", - "File \u001b[0;32m~/anaconda3/envs/paddle-2.5.1/lib/python3.10/site-packages/paddle/nn/layer/layers.py:1254\u001b[0m, in \u001b[0;36mLayer.__call__\u001b[0;34m(self, *inputs, **kwargs)\u001b[0m\n\u001b[1;32m 1245\u001b[0m \u001b[39mif\u001b[39;00m (\n\u001b[1;32m 1246\u001b[0m (\u001b[39mnot\u001b[39;00m in_declarative_mode())\n\u001b[1;32m 1247\u001b[0m \u001b[39mand\u001b[39;00m (\u001b[39mnot\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_forward_pre_hooks)\n\u001b[0;32m (...)\u001b[0m\n\u001b[1;32m 1251\u001b[0m \u001b[39mand\u001b[39;00m (\u001b[39mnot\u001b[39;00m in_profiler_mode())\n\u001b[1;32m 1252\u001b[0m ):\n\u001b[1;32m 1253\u001b[0m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_build_once(\u001b[39m*\u001b[39minputs, \u001b[39m*\u001b[39m\u001b[39m*\u001b[39mkwargs)\n\u001b[0;32m-> 1254\u001b[0m \u001b[39mreturn\u001b[39;00m \u001b[39mself\u001b[39;49m\u001b[39m.\u001b[39;49mforward(\u001b[39m*\u001b[39;49minputs, \u001b[39m*\u001b[39;49m\u001b[39m*\u001b[39;49mkwargs)\n\u001b[1;32m 1255\u001b[0m \u001b[39melse\u001b[39;00m:\n\u001b[1;32m 1256\u001b[0m \u001b[39mreturn\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_dygraph_call_func(\u001b[39m*\u001b[39minputs, \u001b[39m*\u001b[39m\u001b[39m*\u001b[39mkwargs)\n", - "File \u001b[0;32m~/Documents/GitHub/ODE/PaddleXDE/example/CorrSTN/corrstn.py:102\u001b[0m, in \u001b[0;36mCorrSTN.forward\u001b[0;34m(self, src, src_idx, tgt)\u001b[0m\n\u001b[1;32m 98\u001b[0m src \u001b[39m=\u001b[39m src\u001b[39m.\u001b[39mreshape([\u001b[39m-\u001b[39m\u001b[39m1\u001b[39m, \u001b[39mself\u001b[39m\u001b[39m.\u001b[39mtraining_args\u001b[39m.\u001b[39mnum_nodes] \u001b[39m+\u001b[39m src\u001b[39m.\u001b[39mshape[\u001b[39m-\u001b[39m\u001b[39m2\u001b[39m:])\n\u001b[1;32m 99\u001b[0m src_idx \u001b[39m=\u001b[39m src_idx\u001b[39m.\u001b[39mreshape(\n\u001b[1;32m 100\u001b[0m [\u001b[39m-\u001b[39m\u001b[39m1\u001b[39m, \u001b[39mself\u001b[39m\u001b[39m.\u001b[39mtraining_args\u001b[39m.\u001b[39mnum_nodes] \u001b[39m+\u001b[39m src_idx\u001b[39m.\u001b[39mshape[\u001b[39m-\u001b[39m\u001b[39m1\u001b[39m:]\n\u001b[1;32m 101\u001b[0m )\n\u001b[0;32m--> 102\u001b[0m encoder_output \u001b[39m=\u001b[39m \u001b[39mself\u001b[39;49m\u001b[39m.\u001b[39;49mencode(src, src_idx)\n\u001b[1;32m 103\u001b[0m output \u001b[39m=\u001b[39m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39mdecode(tgt, encoder_output)\n\u001b[1;32m 104\u001b[0m \u001b[39mreturn\u001b[39;00m output\n", - "File \u001b[0;32m~/Documents/GitHub/ODE/PaddleXDE/example/CorrSTN/corrstn.py:87\u001b[0m, in \u001b[0;36mCorrSTN.encode\u001b[0;34m(self, src, lookup_index)\u001b[0m\n\u001b[1;32m 85\u001b[0m \u001b[39m# src_tp_embedding = self.encode_temporal_position(src_dense, lookup_index)\u001b[39;00m\n\u001b[1;32m 86\u001b[0m src_sp_embedding \u001b[39m=\u001b[39m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39mencode_spatial_position(src_dense)\n\u001b[0;32m---> 87\u001b[0m encoder_output \u001b[39m=\u001b[39m \u001b[39mself\u001b[39;49m\u001b[39m.\u001b[39;49mencoder(src_sp_embedding)\n\u001b[1;32m 88\u001b[0m \u001b[39mreturn\u001b[39;00m encoder_output\n", - "File \u001b[0;32m~/anaconda3/envs/paddle-2.5.1/lib/python3.10/site-packages/paddle/nn/layer/layers.py:1254\u001b[0m, in \u001b[0;36mLayer.__call__\u001b[0;34m(self, *inputs, **kwargs)\u001b[0m\n\u001b[1;32m 1245\u001b[0m \u001b[39mif\u001b[39;00m (\n\u001b[1;32m 1246\u001b[0m (\u001b[39mnot\u001b[39;00m in_declarative_mode())\n\u001b[1;32m 1247\u001b[0m \u001b[39mand\u001b[39;00m (\u001b[39mnot\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_forward_pre_hooks)\n\u001b[0;32m (...)\u001b[0m\n\u001b[1;32m 1251\u001b[0m \u001b[39mand\u001b[39;00m (\u001b[39mnot\u001b[39;00m in_profiler_mode())\n\u001b[1;32m 1252\u001b[0m ):\n\u001b[1;32m 1253\u001b[0m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_build_once(\u001b[39m*\u001b[39minputs, \u001b[39m*\u001b[39m\u001b[39m*\u001b[39mkwargs)\n\u001b[0;32m-> 1254\u001b[0m \u001b[39mreturn\u001b[39;00m \u001b[39mself\u001b[39;49m\u001b[39m.\u001b[39;49mforward(\u001b[39m*\u001b[39;49minputs, \u001b[39m*\u001b[39;49m\u001b[39m*\u001b[39;49mkwargs)\n\u001b[1;32m 1255\u001b[0m \u001b[39melse\u001b[39;00m:\n\u001b[1;32m 1256\u001b[0m \u001b[39mreturn\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_dygraph_call_func(\u001b[39m*\u001b[39minputs, \u001b[39m*\u001b[39m\u001b[39m*\u001b[39mkwargs)\n", - "File \u001b[0;32m~/Documents/GitHub/ODE/PaddleXDE/example/CorrSTN/endecoder.py:87\u001b[0m, in \u001b[0;36mEncoder.forward\u001b[0;34m(self, x)\u001b[0m\n\u001b[1;32m 82\u001b[0m \u001b[39m\u001b[39m\u001b[39m\"\"\"\u001b[39;00m\n\u001b[1;32m 83\u001b[0m \u001b[39m:param x: src: (batch_size, N, T_in, F_in)\u001b[39;00m\n\u001b[1;32m 84\u001b[0m \u001b[39m:return: (batch_size, N, T_in, F_in)\u001b[39;00m\n\u001b[1;32m 85\u001b[0m \u001b[39m\"\"\"\u001b[39;00m\n\u001b[1;32m 86\u001b[0m \u001b[39mfor\u001b[39;00m layer \u001b[39min\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39mlayers:\n\u001b[0;32m---> 87\u001b[0m x \u001b[39m=\u001b[39m layer(x)\n\u001b[1;32m 88\u001b[0m \u001b[39mreturn\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39mnorm(x)\n", - "File \u001b[0;32m~/anaconda3/envs/paddle-2.5.1/lib/python3.10/site-packages/paddle/nn/layer/layers.py:1254\u001b[0m, in \u001b[0;36mLayer.__call__\u001b[0;34m(self, *inputs, **kwargs)\u001b[0m\n\u001b[1;32m 1245\u001b[0m \u001b[39mif\u001b[39;00m (\n\u001b[1;32m 1246\u001b[0m (\u001b[39mnot\u001b[39;00m in_declarative_mode())\n\u001b[1;32m 1247\u001b[0m \u001b[39mand\u001b[39;00m (\u001b[39mnot\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_forward_pre_hooks)\n\u001b[0;32m (...)\u001b[0m\n\u001b[1;32m 1251\u001b[0m \u001b[39mand\u001b[39;00m (\u001b[39mnot\u001b[39;00m in_profiler_mode())\n\u001b[1;32m 1252\u001b[0m ):\n\u001b[1;32m 1253\u001b[0m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_build_once(\u001b[39m*\u001b[39minputs, \u001b[39m*\u001b[39m\u001b[39m*\u001b[39mkwargs)\n\u001b[0;32m-> 1254\u001b[0m \u001b[39mreturn\u001b[39;00m \u001b[39mself\u001b[39;49m\u001b[39m.\u001b[39;49mforward(\u001b[39m*\u001b[39;49minputs, \u001b[39m*\u001b[39;49m\u001b[39m*\u001b[39;49mkwargs)\n\u001b[1;32m 1255\u001b[0m \u001b[39melse\u001b[39;00m:\n\u001b[1;32m 1256\u001b[0m \u001b[39mreturn\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_dygraph_call_func(\u001b[39m*\u001b[39minputs, \u001b[39m*\u001b[39m\u001b[39m*\u001b[39mkwargs)\n", - "File \u001b[0;32m~/Documents/GitHub/ODE/PaddleXDE/example/CorrSTN/endecoder.py:64\u001b[0m, in \u001b[0;36mEncoderLayer.forward\u001b[0;34m(self, x)\u001b[0m\n\u001b[1;32m 59\u001b[0m \u001b[39mif\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39mresidual_connection \u001b[39mor\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39muse_LayerNorm:\n\u001b[1;32m 60\u001b[0m x \u001b[39m=\u001b[39m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39msublayer[\u001b[39m0\u001b[39m](\n\u001b[1;32m 61\u001b[0m x,\n\u001b[1;32m 62\u001b[0m \u001b[39mlambda\u001b[39;00m x: \u001b[39mself\u001b[39m\u001b[39m.\u001b[39mself_attn(x, x, x),\n\u001b[1;32m 63\u001b[0m ) \u001b[39m# [B,N,T,D]\u001b[39;00m\n\u001b[0;32m---> 64\u001b[0m \u001b[39mreturn\u001b[39;00m \u001b[39mself\u001b[39;49m\u001b[39m.\u001b[39;49msublayer[\u001b[39m1\u001b[39;49m](x, \u001b[39mself\u001b[39;49m\u001b[39m.\u001b[39;49mfeed_forward_gcn) \u001b[39m# [B,N,T,D]\u001b[39;00m\n\u001b[1;32m 65\u001b[0m \u001b[39melse\u001b[39;00m:\n\u001b[1;32m 67\u001b[0m x \u001b[39m=\u001b[39m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39mself_attn(x, x, x) \u001b[39m# [B,N,T,D]\u001b[39;00m\n", - "File \u001b[0;32m~/anaconda3/envs/paddle-2.5.1/lib/python3.10/site-packages/paddle/nn/layer/layers.py:1254\u001b[0m, in \u001b[0;36mLayer.__call__\u001b[0;34m(self, *inputs, **kwargs)\u001b[0m\n\u001b[1;32m 1245\u001b[0m \u001b[39mif\u001b[39;00m (\n\u001b[1;32m 1246\u001b[0m (\u001b[39mnot\u001b[39;00m in_declarative_mode())\n\u001b[1;32m 1247\u001b[0m \u001b[39mand\u001b[39;00m (\u001b[39mnot\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_forward_pre_hooks)\n\u001b[0;32m (...)\u001b[0m\n\u001b[1;32m 1251\u001b[0m \u001b[39mand\u001b[39;00m (\u001b[39mnot\u001b[39;00m in_profiler_mode())\n\u001b[1;32m 1252\u001b[0m ):\n\u001b[1;32m 1253\u001b[0m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_build_once(\u001b[39m*\u001b[39minputs, \u001b[39m*\u001b[39m\u001b[39m*\u001b[39mkwargs)\n\u001b[0;32m-> 1254\u001b[0m \u001b[39mreturn\u001b[39;00m \u001b[39mself\u001b[39;49m\u001b[39m.\u001b[39;49mforward(\u001b[39m*\u001b[39;49minputs, \u001b[39m*\u001b[39;49m\u001b[39m*\u001b[39;49mkwargs)\n\u001b[1;32m 1255\u001b[0m \u001b[39melse\u001b[39;00m:\n\u001b[1;32m 1256\u001b[0m \u001b[39mreturn\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_dygraph_call_func(\u001b[39m*\u001b[39minputs, \u001b[39m*\u001b[39m\u001b[39m*\u001b[39mkwargs)\n", - "File \u001b[0;32m~/Documents/GitHub/ODE/PaddleXDE/example/CorrSTN/endecoder.py:25\u001b[0m, in \u001b[0;36mSublayerConnection.forward\u001b[0;34m(self, x, sublayer)\u001b[0m\n\u001b[1;32m 19\u001b[0m \u001b[39m\u001b[39m\u001b[39m\"\"\"\u001b[39;00m\n\u001b[1;32m 20\u001b[0m \u001b[39m:param x: (batch, N, T, d_model)\u001b[39;00m\n\u001b[1;32m 21\u001b[0m \u001b[39m:param sublayer: nn.Layer\u001b[39;00m\n\u001b[1;32m 22\u001b[0m \u001b[39m:return: (batch, N, T, d_model)\u001b[39;00m\n\u001b[1;32m 23\u001b[0m \u001b[39m\"\"\"\u001b[39;00m\n\u001b[1;32m 24\u001b[0m \u001b[39mif\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39mresidual_connection \u001b[39mand\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39muse_layer_norm:\n\u001b[0;32m---> 25\u001b[0m \u001b[39mreturn\u001b[39;00m x \u001b[39m+\u001b[39m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39mdropout(sublayer(\u001b[39mself\u001b[39;49m\u001b[39m.\u001b[39;49mnorm(x)))\n\u001b[1;32m 26\u001b[0m \u001b[39mif\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39mresidual_connection \u001b[39mand\u001b[39;00m (\u001b[39mnot\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39muse_layer_norm):\n\u001b[1;32m 27\u001b[0m \u001b[39mreturn\u001b[39;00m x \u001b[39m+\u001b[39m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39mdropout(sublayer(x))\n", - "File \u001b[0;32m~/anaconda3/envs/paddle-2.5.1/lib/python3.10/site-packages/paddle/nn/layer/layers.py:1254\u001b[0m, in \u001b[0;36mLayer.__call__\u001b[0;34m(self, *inputs, **kwargs)\u001b[0m\n\u001b[1;32m 1245\u001b[0m \u001b[39mif\u001b[39;00m (\n\u001b[1;32m 1246\u001b[0m (\u001b[39mnot\u001b[39;00m in_declarative_mode())\n\u001b[1;32m 1247\u001b[0m \u001b[39mand\u001b[39;00m (\u001b[39mnot\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_forward_pre_hooks)\n\u001b[0;32m (...)\u001b[0m\n\u001b[1;32m 1251\u001b[0m \u001b[39mand\u001b[39;00m (\u001b[39mnot\u001b[39;00m in_profiler_mode())\n\u001b[1;32m 1252\u001b[0m ):\n\u001b[1;32m 1253\u001b[0m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_build_once(\u001b[39m*\u001b[39minputs, \u001b[39m*\u001b[39m\u001b[39m*\u001b[39mkwargs)\n\u001b[0;32m-> 1254\u001b[0m \u001b[39mreturn\u001b[39;00m \u001b[39mself\u001b[39;49m\u001b[39m.\u001b[39;49mforward(\u001b[39m*\u001b[39;49minputs, \u001b[39m*\u001b[39;49m\u001b[39m*\u001b[39;49mkwargs)\n\u001b[1;32m 1255\u001b[0m \u001b[39melse\u001b[39;00m:\n\u001b[1;32m 1256\u001b[0m \u001b[39mreturn\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_dygraph_call_func(\u001b[39m*\u001b[39minputs, \u001b[39m*\u001b[39m\u001b[39m*\u001b[39mkwargs)\n", - "File \u001b[0;32m~/Documents/GitHub/ODE/PaddleXDE/example/CorrSTN/graphconv.py:66\u001b[0m, in \u001b[0;36mforward\u001b[0;34m(self, x)\u001b[0m\n\u001b[1;32m 64\u001b[0m B, N, T, D \u001b[39m=\u001b[39m x\u001b[39m.\u001b[39mshape\n\u001b[1;32m 65\u001b[0m spatial_attention \u001b[39m=\u001b[39m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39mSAt(x) \u001b[39m# [B, T, N, N]\u001b[39;00m\n\u001b[0;32m---> 66\u001b[0m \u001b[39mif\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39mis_scale:\n\u001b[1;32m 67\u001b[0m spatial_attention \u001b[39m=\u001b[39m spatial_attention \u001b[39m/\u001b[39m math\u001b[39m.\u001b[39msqrt(\u001b[39mself\u001b[39m\u001b[39m.\u001b[39margs\u001b[39m.\u001b[39md_model)\n\u001b[1;32m 68\u001b[0m x \u001b[39m=\u001b[39m x\u001b[39m.\u001b[39mpermute(\u001b[39m0\u001b[39m, \u001b[39m2\u001b[39m, \u001b[39m1\u001b[39m, \u001b[39m3\u001b[39m)\u001b[39m.\u001b[39mreshape((\u001b[39m-\u001b[39m\u001b[39m1\u001b[39m, N, D)) \u001b[39m# [B*T,N,D]\u001b[39;00m\n", - "File \u001b[0;32m~/anaconda3/envs/paddle-2.5.1/lib/python3.10/site-packages/paddle/nn/layer/layers.py:1254\u001b[0m, in \u001b[0;36mLayer.__call__\u001b[0;34m(self, *inputs, **kwargs)\u001b[0m\n\u001b[1;32m 1245\u001b[0m \u001b[39mif\u001b[39;00m (\n\u001b[1;32m 1246\u001b[0m (\u001b[39mnot\u001b[39;00m in_declarative_mode())\n\u001b[1;32m 1247\u001b[0m \u001b[39mand\u001b[39;00m (\u001b[39mnot\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_forward_pre_hooks)\n\u001b[0;32m (...)\u001b[0m\n\u001b[1;32m 1251\u001b[0m \u001b[39mand\u001b[39;00m (\u001b[39mnot\u001b[39;00m in_profiler_mode())\n\u001b[1;32m 1252\u001b[0m ):\n\u001b[1;32m 1253\u001b[0m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_build_once(\u001b[39m*\u001b[39minputs, \u001b[39m*\u001b[39m\u001b[39m*\u001b[39mkwargs)\n\u001b[0;32m-> 1254\u001b[0m \u001b[39mreturn\u001b[39;00m \u001b[39mself\u001b[39;49m\u001b[39m.\u001b[39;49mforward(\u001b[39m*\u001b[39;49minputs, \u001b[39m*\u001b[39;49m\u001b[39m*\u001b[39;49mkwargs)\n\u001b[1;32m 1255\u001b[0m \u001b[39melse\u001b[39;00m:\n\u001b[1;32m 1256\u001b[0m \u001b[39mreturn\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_dygraph_call_func(\u001b[39m*\u001b[39minputs, \u001b[39m*\u001b[39m\u001b[39m*\u001b[39mkwargs)\n", - "File \u001b[0;32m~/Documents/GitHub/ODE/PaddleXDE/example/CorrSTN/graphconv.py:42\u001b[0m, in \u001b[0;36mSpatialAttentionLayer.forward\u001b[0;34m(self, x)\u001b[0m\n\u001b[1;32m 37\u001b[0m \u001b[39m\u001b[39m\u001b[39m\"\"\"\u001b[39;00m\n\u001b[1;32m 38\u001b[0m \u001b[39m:param x: (B, N, T, D)\u001b[39;00m\n\u001b[1;32m 39\u001b[0m \u001b[39m:return: (B, T, N, N)\u001b[39;00m\n\u001b[1;32m 40\u001b[0m \u001b[39m\"\"\"\u001b[39;00m\n\u001b[1;32m 41\u001b[0m B, N, T, D \u001b[39m=\u001b[39m x\u001b[39m.\u001b[39mshape\n\u001b[0;32m---> 42\u001b[0m x \u001b[39m=\u001b[39m x\u001b[39m.\u001b[39mtranposr([\u001b[39m0\u001b[39m, \u001b[39m2\u001b[39m, \u001b[39m1\u001b[39m, \u001b[39m3\u001b[39m])\u001b[39m.\u001b[39mreshape((\u001b[39m-\u001b[39m\u001b[39m1\u001b[39m, N, D)) \u001b[39m# [B*T,N,F_in]\u001b[39;00m\n\u001b[1;32m 43\u001b[0m \u001b[39m# [B*T,N,F_in][B*T,F_in,N]=[B*T,N,N]\u001b[39;00m\n\u001b[1;32m 44\u001b[0m score \u001b[39m=\u001b[39m paddle\u001b[39m.\u001b[39mmatmul(x, x, transpose_y\u001b[39m=\u001b[39m\u001b[39mTrue\u001b[39;00m) \u001b[39m/\u001b[39m math\u001b[39m.\u001b[39msqrt(D)\n", - "\u001b[0;31mAttributeError\u001b[0m: 'Tensor' object has no attribute 'permute'" - ] - } - ], - "source": [ - "from corrstn import CorrSTN\n", - "\n", - "nn.initializer.set_global_initializer(\n", - " nn.initializer.XavierUniform(), nn.initializer.XavierUniform()\n", - ")\n", - "model = CorrSTN(args, adj_matrix)\n", - "model(his_input, his_index, tgt)" - ] - } - ], - "metadata": { - "kernelspec": { - "display_name": "paddle-2.5.1", - "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.10.12" - }, - "orig_nbformat": 4, - "vscode": { - "interpreter": { - "hash": "d5e07ec9d8a63f92ad92c0319a0bf17841e333a8063fcdf0df6e5ed0358e95fb" - } - } - }, - "nbformat": 4, - "nbformat_minor": 2 -} diff --git a/example/CorrSTN/corrstn.py b/example/CorrSTN/corrstn.py index d4dcbde..e516ae3 100644 --- a/example/CorrSTN/corrstn.py +++ b/example/CorrSTN/corrstn.py @@ -1,10 +1,21 @@ from copy import deepcopy -import paddle.nn as nn +import paddle from attention import MultiHeadAttentionAwareTemporalContext -from embedding import SpatialPositionalEmbedding, TemporalPositionalEmbedding +from embedding import TemporalSectionEmbedding + +# SpatialPositionalEmbedding,; TemporalPositionalEmbedding, from endecoder import Decoder, DecoderLayer, Encoder, EncoderLayer -from graphconv import GCN, SpatialAttentionGCN +from graphconv import SpatialAttentionGCN + +# GCN +from paddle import autograd, nn + +from paddlexde.interpolation.interpolate import ( + BezierSpline, + CubicHermiteSpline, + LinearInterpolation, +) class CorrSTN(nn.Layer): @@ -14,23 +25,25 @@ def __init__(self, training_args, adj_matrix, sc_matrix): self.training_args = training_args self.encoder_dense = nn.Linear( - training_args.encoder_input_size, training_args.d_model + training_args.encoder_input_size, training_args.d_proj ) self.decoder_dense = nn.Linear( - training_args.decoder_input_size, training_args.d_model - ) - - self.encode_temporal_position = TemporalPositionalEmbedding( - training_args, max_len=training_args.his_len - ) - self.decode_temporal_position = TemporalPositionalEmbedding( - training_args, max_len=training_args.tgt_len + training_args.decoder_input_size, training_args.d_proj ) - self.encode_spatial_position = SpatialPositionalEmbedding( - training_args, GCN(training_args, adj_matrix, sc_matrix) - ) - self.decode_spatial_position = deepcopy(self.encode_spatial_position) + # self.encode_temporal_position = TemporalPositionalEmbedding( + # training_args, max_len=training_args.his_len + # ) + # self.decode_temporal_position = TemporalPositionalEmbedding( + # training_args, max_len=training_args.tgt_len + # ) + self.temporal_section_week = TemporalSectionEmbedding(training_args, 7, axis=1) + self.temporal_section_day = TemporalSectionEmbedding(training_args, 288, axis=2) + # self.encode_spatial_position = SpatialPositionalEmbedding( + # training_args, + # GCN(training_args, training_args.d_proj, adj_matrix, sc_matrix), + # ) + # self.decode_spatial_position = deepcopy(self.encode_spatial_position) attn_ss = MultiHeadAttentionAwareTemporalContext( args=training_args, @@ -83,29 +96,93 @@ def __init__(self, training_args, adj_matrix, sc_matrix): training_args.d_model, training_args.decoder_output_size ) - def encode(self, src, lookup_index): - src_dense = self.encoder_dense(src) - src_tp_embedding = self.encode_temporal_position(src_dense, lookup_index) - src_sp_embedding = self.encode_spatial_position(src_tp_embedding) - encoder_output = self.encoder(src_sp_embedding) + def encode(self, src_idx, src): + src_dense = self.encoder_dense(src[..., :1]) + # tp_embedding = self.encode_temporal_position(src, src_idx) + # sp_embedding = self.encode_spatial_position(src) + week_embedding = self.temporal_section_week(src) + day_embedding = self.temporal_section_day(src) + + embed = src_dense # + tp_embedding + sp_embedding + embed = paddle.concat([embed, week_embedding, day_embedding], axis=-1) + + encoder_output = self.encoder(embed) return encoder_output - def decode(self, tgt, encoder_output): - tgt_dense = self.decoder_dense(tgt) - tgt_tp_embedding = self.decode_temporal_position(tgt_dense) - tgt_sp_embedding = self.decode_spatial_position(tgt_tp_embedding) - decoder_output = self.decoder(x=tgt_sp_embedding, memory=encoder_output) + def decode(self, encoder_output, tgt_idx=None, tgt=None): + tgt_dense = self.decoder_dense(tgt[..., :1]) + # tp_embedding = self.decode_temporal_position(tgt, tgt_idx) + # sp_embedding = self.decode_spatial_position(tgt) + week_embedding = self.temporal_section_week(tgt) + day_embedding = self.temporal_section_day(tgt) + + embed = tgt_dense # + tp_embedding + sp_embedding + embed = paddle.concat([embed, week_embedding, day_embedding], axis=-1) + + decoder_output = self.decoder(x=embed, memory=encoder_output) return self.generator(decoder_output) - def forward(self, src, src_idx, tgt): - encoder_output = self.encode(src, src_idx) - output = self.decode(tgt, encoder_output) + def forward(self, src_idx, src, tgt_idx=None, tgt=None): + encoder_output = self.encode(src_idx, src) + output = self.decode(encoder_output, tgt_idx, tgt) return output +class DecoderIndex(autograd.PyLayer): + @staticmethod + def forward(ctx, lags, his, his_span, interp_method="cubic"): + """ + 计算给定输入序列的未来值，并返回计算结果。 + 传入lags, history, + 计算序列位置对应位置的梯度, 并保存至backward + + Args: + ctx (): 动态图计算上下文对象。 + xde (): 未来值的输入序列, BaseXDE类型。 + lags (paddle.Tensor): 用多少个过去的值来计算未来的这个值（未来值的滞后量）。 + history (paddle.Tensor): 用于计算未来值的过去输入序列。 + interp_method (str, optional): 插值方法，取值为 "linear"（线性插值）,"cubic"（三次样条插值）或 "bez"（贝塞尔插值）。默认为 "linear"。 + + Returns: + paddle.Tensor: 计算结果，形状为 [batch_size, len_t, dims]。 + + Raises: + NotImplementedError: 如果interp_method不是上述三种情况之一, 将抛出NotImplementedError异常。 + """ + with paddle.no_grad(): + if interp_method == "linear": + interp = LinearInterpolation(his, his_span) + elif interp_method == "cubic": + interp = CubicHermiteSpline(his, his_span) + elif interp_method == "bez": + interp = BezierSpline(his, his_span) + else: + raise NotImplementedError + + y_lags = interp.evaluate(lags) + + derivative_lags = interp.derivative(lags) + ctx.save_for_backward(derivative_lags) + + return y_lags + + @staticmethod + def backward(ctx, grad_y): + # 计算history相应的梯度，并提取forward中保存的梯度，用于计算lag的梯度 + # 在计算的过程中，无需更新history，仅更新lags即可 + (derivative_lags,) = ctx.saved_tensor() + grad = grad_y * derivative_lags + grad = paddle.sum(grad, axis=[0, 1, 3]) + return grad, None, None + # return None, grad_y_lags * derivative_lags, None, None, None + + if __name__ == "__main__": + import os + + # 将日志级别设置为6 + os.environ["GLOG_v"] = "6" import numpy as np - import paddle import paddle.nn as nn from args import args from dataset import TrafficFlowDataset @@ -113,6 +190,9 @@ def forward(self, src, src_idx, tgt): from paddle.nn.initializer import XavierUniform from utils import get_adjacency_matrix_2direction, norm_adj_matrix + from paddlexde.functional import ddeint + from paddlexde.solver import Euler + default_dtype = paddle.get_default_dtype() adj_matrix, _ = get_adjacency_matrix_2direction(args.adj_path, 80) adj_matrix = paddle.to_tensor(norm_adj_matrix(adj_matrix), default_dtype) @@ -123,9 +203,48 @@ def forward(self, src, src_idx, tgt): nn.initializer.set_global_initializer(XavierUniform(), XavierUniform()) model = CorrSTN(args, adj_matrix, sc_matrix) + def collate_func(batch_data): + src_list, tgt_list = [], [] + + for item in batch_data: + if item[2]: + src_list.append(item[0]) + tgt_list.append(item[1]) + + if len(src_list) == 0: + src_list.append(item[0]) + tgt_list.append(item[1]) + + return paddle.stack(src_list), paddle.stack(tgt_list) + train_dataset = TrafficFlowDataset(args, "train") - train_dataloader = DataLoader(train_dataset, batch_size=args.batch_size) - his, tgt = next(iter(train_dataloader)) - his_index = paddle.randint(shape=[36], high=his.shape[-2], low=0) - his_input = paddle.index_select(his, his_index, axis=-2) - output = model(his_input, his_index, tgt) + train_dataloader = DataLoader( + train_dataset, batch_size=args.batch_size, collate_fn=collate_func + ) + src, tgt = next(iter(train_dataloader)) + src_index = paddle.randint(shape=[12], high=src.shape[-2], low=0) + src_input = paddle.index_select(src, src_index, axis=-2) + decoder_input = paddle.concat([src[:, :, -1:, :], tgt[:, :, :-1, :]], axis=-2) + + # 1. call model foreward (choose 1 or 2) + preds = model(src_index, src_input, None, decoder_input) + preds.backward() + + # 2. call ddeint (choose 1 or 2) + y0 = decoder_input + preds = ddeint( + func=model, + y0=y0, + t_span=paddle.arange(args.tgt_len + 1), + lags=src_index, + his=src, + his_span=paddle.arange(args.his_len), + solver=Euler, + ) + preds = preds[:, :, -args.tgt_len :, :] + preds.backward() + + from paddleviz.paddleviz.viz import make_graph + + dot = make_graph(preds, dpi="600") + dot.render("viz-result.gv", format="png", view=False) diff --git a/example/CorrSTN/dataset.py b/example/CorrSTN/dataset.py index 5ec525e..7d98d1e 100644 --- a/example/CorrSTN/dataset.py +++ b/example/CorrSTN/dataset.py @@ -135,8 +135,9 @@ def __init__(self, training_args, data_type="train"): # [T, N, D] # D=3 for PEMS04 and PEMS08, D=1 for others - self.origin_data = np.load(training_args.data_path)["data"].transpose([1, 0, 2]) - self.num_nodes, self.seq_len, self.dims = self.origin_data.shape + origin_data = np.load(training_args.data_path)["data"].transpose([1, 0, 2]) + origin_data = origin_data[:, :, :1] + self.num_nodes, self.seq_len, self.dims = origin_data.shape self.train_ratio, self.val_ratio, self.test_ratio = map( int, training_args.split.split(":") @@ -153,61 +154,73 @@ def __init__(self, training_args, data_type="train"): self.test_size = int(self.seq_len * self.test_ratio) self.data_type = data_type + # Scaler if training_args.scale: self.scaler = ScalerMinMax() - train_data = self.origin_data[: self.train_size, :, :] + train_data = origin_data[: self.train_size, :, :] self.scaler.fit(train_data.reshape(-1, train_data.shape[-1])) - self.data = self.scaler.transform(self.origin_data).reshape( + self.data = self.scaler.transform(origin_data).reshape( self.num_nodes, self.seq_len, self.dims ) else: - self.data = self.origin_data + self.data = origin_data - self.data = paddle.to_tensor(self.data, dtype=paddle.get_default_dtype()) + # Concat day of week and hour of day index + index = np.arange(0, self.seq_len, step=1).reshape([1, -1, 1]) + index = index.repeat(self.num_nodes, axis=0) + day_of_week_index = (index // 288) % 7 + hour_of_day_index = index % 288 - def __getitem__(self, index): + self.data = np.concatenate( + [self.data, day_of_week_index, hour_of_day_index], axis=-1 + ) if self.data_type == "train": - index += 0 + data_len = ( + self.train_size + - self.training_args.his_len + - self.training_args.tgt_len + ) elif self.data_type == "val": - index += self.train_size - self.training_args.his_len + data_len = self.val_size - self.training_args.tgt_len else: - index += self.train_size + self.val_size - self.training_args.his_len - - his_begin = index - his_end = his_begin + self.training_args.his_len - tgt_begin = his_end - tgt_end = tgt_begin + self.training_args.tgt_len + data_len = self.test_size - self.training_args.tgt_len - if "HZME" in self.training_args.dataset_name: - if tgt_begin % 288 < 72: - offset = 72 - his_begin += offset - his_end += offset - tgt_begin += offset - tgt_end += offset + self.his_pair = [] + self.tgt_pair = [] + for i in range(data_len): + if self.data_type == "train": + i += 0 + elif self.data_type == "val": + i += self.train_size - self.training_args.his_len + else: + i += self.train_size + self.val_size - self.training_args.his_len + + his_begin = i + his_end = his_begin + self.training_args.his_len + tgt_begin = his_end + tgt_end = tgt_begin + self.training_args.tgt_len + + if "HZME" in self.training_args.dataset_name and ( + tgt_begin % 288 < 72 or tgt_end % 288 < 72 + ): + continue + self.his_pair.append((his_begin, his_end)) + self.tgt_pair.append((tgt_begin, tgt_end)) - # print(self.data_type, his_begin, his_end, tgt_begin, tgt_end) + def __getitem__(self, index): + his_begin, his_end = self.his_pair[index] + tgt_begin, tgt_end = self.tgt_pair[index] - # [N, T, F] his = self.data[:, his_begin:his_end, :] tgt = self.data[:, tgt_begin:tgt_end, :] + # his_extend = his[:,-12:, :] + (his[:,-1:,:] - his[:,-12:-11,:]) + # his = np.concatenate([his, his_extend], axis=-2) return his, tgt def __len__(self): - if self.data_type == "train": - data_len = ( - self.train_size - - self.training_args.his_len - - self.training_args.tgt_len - ) - elif self.data_type == "val": - data_len = self.val_size - self.training_args.tgt_len - else: - data_len = self.test_size - self.training_args.tgt_len - - return data_len + return len(self.his_pair) def inverse_transform(self, data, axis=None): if self.training_args.scale: @@ -223,14 +236,27 @@ def inverse_transform(self, data, axis=None): val_dataset = TrafficFlowDataset(args, "val") test_dataset = TrafficFlowDataset(args, "test") - traing_dataloader = DataLoader(train_dataset, batch_size=10, shuffle=False) - val_dataloader = DataLoader(val_dataset, batch_size=10, shuffle=False) - test_dataloader = DataLoader(test_dataset, batch_size=10, shuffle=False) - - print(test_dataloader[0]) - - # for item in traing_dataloader: - # print(item[0].shape, item[1].shape) + traing_dataloader = DataLoader( + train_dataset, + batch_size=10, + shuffle=False, + num_workers=4, + ) + val_dataloader = DataLoader( + val_dataset, + batch_size=10, + shuffle=False, + num_workers=4, + ) + test_dataloader = DataLoader( + test_dataset, + batch_size=10, + shuffle=False, + num_workers=4, + ) + + for item in traing_dataloader: + print(item[0].shape, item[1].shape) # for item in val_dataloader: # print(item[0].shape, item[1].shape) diff --git a/example/CorrSTN/embedding.py b/example/CorrSTN/embedding.py index ddda699..202f373 100644 --- a/example/CorrSTN/embedding.py +++ b/example/CorrSTN/embedding.py @@ -1,19 +1,14 @@ import math -from typing import List -import numpy as np import paddle import paddle.nn as nn -import pandas as pd -from pandas.tseries import offsets -from pandas.tseries.frequencies import to_offset class SpatialPositionalEmbedding(nn.Layer): def __init__(self, args, gcn=None): super(SpatialPositionalEmbedding, self).__init__() self.dropout = nn.Dropout(p=args.dropout) - self.embedding = paddle.nn.Embedding(args.num_nodes, args.d_model) + self.embedding = paddle.nn.Embedding(args.num_nodes, args.d_proj) self.gcn_smooth_layers = None if (gcn is not None) and (args.smooth_layer_num > 0): self.gcn_smooth_layers = nn.LayerList( @@ -31,8 +26,7 @@ def forward(self, x): if self.gcn_smooth_layers is not None: for _, l in enumerate(self.gcn_smooth_layers): embed = l(embed) # [1,N,D] -> [1,N,D] - x = x + embed.unsqueeze(-2) # [B,N,T,D]+[1,N,1,D] - return self.dropout(x) + return embed.unsqueeze(-2) # [B,N,T,D]+[1,N,1,D] class TemporalPositionalEmbedding(nn.Layer): @@ -41,7 +35,7 @@ def __init__(self, args, max_len): self.args = args self.dropout = nn.Dropout(p=args.dropout) self.max_len = max_len - self.d_model = args.d_model + self.d_model = args.d_proj # computing the positional encodings once in log space pe = paddle.zeros([max_len, self.d_model]) for pos in range(max_len): @@ -71,70 +65,27 @@ def forward(self, x, lookup_index=None): if lookup_index is None: lookup_index = paddle.arange(x.shape[-2]) # [B,N,T,D] + [1,1,T,D] - x = x + paddle.index_select(self.pe, lookup_index, axis=-2) + embed = paddle.index_select(self.pe, lookup_index, axis=-2) else: + if lookup_index.dtype != paddle.int64: + lookup_index = paddle.cast(lookup_index, dtype="int64") + embed = paddle.index_select(self.pe, lookup_index, axis=-2) - x = x + paddle.index_select(self.pe, lookup_index, axis=-2) + return embed - return self.dropout(x.detach()) - - -class TimeEmbedding: - def __init__(self): - pass - - def __call__(self, index: pd.DatetimeIndex) -> np.ndarray: - pass - - def __repr__(self): - return self.__class__.__name__ + "()" - - -class MinuteOfHour(TimeEmbedding): - """Minute of hour encoded as value between [-1.0, 1.0]""" - - def __call__(self, index: pd.DatetimeIndex) -> np.ndarray: - return (index.minute / 59.0 - 0.5) * 2.0 - - -class HourOfDay(TimeEmbedding): - """Hour of day encoded as value between [-1.0, 1.0]""" - - def __call__(self, index: pd.DatetimeIndex) -> np.ndarray: - return (index.hour / 23.0 - 0.5) * 2.0 - - -def time_features_from_frequency_str(freq_str: str) -> List[TimeEmbedding]: - """ - Returns a list of time features that will be appropriate for the given frequency string. - Parameters - ---------- - freq_str - Frequency string of the form [multiple][granularity] such as "12H", "5min", "1D" etc. - """ - - features_by_offsets = { - offsets.Hour: [HourOfDay], - offsets.Minute: [MinuteOfHour, HourOfDay], - } - offset = to_offset(freq_str) - - for offset_type, feature_classes in features_by_offsets.items(): - if isinstance(offset, offset_type): - return [cls() for cls in feature_classes] - - -def time_embedding(dates, freq="10T"): - dates = pd.to_datetime(list(dates)) - res = [] - for feat in time_features_from_frequency_str(freq): - res.append(feat(dates)) - - return np.vstack(res).transpose(1, 0) +class TemporalSectionEmbedding(nn.Layer): + def __init__(self, args, section_nums, axis=1): + """ + axis=1 indicate day of week + axis=2 indicate hour of day + """ + super(TemporalSectionEmbedding, self).__init__() + self.axis = axis + self.embedding = paddle.nn.Embedding(section_nums, args.d_sect) -if __name__ == "__main__": - dataStrs = ["00:10", "12:30", "13:30", "22:20"] - time_embedding_res = time_embedding(dates=dataStrs) - print(time_embedding_res) - print() + def forward(self, x): + input = x[..., self.axis] + input = paddle.clip(input,min=0,max=self.embedding._num_embeddings-1) + input = paddle.cast(input, dtype=paddle.int32) + return self.embedding(input) diff --git a/example/CorrSTN/graphconv.py b/example/CorrSTN/graphconv.py index b5b291e..494053f 100644 --- a/example/CorrSTN/graphconv.py +++ b/example/CorrSTN/graphconv.py @@ -6,19 +6,22 @@ class GCN(nn.Layer): - def __init__(self, training_args, norm_adj_matrix, norm_sc_matrix): + def __init__(self, training_args, d_model, norm_adj_matrix, norm_sc_matrix): super(GCN, self).__init__() - self.register_buffer("norm_adj_matrix", norm_adj_matrix) - self.register_buffer("norm_sc_matrix", norm_sc_matrix) + self.norm_adj_matrix = norm_adj_matrix + self.norm_sc_matrix = norm_sc_matrix self.Theta = nn.Linear( - training_args.d_model, - training_args.d_model, + d_model, + d_model, bias_attr=False, ) self.alpha = paddle.create_parameter( shape=[1], dtype=paddle.get_default_dtype(), ) + if training_args.no_adj: + self.alpha.set_value(paddle.to_tensor([0.0])) + self.alpha.stop_gradient = True self.beta = paddle.create_parameter(shape=[1], dtype=paddle.get_default_dtype()) def forward(self, x): @@ -61,8 +64,8 @@ def forward(self, x): class SpatialAttentionGCN(nn.Layer): def __init__(self, args, adj_matrix, sc_matrix, is_scale=True): super(SpatialAttentionGCN, self).__init__() - self.register_buffer("norm_adj", adj_matrix) - self.register_buffer("norm_sc", sc_matrix) + self.norm_adj = adj_matrix + self.norm_sc = sc_matrix self.args = args self.linear = nn.Linear(args.d_model, args.d_model, bias_attr=False) self.is_scale = is_scale @@ -70,6 +73,9 @@ def __init__(self, args, adj_matrix, sc_matrix, is_scale=True): self.alpha = paddle.create_parameter( shape=[1], dtype=paddle.get_default_dtype() ) + if args.no_adj: + self.alpha.set_value(paddle.to_tensor([0.0])) + self.alpha.stop_gradient = True self.beta = paddle.create_parameter(shape=[1], dtype=paddle.get_default_dtype()) def forward(self, x): diff --git a/example/CorrSTN/train.py b/example/CorrSTN/train.py index b920854..89853d4 100644 --- a/example/CorrSTN/train.py +++ b/example/CorrSTN/train.py @@ -65,6 +65,8 @@ def __init__(self, training_args): self._build_data() self._build_model() self._build_optim() + if training_args.distribute: + self._build_distribute() def _build_data(self): self.train_dataset = TrafficFlowDataset(self.training_args, "train") @@ -172,6 +174,51 @@ def _build_optim(self): for var_name in self.optimizer.state_dict(): self.logger.info(f"{var_name} \t {self.optimizer.state_dict()[var_name]}") + def _build_distribute(self): + # 一、导入分布式专用 Fleet API + from paddle.distributed import fleet + + # 构建分布式数据加载器所需 API + from paddle.io import DataLoader, DistributedBatchSampler + + # 二、初始化 Fleet 环境 + fleet.init(is_collective=True) + + # 三、构建分布式训练使用的网络模型 + self.net = fleet.distributed_model(self.net) + + # 四、构建分布式训练使用的优化器 + self.optimizer = fleet.distributed_optimizer(self.optimizer) + + # 五、构建分布式训练使用的数据集 + train_sampler = DistributedBatchSampler( + self.train_dataset, + batch_size=self.training_args.batch_size, + shuffle=True, + drop_last=False, + ) + self.train_dataloader = DataLoader( + self.train_dataset, batch_sampler=train_sampler, num_workers=2 + ) + eval_sampler = DistributedBatchSampler( + self.val_dataset, + batch_size=self.training_args.batch_size * 32, + shuffle=False, + drop_last=False, + ) + self.eval_dataloader = DataLoader( + self.val_dataset, batch_sampler=eval_sampler, num_workers=2 + ) + test_sampler = DistributedBatchSampler( + self.test_dataset, + batch_size=self.training_args.batch_size * 32, + shuffle=False, + drop_last=False, + ) + self.test_dataloader = DataLoader( + self.test_dataset, batch_sampler=test_sampler, num_workers=2 + ) + def train(self): self.logger.info("start train...") @@ -187,6 +234,7 @@ def train(self): # finetune => load best trainging model if epoch == self.training_args.train_epochs: self._init_finetune() + self.compute_test_loss() self.net.train() # ensure dropout layers are in train mode tr_s_time = time() @@ -209,7 +257,7 @@ def train(self): best_epoch = epoch self.logger.info(f"best_epoch: {best_epoch}") self.logger.info(f"eval_loss: {eval_loss}") - self.compute_test_loss() + # self.compute_test_loss() # save parameters # params_filename = os.path.join(self.save_path, f"epoch_{epoch}.params") params_filename = os.path.join(self.save_path, "epoch_best.params") @@ -264,8 +312,11 @@ def train_one_step(self, src, tgt): with amp_guard_context(self.training_args.fp16): if not self.finetune: + decoder_input = paddle.concat( + [src[:, :, -1:, :], tgt[:, :, :-1, :]], axis=-2 + ) decoder_output = self.net( - src=encoder_input, src_idx=self.encoder_idx, tgt=tgt + src=encoder_input, src_idx=self.encoder_idx, tgt=decoder_input ) else: decoder_start_inputs = encoder_input[:, :, -1:, :] diff --git a/example/CorrSTN/train_dde.py b/example/CorrSTN/train_dde.py new file mode 100644 index 0000000..a272c73 --- /dev/null +++ b/example/CorrSTN/train_dde.py @@ -0,0 +1,592 @@ +import contextlib +import os +from time import time + +import args +import numpy as np +import paddle +import paddle.nn as nn +import paddle.optimizer as optim +from corrstn import CorrSTN, DecoderIndex +from dataset import TrafficFlowDataset +from paddle.distributed import fleet +from paddle.io import DataLoader, DistributedBatchSampler +from paddle.nn.initializer import Constant, XavierUniform +from sklearn.metrics import mean_absolute_error, mean_squared_error +from utils import ( + CosineAnnealingWithWarmupDecay, + EarlyStopping, + Logger, + get_adjacency_matrix_2direction, + masked_mape_np, + norm_adj_matrix, +) +from visualdl import LogWriter + +from paddlexde.functional import ddeint +from paddlexde.solver.fixed_solver import Euler, RK4 + + +def amp_guard_context(fp16=False): + if fp16: + return paddle.amp.auto_cast(level="O2") + else: + return contextlib.nullcontext() + + +class Trainer: + def __init__(self, training_args): + + self.training_args = training_args + + self.folder_dir = ( + f"MAE_{training_args.model_name}_elayer{training_args.encoder_num_layers}_" + + f"dlayer{training_args.decoder_num_layers}_head{training_args.head}_dm{training_args.d_model}_" + + f"einput{training_args.encoder_input_size}_dinput{training_args.decoder_input_size}_" + + f"doutput{training_args.decoder_output_size}_drop{training_args.dropout}_" + + f"lr{training_args.learning_rate}_wd{training_args.weight_decay}_bs{training_args.batch_size}_" + + f"topk{training_args.top_k}_att{training_args.attention}_trepoch{training_args.train_epochs}_" + + f"finepoch{training_args.finetune_epochs}_dde" + ) + + self.save_path = os.path.join( + "experiments", training_args.dataset_name, self.folder_dir + ) + os.makedirs(self.save_path, exist_ok=True) + self.logger = Logger("CorrSTN", os.path.join(self.save_path, "log.txt")) + self.writer = LogWriter(logdir=os.path.join(self.save_path, "visualdl")) + + if training_args.start_epoch == 0: + self.logger.info(f"create params directory {self.save_path}") + elif training_args.start_epoch > 0: + self.logger.info(f"train from params directory {self.save_path}") + + self.logger.info(f"save folder: {self.folder_dir}") + self.logger.info(f"save path : {self.save_path}") + self.logger.info(f"log file : {self.logger.log_file}") + + args_message = "\n".join( + [f"{k:<20}: {v}" for k, v in vars(training_args).items()] + ) + self.logger.info(f"training_args : \n{args_message}") + self.finetune = False + self.early_stopping = EarlyStopping(patience=training_args.patience, delta=0.0) + + self._build_data() + self._build_model() + self._build_optim() + if training_args.distribute: + self._build_distribute() + + def _build_data(self): + self.train_dataset = TrafficFlowDataset(self.training_args, "train") + self.val_dataset = TrafficFlowDataset(self.training_args, "val") + self.test_dataset = TrafficFlowDataset(self.training_args, "test") + + self.train_dataloader = DataLoader( + self.train_dataset, + batch_size=self.training_args.batch_size, + shuffle=True, + drop_last=True, + num_workers=4, + ) + self.eval_dataloader = DataLoader( + self.val_dataset, + batch_size=self.training_args.batch_size, + shuffle=False, + drop_last=False, + num_workers=4, + ) + self.test_dataloader = DataLoader( + self.test_dataset, + batch_size=self.training_args.batch_size, + shuffle=False, + drop_last=False, + num_workers=4, + ) + + # 保持输入序列长度为12 + self.fix_week = paddle.arange( + start=self.training_args.his_len - 2016, + end=self.training_args.his_len - 2016 + 12, + ) + self.fix_day = paddle.arange( + start=self.training_args.his_len - 288, + end=self.training_args.his_len - 288 + 12, + ) + self.fix_hour = paddle.arange( + start=self.training_args.his_len - 12, + end=self.training_args.his_len, + ) + self.fix_pred = paddle.arange( + start=self.training_args.his_len, + end=self.training_args.his_len + 12, + ) + self.fix_pred = paddle.ones(shape=[self.training_args.tgt_len]) * ( + self.training_args.his_len - 1 + ) + + encoder_idx = [] + decoder_idx = [self.fix_pred] + + # for week + if self.training_args.his_len >= 2016: + encoder_idx.append(self.fix_week) + + # for day + elif self.training_args.his_len >= 288: + encoder_idx.append(self.fix_day) + + # for hour + elif self.training_args.his_len >= 12: + encoder_idx.append(self.fix_hour) + + # concat all + encoder_idx = paddle.concat(encoder_idx) + decoder_idx = paddle.concat(decoder_idx) + + if self.training_args.fp16: + self.encoder_idx = paddle.create_parameter( + shape=encoder_idx.shape, dtype="float16" + ) + self.decoder_idx = paddle.create_parameter( + shape=decoder_idx.shape, dtype="float16" + ) + self.encoder_idx.set_value(paddle.cast(encoder_idx, "float16")) + self.decoder_idx.set_value(paddle.cast(decoder_idx, "float16")) + else: + self.encoder_idx = paddle.create_parameter( + shape=encoder_idx.shape, dtype="float32" + ) + self.decoder_idx = paddle.create_parameter( + shape=decoder_idx.shape, dtype="float32" + ) + self.encoder_idx.set_value(paddle.cast(encoder_idx, "float32")) + self.decoder_idx.set_value(paddle.cast(decoder_idx, "float32")) + + self.logger.info(f"encoder_idx: {self.encoder_idx}") + self.logger.info(f"decoder_idx: {self.decoder_idx}") + + def _build_model(self): + default_dtype = paddle.get_default_dtype() + adj_matrix, _ = get_adjacency_matrix_2direction( + self.training_args.adj_path, self.training_args.num_nodes + ) + adj_matrix = paddle.to_tensor(norm_adj_matrix(adj_matrix), default_dtype) + + sc_matrix = np.load(self.training_args.sc_path)[0, :, :] + sc_matrix = paddle.to_tensor(norm_adj_matrix(sc_matrix), default_dtype) + + nn.initializer.set_global_initializer(XavierUniform(), Constant(value=0.0)) + + self.net = CorrSTN( + self.training_args, + adj_matrix=adj_matrix, + sc_matrix=sc_matrix, + ) + + if self.training_args.fp16: + self.net = paddle.amp.decorate(models=self.net, level="O2") + self.scaler = paddle.amp.GradScaler(init_loss_scaling=1024) + + if self.training_args.continue_training: + self.load() + + self.logger.debug(self.net) + + total_param = 0 + self.logger.debug("Net's state_dict:") + for param_tensor in self.net.state_dict(): + self.logger.debug( + f"{param_tensor} \t {self.net.state_dict()[param_tensor].shape}" + ) + total_param += np.prod(self.net.state_dict()[param_tensor].shape) + self.logger.debug(f"Net's total params: {total_param}.") + + self.criterion = nn.L1Loss() # 定义损失函数 + + def _build_optim(self): + self.lr_scheduler = CosineAnnealingWithWarmupDecay( + max_lr=1, + min_lr=0.1, + warmup_step=0.2 * self.training_args.train_epochs, + decay_step=0.8 * self.training_args.train_epochs, + ) + + parameters = [ + { + "params": self.net.parameters(), + "learning_rate": self.training_args.learning_rate, + }, + { + "params": [self.decoder_idx], + "learning_rate": self.training_args.learning_rate * 0.1, + }, + { + "params": [self.encoder_idx], + "learning_rate": self.training_args.learning_rate * 0.1, + }, + ] + + # 定义优化器，传入所有网络参数 + self.optimizer = optim.Adam( + parameters=parameters, + learning_rate=self.lr_scheduler, + weight_decay=self.training_args.weight_decay, + multi_precision=True, + ) + + self.logger.info("Optimizer's state_dict:") + for var_name in self.optimizer.state_dict(): + self.logger.info(f"{var_name} \t {self.optimizer.state_dict()[var_name]}") + + self.dde_solver = Euler + + def _build_distribute(self): + # 二、初始化 Fleet 环境 + fleet.init(is_collective=True) + + # 三、构建分布式训练使用的网络模型 + self.net = fleet.distributed_model(self.net) + + # 四、构建分布式训练使用的优化器 + self.optimizer = fleet.distributed_optimizer(self.optimizer) + + # 五、构建分布式训练使用的数据集 + train_sampler = DistributedBatchSampler( + self.train_dataset, + batch_size=self.training_args.batch_size, + shuffle=True, + drop_last=False, + ) + self.train_dataloader = DataLoader( + self.train_dataset, batch_sampler=train_sampler, num_workers=12 + ) + eval_sampler = DistributedBatchSampler( + self.val_dataset, + batch_size=self.training_args.batch_size, + shuffle=False, + drop_last=False, + ) + self.eval_dataloader = DataLoader( + self.val_dataset, batch_sampler=eval_sampler, num_workers=12 + ) + test_sampler = DistributedBatchSampler( + self.test_dataset, + batch_size=self.training_args.batch_size, + shuffle=False, + drop_last=False, + ) + self.test_dataloader = DataLoader( + self.test_dataset, batch_sampler=test_sampler, num_workers=12 + ) + + def save(self, epoch=None): + if epoch is not None: + params_filename = os.path.join(self.save_path, f"epoch_{epoch}.params") + encoder_idx_filename = os.path.join(self.save_path, f"epoch_{epoch}.enidx") + decoder_idx_filename = os.path.join(self.save_path, f"epoch_{epoch}.deidx") + else: + params_filename = os.path.join(self.save_path, "epoch_best.params") + encoder_idx_filename = os.path.join(self.save_path, "epoch_best.enidx") + decoder_idx_filename = os.path.join(self.save_path, "epoch_best.deidx") + paddle.save(self.net.state_dict(), params_filename) + paddle.save(self.encoder_idx, encoder_idx_filename) + paddle.save(self.decoder_idx, decoder_idx_filename) + self.logger.info(f"save parameters to file: {params_filename}") + + def load(self, epoch=None): + if epoch is not None: + params_filename = os.path.join(self.save_path, f"epoch_{epoch}.params") + encoder_idx_filename = os.path.join(self.save_path, f"epoch_{epoch}.enidx") + decoder_idx_filename = os.path.join(self.save_path, f"epoch_{epoch}.deidx") + else: + params_filename = os.path.join(self.save_path, "epoch_best.params") + encoder_idx_filename = os.path.join(self.save_path, "epoch_best.enidx") + decoder_idx_filename = os.path.join(self.save_path, "epoch_best.deidx") + + self.net.set_state_dict(paddle.load(params_filename)) + self.encoder_idx.set_value(paddle.load(encoder_idx_filename)) + self.decoder_idx.set_value(paddle.load(decoder_idx_filename)) + self.logger.info(f"load weight from: {params_filename}") + + def train(self): + self.logger.info("start train...") + + s_time = time() + best_eval_loss = np.inf + best_epoch = 0 + global_step = 0 + epoch = self.training_args.start_epoch + + while ( + epoch < self.training_args.train_epochs + self.training_args.finetune_epochs + ): + # finetune => load best trainging model + if epoch == self.training_args.train_epochs: + self._init_finetune() + self.compute_test_loss(epoch) + + self.net.train() # ensure dropout layers are in train mode + tr_s_time = time() + epoch_step = 0 + self.lr_scheduler.step() + for batch_index, batch_data in enumerate(self.train_dataloader): + src, tgt = batch_data + src = paddle.cast(src, paddle.get_default_dtype()) + tgt = paddle.cast(tgt, paddle.get_default_dtype()) + _, training_loss = self.train_one_step(src, tgt) + self.writer.add_scalar("train/loss", training_loss, global_step) + self.writer.add_scalar("train/lr", self.optimizer.get_lr(), global_step) + epoch_step += 1 + global_step += 1 + self.logger.info(f"learning_rate: {self.optimizer.get_lr()}") + self.logger.info(f"epoch: {epoch}, train time cost:{time() - tr_s_time}") + self.logger.info(f"epoch: {epoch}, total time cost:{time() - s_time}") + + # apply model on the validation data set + eval_loss = self.compute_eval_loss(epoch) + if eval_loss < best_eval_loss: + best_eval_loss = eval_loss + best_epoch = epoch + self.logger.info(f"best_epoch: {best_epoch}") + self.logger.info(f"eval_loss: {float(eval_loss)}") + self.compute_test_loss(epoch) + # save parameters + self.save(epoch=epoch) + self.save() + + self.early_stopping(val_loss=eval_loss) + if self.early_stopping.early_stop: + self.logger.info("Early stopping") + if epoch < self.training_args.train_epochs: + epoch = self.training_args.train_epochs + else: + break + else: + epoch += 1 + + self.logger.info(f"best epoch: {best_epoch}") + self.logger.info("apply the best val model on the test dataset ...") + + self.load() + self.compute_test_loss(epoch) + + def _init_finetune(self): + self.logger.info("Start FineTune Training") + self.load() + + self.early_stopping.reset() + + self.lr_scheduler = CosineAnnealingWithWarmupDecay( + max_lr=1, + min_lr=0.1, + warmup_step=0.2 * self.training_args.finetune_epochs, + decay_step=0.8 * self.training_args.finetune_epochs, + ) + + parameters = [ + { + "params": self.net.parameters(), + "learning_rate": self.training_args.learning_rate * 0.1, + }, + { + "params": [self.decoder_idx], + "learning_rate": self.training_args.learning_rate, + }, + { + "params": [self.encoder_idx], + "learning_rate": self.training_args.learning_rate, + }, + ] + + # 定义优化器，传入所有网络参数 + self.optimizer = optim.Adam( + parameters=parameters, + learning_rate=self.lr_scheduler, + weight_decay=self.training_args.weight_decay, + multi_precision=True, + ) + self.dde_solver = RK4 + if self.training_args.distribute: + self.optimizer = fleet.distributed_optimizer(self.optimizer) + self.finetune = True + + def train_one_step(self, src, tgt): + """_summary_ + + Args: + src (_type_): [B,N,T,D] + tgt (_type_): [B,N,T,D] + + Returns: + _type_: _description_ + """ + self.net.train() + + with amp_guard_context(self.training_args.fp16): + # self.decoder_idx.stop_gradient= True + y0 = DecoderIndex.apply( + lags=self.decoder_idx, + his=src, + his_span=paddle.arange(self.training_args.his_len), + ) + + preds = ddeint( + func=self.net, + y0=y0, + t_span=paddle.arange(1 + 1), + lags=self.encoder_idx, + his=src, + his_span=paddle.arange(self.training_args.his_len), + solver=self.dde_solver, + ) + pred_len = y0.shape[-2] + preds = preds[:, :, -pred_len:, :1] + + loss = self.criterion(preds, tgt[..., :1]) + if self.net.training: + if self.training_args.fp16: + scaled = self.scaler.scale(loss) # loss 缩放，乘以系数 loss_scaling + scaled.backward() # 反向传播 + self.scaler.step(self.optimizer) # 更新参数（参数梯度先除系数 loss_scaling 再更新参数） + self.scaler.update() # 基于动态 loss_scaling 策略更新 loss_scaling 系数 + self.optimizer.clear_grad(set_to_zero=False) + else: + loss.backward() + self.optimizer.step() + self.optimizer.clear_grad() + return preds, loss + + def eval_one_step(self, src, tgt): + self.net.eval() + with amp_guard_context(self.training_args.fp16): + y0 = DecoderIndex.apply( + lags=self.decoder_idx, + his=src, + his_span=paddle.arange(self.training_args.his_len), + ) + preds = ddeint( + func=self.net, + y0=y0, + t_span=paddle.arange(1 + 1), + lags=self.encoder_idx, + his=src, + his_span=paddle.arange(self.training_args.his_len), + solver=self.dde_solver, + ) + pred_len = y0.shape[-2] + preds = preds[:, :, -pred_len:, :1] + + loss = self.criterion(preds, tgt[..., :1]) + + return preds, loss + + def test_one_step(self, src, tgt): + self.net.eval() + with amp_guard_context(self.training_args.fp16): + y0 = DecoderIndex.apply( + lags=self.decoder_idx, + his=src, + his_span=paddle.arange(self.training_args.his_len), + ) + preds = ddeint( + func=self.net, + y0=y0, + t_span=paddle.arange(1 + 1), + lags=self.encoder_idx, + his=src, + his_span=paddle.arange(self.training_args.his_len), + solver=self.dde_solver, + ) + pred_len = y0.shape[-2] + preds = preds[:, :, -pred_len:, :1] + + loss = self.criterion(preds, tgt[..., :1]) + + return preds, loss + + def compute_eval_loss(self, epoch=-1): + with paddle.no_grad(): + all_eval_loss = paddle.zeros([1], dtype=paddle.get_default_dtype()) + start_time = time() + for batch_index, batch_data in enumerate(self.eval_dataloader): + src, tgt = batch_data + src = paddle.cast(src, paddle.get_default_dtype()) + tgt = paddle.cast(tgt, paddle.get_default_dtype()) + predict_output, eval_loss = self.eval_one_step(src, tgt) + self.writer.add_scalar(f"eval/loss-{epoch}", eval_loss, batch_index) + + all_eval_loss += eval_loss + + eval_loss = all_eval_loss / len(self.eval_dataloader) + self.logger.info(f"eval cost time: {time() - start_time} s") + self.logger.info(f"eval_loss: {float(eval_loss)}") + return eval_loss + + def compute_test_loss(self, epoch=-1): + with paddle.no_grad(): + preds = [] + tgts = [] + start_time = time() + for batch_index, batch_data in enumerate(self.test_dataloader): + src, tgt = batch_data + src = paddle.cast(src, paddle.get_default_dtype()) + tgt = paddle.cast(tgt, paddle.get_default_dtype()) + predict_output, test_loss = self.test_one_step(src, tgt) + self.writer.add_scalar(f"test/loss-{epoch}", test_loss, batch_index) + + preds.append(predict_output) + tgts.append(tgt[..., :1]) + self.logger.info(f"test time on whole data: {time() - start_time} s") + + preds = paddle.concat(preds, axis=0) # [B,N,T,1] + trues = paddle.concat(tgts, axis=0) # [B,N,T,F] + # [B,N,T,1] + preds = self.test_dataset.inverse_transform(preds, axis=-1).numpy() + # [B,N,T,1] + trues = self.test_dataset.inverse_transform(trues, axis=-1).numpy() + + self.logger.info(f"preds: {preds.shape}") + self.logger.info(f"tgts: {trues.shape}") + + for index in range(trues.shape[0]): + scalar_dict = { + "true": trues[index, 0, 6, 0], + "pred": preds[index, 0, 6, 0], + } + self.writer.add_scalars(f"test/line-{epoch}", scalar_dict, index) + + # 计算误差 + excel_list = [] + prediction_length = trues.shape[2] + + for i in range(prediction_length): + assert preds.shape[0] == trues.shape[0] + mae = mean_absolute_error(trues[:, :, i, 0], preds[:, :, i, 0]) + rmse = mean_squared_error(trues[:, :, i, 0], preds[:, :, i, 0]) ** 0.5 + mape = masked_mape_np(trues[:, :, i, 0], preds[:, :, i, 0], 0) + self.logger.info(f"{i} MAE: {mae}") + self.logger.info(f"{i} RMSE: {rmse}") + self.logger.info(f"{i} MAPE: {mape}") + excel_list.extend([mae, rmse, mape]) + + # print overall results + mae = mean_absolute_error(trues.reshape(-1, 1), preds.reshape(-1, 1)) + rmse = mean_squared_error(trues.reshape(-1, 1), preds.reshape(-1, 1)) ** 0.5 + mape = masked_mape_np(trues.reshape(-1, 1), preds.reshape(-1, 1), 0) + self.logger.info(f"all MAE: {mae}") + self.logger.info(f"all RMSE: {rmse}") + self.logger.info(f"all MAPE: {mape}") + excel_list.extend([mae, rmse, mape]) + self.logger.info(excel_list) + + def run_test(self): + self.load() + self.compute_test_loss() + + +if __name__ == "__main__": + trainer = Trainer(training_args=args.args) + trainer.train() + trainer.run_test() diff --git a/example/CorrSTN/utils.py b/example/CorrSTN/utils.py index ba32dbb..85a394b 100644 --- a/example/CorrSTN/utils.py +++ b/example/CorrSTN/utils.py @@ -73,7 +73,7 @@ def _get_logger(self): # 创建一个logger logger = logging.getLogger(self.name) # 设置logger的日志级别 - logger.setLevel(logging.DEBUG) + logger.setLevel(logging.INFO) # 创建一个handler，用于写入日志文件 file_handler = logging.FileHandler(self.log_file) # 定义handler的输出格式 diff --git a/example/dde_demo.py b/example/dde_demo.py index cf444a5..446c33e 100644 --- a/example/dde_demo.py +++ b/example/dde_demo.py @@ -1,9 +1,9 @@ import paddle import paddle.nn as nn +from demo_utils import DemoUtils, SimpleDemoData -from example.demo_utils import DemoUtils, SimpleDemoData from paddlexde.functional import ddeint, ddeint_adjoint -from paddlexde.solver.fixed_solver import Euler +from paddlexde.solver.fixed_solver import RK4 demo_utils = DemoUtils() if demo_utils.args.adjoint: @@ -21,6 +21,7 @@ def __init__(self, his_len): def __getitem__(self, idx): his = self.true_y[idx : idx + self.his_len, :] his_span = self.t_span[idx : idx + self.his_len] + # his_span = paddle.arange(self.his_len) return super().__getitem__(idx + self.his_len) + (his, his_span) @@ -34,6 +35,13 @@ def __init__(self): self.linear1 = nn.Linear(2, 128) self.linear2 = nn.Linear(128, 2) + self.gru = paddle.nn.GRU(2, 128, 2, time_major=False) + + self.linear1.weight.set_value(0.1 * paddle.randn(self.linear1.weight.shape)) + self.linear1.bias.set_value(paddle.zeros(self.linear1.bias.shape)) + self.linear2.weight.set_value(0.1 * paddle.randn(self.linear2.weight.shape)) + self.linear2.bias.set_value(paddle.zeros(self.linear2.bias.shape)) + def forward(self, t, y0, lags, y_lags): """_summary_ @@ -46,22 +54,23 @@ def forward(self, t, y0, lags, y_lags): Returns: _type_: [B, D], the shape should match y0 shape """ - h = self.linear1(y0).unsqueeze(-2) # [10, 1, 256] - h_his = self.linear1(y_lags) # [10, 20, 256] - h_his = paddle.sum(h_his, axis=1, keepdim=True) # [10, 1, 256] - re = self.linear2(h + h_his).squeeze(-2) # [10, 128] + + h = self.linear1(y0**3) # [B, D] => [B, D] + h_lags = self.gru(y_lags)[0][:, -1, :] # [B, T, D] => [B, D] + h = (h + h_lags) / 2 + re = self.linear2(paddle.tanh(h)) # [10, 128] return re if __name__ == "__main__": his_len = demo_utils.args.his_len + pred_len = demo_utils.args.pred_len dde_dataset = DDEDataset(his_len=his_len) demo_utils.make_dataloader(dde_dataset) - # [B, T], T is pred_len - lags = paddle.randint( - low=0, high=his_len, shape=[demo_utils.args.batch_size, 20] - ).astype("float32") + # [T], T is pred_len + lags = paddle.randint(low=0, high=his_len, shape=[32]) + lags = paddle.cast(lags, dtype=paddle.float32) lags.stop_gradient = False func = DDEFunc() @@ -80,26 +89,29 @@ def forward(self, t, y0, lags, y_lags): batch_his_span, ) in demo_utils.dataloader: # batch_y0 : [B, D] - # batch_t_span : [B, T], T is pres_len + # t_span : [T], T is pres_len # batch_y : [B, T, D], T is pred_len # batch_his : [B, T, D], T is his_len - # batch_his_span : [B, T], T is his_len + # his_span : [T], T is his_len + # print(lags.numpy()) + t_span = paddle.linspace(0.0, 25.0, demo_utils.args.data_len)[:pred_len] + his_span = paddle.arange(his_len) pred_y = xdeint( func, batch_y0, - batch_t_span, + t_span, lags, batch_his, - batch_his_span, - solver=Euler, + his_span, + solver=RK4, ) loss = paddle.mean(paddle.abs(pred_y - batch_y)) loss.backward() optimizer.step() optimizer.clear_grad() - lags = paddle.assign(lags.detach().clip(0, his_len)) - lags.stop_gradient = False + # lags = paddle.assign(lags.detach().clip(0, his_len)) + # lags.stop_gradient = False global_step += 1 # retain_graph=True @@ -111,13 +123,11 @@ def forward(self, t, y0, lags, y_lags): with paddle.no_grad(): data = demo_utils.data y0 = data.true_y[his_len].unsqueeze(0) # [1, D] - t_span = data.t_span[his_len + 1 :].unsqueeze(0) # [1, T] + t_span = data.t_span[his_len + 1 :] # [T] true_y = data.true_y[his_len + 1 :].unsqueeze(0) # [1, T, D] his = data.true_y[:his_len].unsqueeze(0) # [1, T, D] - his_span = data.t_span[:his_len].unsqueeze(0) # [1, T] - pred_y = xdeint( - func, y0, t_span, lags[0:1, ...], his, his_span, solver=Euler - ) + his_span = data.t_span[:his_len] # [T] + pred_y = xdeint(func, y0, t_span, lags, his, his_span, solver=RK4) loss = paddle.mean(paddle.abs(pred_y - true_y)) print( "Iter {:04d} | Total Loss {:.6f}".format( diff --git a/example/demo_utils.py b/example/demo_utils.py index 8a81d57..d8cc529 100644 --- a/example/demo_utils.py +++ b/example/demo_utils.py @@ -41,7 +41,7 @@ def make_args(self): parser.add_argument("--adjoint", type=bool, default=False) # DDE - parser.add_argument("--his_len", type=int, default=64) + parser.add_argument("--his_len", type=int, default=288) self.args = parser.parse_args() def make_dataloader(self, dataset=None): @@ -159,20 +159,15 @@ def __init__(self, config): self.true_y = odeint( Lambda(self.trans_matrix), self.true_y0.unsqueeze(0), - self.t_span.unsqueeze(0), + self.t_span, solver=RK4, ).squeeze(0) def __getitem__(self, idx): - y0, t_span, tgt_y = self.true_y[idx, :], [], [] - for i in range(self.config.pred_len): - t_span.append(self.t_span[idx + i]) - tgt_y.append(self.true_y[idx + i, :]) - try: - t_span = paddle.concat(t_span) - except: - t_span = paddle.stack(t_span) - tgt_y = paddle.stack(tgt_y) + y0 = self.true_y[idx, :] + selcect_idx = paddle.arange(idx, idx + self.config.pred_len) + t_span = paddle.index_select(self.t_span, selcect_idx) + tgt_y = paddle.index_select(self.true_y, selcect_idx, axis=0) # [D], [T], [T, D], T is pred_len return y0, t_span, tgt_y diff --git a/example/ode_demo.py b/example/ode_demo.py index ca0dab8..5afa07a 100644 --- a/example/ode_demo.py +++ b/example/ode_demo.py @@ -1,10 +1,12 @@ import paddle import paddle.nn as nn +from demo_utils import DemoUtils -from example.demo_utils import DemoUtils from paddlexde.functional import odeint, odeint_adjoint from paddlexde.solver.fixed_solver import RK4 +paddle.seed(42) + demo_utils = DemoUtils() if demo_utils.args.adjoint: xdeint = odeint_adjoint @@ -32,6 +34,7 @@ def forward(self, t, y): if __name__ == "__main__": + pred_len = demo_utils.args.pred_len func = ODEFunc() optimizer = paddle.optimizer.RMSProp( parameters=func.parameters(), learning_rate=1e-3 @@ -44,7 +47,8 @@ def forward(self, t, y): # batch_y0 : [B, D] # batch_t : [B, T] # batch_y : [B, T, D] - pred_y = odeint(func, batch_y0, batch_t, solver=RK4) + t_span = paddle.linspace(0.0, 25.0, demo_utils.args.data_len)[:pred_len] + pred_y = odeint(func, batch_y0, t_span, solver=RK4) loss = paddle.mean(paddle.abs(pred_y - batch_y)) loss.backward() optimizer.step() @@ -58,7 +62,7 @@ def forward(self, t, y): if global_step % demo_utils.args.test_steps == 0: with paddle.no_grad(): y0 = demo_utils.data.true_y0.unsqueeze(0) # [1, D] - t_span = demo_utils.data.t_span.unsqueeze(0) # [1, T] + t_span = demo_utils.data.t_span # [T] true_y = demo_utils.data.true_y.unsqueeze(0) # [1, T, D] pred_y = xdeint(func, y0, t_span, solver=RK4) loss = paddle.mean(paddle.abs(pred_y - true_y)) diff --git a/paddlexde/functional/ddeint.py b/paddlexde/functional/ddeint.py index f9ecede..abf6448 100644 --- a/paddlexde/functional/ddeint.py +++ b/paddlexde/functional/ddeint.py @@ -35,6 +35,6 @@ def ddeint( s = solver(xde=xde, y0=xde.y0, rtol=rtol, atol=atol, **options) solution = s.integrate(t_span) - solution = xde.format(solution) + # solution = xde.format(solution) return solution diff --git a/paddlexde/interpolation/interpolate.py b/paddlexde/interpolation/interpolate.py index c246540..36a2b22 100644 --- a/paddlexde/interpolation/interpolate.py +++ b/paddlexde/interpolation/interpolate.py @@ -6,13 +6,13 @@ class LinearInterpolation(InterpolationBase): """Calculates the linear interpolation to the batch of controls given. Also calculates its derivative.""" - def __init__(self, series, t=None, **kwargs): + def __init__(self, series, t=None): """_summary_ Arguments: ----------- series (tensor): B T D - t (_type_, optional): _description_. Defaults to None. + t (_type_, optional): _description_. Defaults to None. [T] Example: ----------- @@ -29,7 +29,7 @@ def __init__(self, series, t=None, **kwargs): print(chp.derivative(99)) ``` """ - super().__init__(series, t, **kwargs) + super().__init__(series, t) # build cubic hemite spline matrix H indices = [[0, 0, 1], [0, 1, 0]] @@ -44,24 +44,23 @@ def _make_series(self, series, t): Args: series (_type_): [B, T, D] - t (_type_): [B, T] + t (_type_): [T] Note: B can be ignored Returns: _type_: _description_ """ - scale = t[..., 1:] - t[..., :-1] - scale1 = paddle.concat([scale, scale[..., -1:]], axis=-1) - scale2 = paddle.concat([scale[..., :1], scale1[..., :-1]], axis=-1) + scale = t[1:] - t[:-1] + scale1 = paddle.concat([scale, scale[-1:]]) + scale2 = paddle.concat([scale[:1], scale1[:-1]]) series1 = series series2 = paddle.concat([series1[..., 1:, :], series[..., -1:, :]], axis=-2) # [B, T, 2, D] series_r = paddle.stack( - [series1 / scale1.unsqueeze(-1), series2 / scale2.unsqueeze(-1)], - axis=-2, + [series1 / scale1.unsqueeze(-1), series2 / scale2.unsqueeze(-1)], axis=-2 ) return series_r, scale1 @@ -82,19 +81,15 @@ def ps(self, index): """_summary_ Args: - index (_type_): [B, T] + index (_type_): [T] Returns: _type_: _description_ """ - axis_b = paddle.arange(self._batch_size)[:, None, None] - axis_index = index[:, :, None] - axis_d = paddle.arange(self._dims)[None, None, :] - p_tensor = paddle.stack( [ - (self._series_arr[..., 0, :])[axis_b, axis_index, axis_d], - (self._series_arr[..., 1, :])[axis_b, axis_index, axis_d], + paddle.index_select(self._series_arr[..., 0, :], index, axis=-2), + paddle.index_select(self._series_arr[..., 1, :], index, axis=-2), ], axis=-2, ) @@ -141,24 +136,23 @@ def _make_series(self, series, t): Args: series (_type_): [B, T, D] - t (_type_): [B, T] + t (_type_): [T] Note: B can be ignored Returns: _type_: _description_ """ - scale = t[..., 1:] - t[..., :-1] - scale1 = paddle.concat([scale, scale[..., -1:]], axis=-1) - scale2 = paddle.concat([scale[..., :1], scale1[..., :-1]], axis=-1) + scale = t[1:] - t[:-1] + scale1 = paddle.concat([scale, scale[-1:]]) + scale2 = paddle.concat([scale[:1], scale1[:-1]]) series1 = series series2 = paddle.concat([series1[..., 1:, :], series[..., -1:, :]], axis=-2) # [B, T, S, D] series_r = paddle.stack( - [series1 / scale1.unsqueeze(-1), series2 / scale2.unsqueeze(-1)], - axis=-2, + [series1 / scale1.unsqueeze(-1), series2 / scale2.unsqueeze(-1)], axis=-2 ) return series_r, scale1 @@ -168,15 +162,15 @@ def _make_derivative(self, series, t): Args: series (_type_): [B, T, D] - t (_type_): [B, T] + t (_type_): [T] Note: B can be ignored Returns: _type_: _description_ """ - diffs_t = t[..., 1:] - t[..., :-1] - diffs_t1 = paddle.concat([diffs_t, diffs_t[..., -1:]], axis=-1) + diffs_t = t[1:] - t[:-1] + diffs_t1 = paddle.concat([diffs_t, diffs_t[-1:]]) diffs_series = series[..., 1:, :] - series[..., :-1, :] diffs_series = paddle.concat([diffs_series, diffs_series[..., -1:, :]], axis=-2) @@ -197,16 +191,12 @@ def ts(self, t, der=False): return t_tensor # [B, T, 1, M] def ps(self, index): - axis_b = paddle.arange(self._batch_size)[:, None, None] - axis_index = index[:, :, None] - axis_d = paddle.arange(self._dims)[None, None, :] - p_tensor = paddle.stack( [ - (self._series_arr[..., 0, :])[axis_b, axis_index, axis_d], - (self._series_arr[..., 1, :])[axis_b, axis_index, axis_d], - self._derivs[axis_b, axis_index, axis_d], - self._derivs[axis_b, axis_index, axis_d], + paddle.index_select(self._series_arr[..., 0, :], index, axis=-2), + paddle.index_select(self._series_arr[..., 1, :], index, axis=-2), + paddle.index_select(self._derivs, index, axis=-2), + paddle.index_select(self._derivs, index + 1, axis=-2), ], axis=-2, ) @@ -258,13 +248,11 @@ def _make_series(self, series, t): Returns: _type_: _description_ """ - scale = t[..., 3:] - t[..., :-3] - scale1 = paddle.concat( - [scale, scale[..., -1:], scale[..., -1:], scale[..., -1:]], axis=-1 - ) - scale2 = paddle.concat([scale[..., :1], scale1[..., :-1]], axis=-1) - scale3 = paddle.concat([scale[..., :1], scale2[..., :-1]], axis=-1) - scale4 = paddle.concat([scale[..., :1], scale3[..., :-1]], axis=-1) + scale = t[3:] - t[:-3] + scale1 = paddle.concat([scale, scale[-1:], scale[-1:], scale[-1:]]) + scale2 = paddle.concat([scale[:1], scale1[:-1]]) + scale3 = paddle.concat([scale[:1], scale2[:-1]]) + scale4 = paddle.concat([scale[:1], scale3[:-1]]) series1 = series series2 = paddle.concat([series1[..., 1:, :], series[..., -1:, :]], axis=-2) @@ -297,15 +285,12 @@ def ts(self, t, der=False): return t_tensor # [B, T, 1, M] def ps(self, index): - axis_b = paddle.arange(self._batch_size)[:, None, None] - axis_index = index[:, :, None] - axis_d = paddle.arange(self._dims)[None, None, :] p_tensor = paddle.stack( [ - (self._series_arr[..., 0, :])[axis_b, axis_index, axis_d], - (self._series_arr[..., 1, :])[axis_b, axis_index, axis_d], - (self._series_arr[..., 2, :])[axis_b, axis_index, axis_d], - (self._series_arr[..., 3, :])[axis_b, axis_index, axis_d], + paddle.index_select(self._series_arr[..., 0, :], index, axis=-2), + paddle.index_select(self._series_arr[..., 1, :], index, axis=-2), + paddle.index_select(self._series_arr[..., 2, :], index, axis=-2), + paddle.index_select(self._series_arr[..., 3, :], index, axis=-2), ], axis=-2, ) diff --git a/paddlexde/interpolation/interpolate_base.py b/paddlexde/interpolation/interpolate_base.py index 939fc76..5c8c47c 100644 --- a/paddlexde/interpolation/interpolate_base.py +++ b/paddlexde/interpolation/interpolate_base.py @@ -5,26 +5,27 @@ class InterpolationBase(nn.Layer, metaclass=abc.ABCMeta): - def __init__(self, series, t=None, **kwargs): + def __init__(self, series, t=None): """_summary_ Args: series (_type_): [B, T, D] - t (_type_, optional): [B, T]. Defaults to None. + t (_type_, optional): [T]. Defaults to None. the B dim can be ignored """ super().__init__() + self.default_type = paddle.get_default_dtype() if t is None: t = paddle.linspace( 0, series.shape[-2], series.shape[-2] + 1, - dtype=series.dtype, - ).expand(series.shape[:-1]) - series = paddle.cast(series, dtype="float32") - t = paddle.cast(t, dtype="float32") + dtype=self.default_type, + ) # [T] + series = paddle.cast(series, dtype=self.default_type) + t = paddle.cast(t, dtype=self.default_type) series_arr, scale_t = self._make_series(series=series, t=t) derivs = self._make_derivative(series, t) @@ -35,14 +36,6 @@ def __init__(self, series, t=None, **kwargs): self._series = series self._derivs = derivs - if len(series.shape) == 2: - self._seq_len, self._dims = series.shape - self._batch_size = None - elif len(series.shape) == 3: - self._batch_size, self._seq_len, self._dims = series.shape - else: - raise ValueError - @property def grid_points(self): """The time points.""" @@ -57,8 +50,7 @@ def interpolate(self, t, der=False): """Calculates the index of the given time point t in the list of time points. Args: - t (_type_): time point t [B, T] - B can be ignored => [T] + t (_type_): time point t [T] Raises: NotImplementedError: @@ -66,32 +58,14 @@ def interpolate(self, t, der=False): Retuns: The index of the given time point t in the list of time points. """ - t = paddle.to_tensor(t, dtype=self._series.dtype) - if len(t.shape) == 1: - t = t.expand([self._batch_size, t.shape[-1]]) + t = paddle.to_tensor(t, dtype=self.default_type) maxlen = self._series.shape[-2] - 1 - # [B, T], [B, T], [B, T] - t_shape, _t_shape, _scale_t_shape = t.shape, self._t.shape, self._scale_t.shape - - index = [] - norm_t = [] - - t = paddle.reshape(t, shape=[-1, t_shape[-1]]) - _t = paddle.reshape(self._t, shape=[-1, _t_shape[-1]]) - _scale_t = paddle.reshape(self._scale_t, shape=[-1, _scale_t_shape[-1]]) - - for bs in range(self._batch_size): - t_i, _t_i, _scale_t_i = t[bs], _t[bs], _scale_t[bs] - # clamp because t may go outside of [t[0], t[-1]]; this is fine - # will never access the last element of self._t; this is correct behaviour - index_i = (paddle.bucketize(t_i, _t_i) - 1).clip(0, maxlen) # [T] - norm_t_i = (t_i - _t_i[index_i]) / _scale_t_i[index_i] # [T] - index.append(index_i) - norm_t.append(norm_t_i) - - index = paddle.stack(index, axis=0).reshape(t_shape) # [B, T] - norm_t = paddle.stack(norm_t, axis=0).reshape(t_shape) # [B, T] + # clamp because t may go outside of [t[0], t[-1]]; this is fine + # will never access the last element of self._t; this is correct behaviour + index = (paddle.bucketize(t, self._t) - 1).clip(0, maxlen) # [T] + norm_t = t - paddle.index_select(self._t, index) + norm_t /= paddle.index_select(self._scale_t, index) # [T] # [B, T, 1, M], M is 2 for linear, 4 for cubic ts_tensor = self.ts(norm_t, der=der) @@ -116,10 +90,8 @@ def evaluate(self, t): ts_tensor, ps_tensor, index = self.interpolate(t, der=False) # [B, T, 1, D] => [B, T, D] result = (ts_tensor @ self._h.to_dense() @ ps_tensor).squeeze(-2) - axis_b = paddle.arange(self._batch_size)[:, None] - axis_index = index[:, :] - scale = self._scale_t[axis_b, axis_index].unsqueeze(-1) # [B, T, 1] - result *= scale + scale = paddle.index_select(self._scale_t, index) # [T] + result *= scale.unsqueeze(-1) return result def derivative(self, t): @@ -143,6 +115,15 @@ def derivative(self, t): @abc.abstractmethod def _make_series(self, series, t): + """_summary_ + + Args: + series (_type_): [B, T, D] + t (_type_): [B, T] + + Raises: + NotImplementedError: _description_ + """ raise NotImplementedError @abc.abstractmethod diff --git a/paddlexde/solver/base_fixed_solver.py b/paddlexde/solver/base_fixed_solver.py index 4bdb581..ef3fc61 100644 --- a/paddlexde/solver/base_fixed_solver.py +++ b/paddlexde/solver/base_fixed_solver.py @@ -61,6 +61,7 @@ def __init__( self.move = self.xde.move self.fuse = self.xde.fuse + self.on_integrate_step_end = self.xde.on_integrate_step_end @staticmethod def _grid_constructor_from_step_size(step_size): @@ -74,8 +75,8 @@ def _grid_constructor(y0, t): Returns: _type_: _description_ """ - start_time = t[..., 0] - end_time = t[..., -1] + start_time = t[0] + end_time = t[-1] niters = paddle.ceil((end_time - start_time) / step_size + 1).item() t_infer = ( @@ -92,8 +93,8 @@ def step(self, t0, t1, y0): """Propose a step with step size from time t to time next_t, with current state y. - :param t0: [B, 1] - :param t1: [B, 1] + :param t0: [1] + :param t1: [1] :param y0: [B, 1, D] :return: """ @@ -103,7 +104,7 @@ def integrate(self, t_span: paddle.Tensor): """_summary_ Args: - t_span (paddle.Tensor): [batch_size, pred_len] + t_span (paddle.Tensor): [pred_len] Raises: ValueError: _description_ @@ -112,35 +113,34 @@ def integrate(self, t_span: paddle.Tensor): _type_: _description_ """ - batch_size, pred_len = t_span.shape + pred_len = len(t_span) time_grid = self.grid_constructor(self.y0, t_span) # time_grid = t_span - assert paddle.equal_all(time_grid[..., 0], t_span[..., 0]) - assert paddle.equal_all(time_grid[..., -1], t_span[..., -1]) - - # sol solution [pred_len, batch_size, dims] - sol = paddle.empty(shape=[pred_len] + self.y0.shape, dtype=self.y0.dtype) - sol[0] = self.y0 + assert paddle.equal_all(time_grid[0], t_span[0]) + assert paddle.equal_all(time_grid[-1], t_span[-1]) + # sol solution pred_len * [batch_size, 1, dims] + sol = [self.y0] y0 = self.y0 for i in range(1, pred_len): - t0, t1 = time_grid[..., i - 1 : i], time_grid[..., i : i + 1] + t0, t1 = time_grid[i - 1 : i], time_grid[i : i + 1] y1, dy0 = self.step(t0, t1, y0) # while j < pred_len and paddle.greater_equal(t1, t_span[..., j]): if self.interp == "linear": - sol[i] = linear_interp(t0, t1, y0, y1, t_span[..., i : i + 1]) + sol.append(linear_interp(t0, t1, y0, y1, t_span[i : i + 1])) elif self.interp == "cubic": y2, dy1 = self.step(t1, t1, y1) - sol[i] = cubic_hermite_interp( - t0, y0, dy0, t1, y1, dy1, t_span[..., i : i + 1] + sol.append( + cubic_hermite_interp(t0, y0, dy0, t1, y1, dy1, t_span[i : i + 1]) ) else: raise ValueError(f"Unknown interpolation method {self.interp}") y0 = y1 + sol = paddle.concat(sol, axis=-2) return sol def rk4_step_func(self, t0, t1, y0, f0=None): diff --git a/paddlexde/xde/base_dde.py b/paddlexde/xde/base_dde.py index 9163d6a..39632d5 100644 --- a/paddlexde/xde/base_dde.py +++ b/paddlexde/xde/base_dde.py @@ -25,16 +25,17 @@ def __init__( his: paddle.Tensor, his_span: paddle.Tensor, ): - # TODO 此处传入的数据值需要进行改变 - # 如果lags不存在梯度值，则不需要进行初始化和更新，采取固定lags的形式 - # 如果lags存在梯度，证明lags可以进行更新 - # 如果lags为None，则选择动态初始化lags super(BaseDDE, self).__init__(name="DDE", var_nums=1, y0=y0, t_span=t_span) self.func = func self.lags = lags + self.y_lags = HistoryIndex.apply(lags=lags, his=his, his_span=his_span) self.his = his self.his_span = his_span + self.init_y0(y0) + + def init_y0(self, input): + self.y0 = input def handle(self, h, ts): pass @@ -43,12 +44,8 @@ def move(self, t0, dt, y0): # self.init_lags() # input_history = paddle.index_select(self.history, self.lags) - y_lags = HistoryIndex.apply( - xde=self, t0=t0, y0=y0, lags=self.lags, his=self.his, his_span=self.his_span - ) # y_lags [B, T, D] T是选择后的序列长度 - - dy = self.call_func(t0, y0, self.lags, y_lags) + dy = self.func(self.lags, self.y_lags, None, y0) return dy def fuse(self, dy, dt, y0): @@ -59,18 +56,28 @@ def fuse(self, dy, dt, y0): # return dy * dt + y0 def call_func(self, t, y0, lags, y_lags): - y0 = self.unflatten(y0, length=1) + # y0 = self.unflatten(y0, length=1) dy = self.func(t, y0, lags, y_lags) - dy = self.flatten(dy) + # dy = self.flatten(dy) return dy def init_lags(self): # TODO 不同时刻初始化不同lags pass + def flatten(self, input): + return input + + def unflatten(self, input, length): + return input + + def on_integrate_step_end(self, y0=None, y1=None, t0=None, t1=None): + return super().on_integrate_step_end(y0, y1, t0, t1) + class HistoryIndex(autograd.PyLayer): - def forward(ctx, xde: BaseXDE, t0, y0, lags, his, his_span, interp_method="linear"): + @staticmethod + def forward(ctx, lags, his, his_span, interp_method="cubic"): """ 计算给定输入序列的未来值，并返回计算结果。 传入lags, history, @@ -89,53 +96,29 @@ def forward(ctx, xde: BaseXDE, t0, y0, lags, his, his_span, interp_method="linea Raises: NotImplementedError: 如果interp_method不是上述三种情况之一, 将抛出NotImplementedError异常。 """ - if interp_method == "linear": - interp = LinearInterpolation(his, his_span) - elif interp_method == "cubic": - interp = CubicHermiteSpline(his, his_span) - elif interp_method == "bez": - interp = BezierSpline(his, his_span) - else: - raise NotImplementedError - - batch_size, len_t, dims = his.shape # [B, T, D] - axis_b = paddle.arange(batch_size)[:, None, None] - axis_index = lags[:, :, None].astype("int64") - axis_d = paddle.arange(dims)[None, None, :] - y_lags = his[axis_b, axis_index, axis_d] - assert isinstance(y_lags, paddle.Tensor) - y_lags.stop_gradient = False - - ctx.t0 = t0 - ctx.y0 = y0 - ctx.lags = lags - ctx.derivative_lags = interp.derivative(lags) - ctx.xde = xde - ctx.y_lags = y_lags + with paddle.no_grad(): + if interp_method == "linear": + interp = LinearInterpolation(his, his_span) + elif interp_method == "cubic": + interp = CubicHermiteSpline(his, his_span) + elif interp_method == "bez": + interp = BezierSpline(his, his_span) + else: + raise NotImplementedError + + y_lags = interp.evaluate(lags) + + derivative_lags = interp.derivative(lags) + ctx.save_for_backward(derivative_lags) return y_lags + @staticmethod def backward(ctx, grad_y): # 计算history相应的梯度，并提取forward中保存的梯度，用于计算lag的梯度 # 在计算的过程中，无需更新history，仅更新lags即可 - - t0 = ctx.t0 - y0 = ctx.y0 - lags = ctx.lags - derivative_lags = ctx.derivative_lags - xde = ctx.xde - y_lags = ctx.y_lags - y_lags.stop_gradient = False - - with paddle.set_grad_enabled(True): - eval = xde.call_func(t0, y0, lags, y_lags) - - grad_y_lags = paddle.grad( - outputs=[eval], - inputs=[y_lags], - grad_outputs=-grad_y, - allow_unused=True, - retain_graph=True, - )[0] - return None, None, grad_y_lags * derivative_lags * 100000, None, None + (derivative_lags,) = ctx.saved_tensor() + grad = grad_y * derivative_lags + grad = paddle.sum(grad, axis=[0, 1, 3]) + return grad, None, None # return None, grad_y_lags * derivative_lags, None, None, None diff --git a/paddlexde/xde/base_ode.py b/paddlexde/xde/base_ode.py index 212e782..f82c3d3 100644 --- a/paddlexde/xde/base_ode.py +++ b/paddlexde/xde/base_ode.py @@ -22,7 +22,7 @@ def __init__( Args: func (Union[nn.Layer, callable]): _description_ y0 (Union[tuple, paddle.Tensor]): paddle.Tensor shape is (B, T, D), T=1, tuple shape is TODO - t_span (Union[list, paddle.Tensor]): shape is (B, T), T=pred_len + t_span (Union[list, paddle.Tensor]): shape is (T), T=pred_len """ super(BaseODE, self).__init__(name="ODE", var_nums=1, y0=y0, t_span=t_span) """ @@ -36,6 +36,10 @@ def __init__( """ self.func = func + self.init_y0(y0) + + def init_y0(self, y0): + self.y0 = y0 def handle(self, h, ts): pass @@ -45,15 +49,14 @@ def move(self, t0, dt, y0): return dy def fuse(self, dy, dt, y0): - # 测试是否存在振动 - y = dy * dt + y0 - _lambda = 0.001 - return (dy - _lambda * y) * dt + y0 + # 增加负反馈，减少振动，加快收敛 + # y = dy * dt + y0 + # _lambda = 0.001 + # return (dy - _lambda * y) * dt + y0 - # return dy * dt + y0 + # 原始写法 + return dy * dt + y0 def call_func(self, t, y0): - y0 = self.unflatten(y0, length=1) dy = self.func(t, y0) - dy = self.flatten(dy) return dy diff --git a/paddlexde/xde/base_xde.py b/paddlexde/xde/base_xde.py index a33af61..382bf7b 100644 --- a/paddlexde/xde/base_xde.py +++ b/paddlexde/xde/base_xde.py @@ -22,25 +22,31 @@ def __init__( self.var_nums = var_nums # 返回值数量 self.t_span = t_span - self.batch_size, self.pred_len = t_span.shape + self.pred_len = t_span.shape - self.init_y0(y0) # shapes, numels, y0 + # self.init_y0(y0) # shapes, numels, y0 + def method(self): + print(f"current method is {self.name}.") + return self.name + + @abstractmethod def init_y0(self, input): - if isinstance(input, tuple) or isinstance(input, list): - self.shapes = [_tensor.shape for _tensor in input] - self.num_elements = [ - paddle.numel(_tensor) / self.batch_size for _tensor in input - ] - self.y0 = paddle.concat( - [_tensor.reshape([self.batch_size, -1]) for _tensor in input], axis=-1 - ) # [batch_size, -1] - elif isinstance(input, paddle.Tensor): - self.shapes = [input.shape] - self.num_elements = [paddle.numel(input) / self.batch_size] - self.y0 = input.reshape([self.batch_size, -1]) - else: - raise NotImplementedError + # if isinstance(input, tuple) or isinstance(input, list): + # self.shapes = [_tensor.shape for _tensor in input] + # self.num_elements = [ + # paddle.numel(_tensor) / self.batch_size for _tensor in input + # ] + # self.y0 = paddle.concat( + # [_tensor.reshape([self.batch_size, -1]) for _tensor in input], axis=-1 + # ) # [batch_size, -1] + # elif isinstance(input, paddle.Tensor): + # self.shapes = [input.shape] + # self.num_elements = [paddle.numel(input) / self.batch_size] + # self.y0 = input.reshape([self.batch_size, -1]) + # else: + # raise NotImplementedError + raise NotImplementedError @abstractmethod def handle(self, h, ts): @@ -74,48 +80,27 @@ def fuse(self, dy, dt, y0): """ raise NotImplementedError - def format(self, sol): - # [pred_len, batch_size, D] - sol = self.unflatten(sol, self.pred_len) - sol = paddle.transpose(sol, perm=[1, 0, 2]) - return sol - - def method(self): - print(f"current method is {self.name}.") - return self.name - def unflatten(self, input, length): - batch_size = input.shape[0] # 默认第一维是batch_size - - length_shape = [] if length == 1 else [length] - if len(self.shapes) == 1: - return input.reshape(length_shape + self.shapes[0]) - - output = [] - total = 0 - - # TODO:此处需要验证 - for shape, num_ele in zip(self.shapes, self.num_elements): - next_total = total + num_ele - output.append( - input[..., total:next_total].reshape( - [batch_size] + length_shape + shape - ) - ) - total = next_total - return tuple(output) + raise NotImplementedError def flatten(self, input): - if isinstance(input, tuple) or isinstance(input, list): - output = paddle.concat( - [_tensor.reshape([self.batch_size, -1]) for _tensor in input] - ) - elif isinstance(input, paddle.Tensor): - output = input.reshape([self.batch_size, -1]) - else: - raise NotImplementedError - - return output + raise NotImplementedError + + # def format(self, sol): + # """_summary_ + + # Args: + # sol (list): pred_len * [B, 1, D] + + # Returns: + # _type_: _description_ + # """ + # # [pred_len, batch_size, D] + # sol = paddle.concat(sol, axis=-2) + # return sol + + def on_integrate_step_end(self, y0=None, y1=None, t0=None, t1=None): + pass @abstractmethod def call_func(self, **kwargs): diff --git a/requirements.txt b/requirements.txt index ec4f643..6370c35 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,4 +1,7 @@ -numpy +paddle boltons trampoline -matplotlib \ No newline at end of file +matplotlib +scikit-learn +pandas +visualdl \ No newline at end of file diff --git a/tests/functional/test_fixed_solver.py b/tests/functional/test_fixed_solver.py index ecf23fd..1b8fdd5 100644 --- a/tests/functional/test_fixed_solver.py +++ b/tests/functional/test_fixed_solver.py @@ -15,10 +15,10 @@ def setUpClass(self): ode="constant", reverse=False, ) - self.f = f - self.y0 = y0 - self.t = t - self.sol = sol + self.f = f # paddle model + self.y0 = y0 # [B, D] + self.t = t # [T] + self.sol = sol # [B, T, D] self.xdeints = [odeint, odeint_adjoint] self.solvers = [Euler, Midpoint, RK4, AdamsBashforthMoulton] diff --git a/tests/utils/interpolation/__init__.py b/tests/interpolation/__init__.py similarity index 100% rename from tests/utils/interpolation/__init__.py rename to tests/interpolation/__init__.py diff --git a/tests/utils/interpolation/test_interpolation.py b/tests/interpolation/test_interpolation.py similarity index 91% rename from tests/utils/interpolation/test_interpolation.py rename to tests/interpolation/test_interpolation.py index 5575a09..710ba9e 100644 --- a/tests/utils/interpolation/test_interpolation.py +++ b/tests/interpolation/test_interpolation.py @@ -21,10 +21,10 @@ def setUpClass(self): ], axis=-1, ).unsqueeze(0) - # [B, T] - self.t = paddle.arange(0, 2000, 1).unsqueeze(0) - # [B, T] - self.t_eval = paddle.to_tensor([[21.12]]) + # [T] + self.t = paddle.arange(0, 2000, 1) + # [T] + self.t_eval = paddle.to_tensor([21.12]) # [B, T, D] self.val_tgt = paddle.to_tensor([21.12 * 0.5, 0]).unsqueeze(0).unsqueeze(0) # [B, T, D] @@ -58,9 +58,9 @@ def setUpClass(self): axis=-1, ).unsqueeze(0) self.series = paddle.sin(self.series) - self.t = paddle.arange(0, 20, 0.01).unsqueeze(0) # [B, T] + self.t = paddle.arange(0, 20, 0.01) # [T] - self.t_eval = paddle.to_tensor([[16.5]]) # [B, T] + self.t_eval = paddle.to_tensor([16.5]) # [T] self.val_tgt = paddle.sin( paddle.to_tensor([16.5, 0]).unsqueeze(0).unsqueeze(0) ) # [B, T, D] @@ -76,7 +76,7 @@ def test_LinearInterpolation(self): def test_CubicHermiteSpline(self): interp = CubicHermiteSpline(self.series, self.t) - assert paddle.allclose(self.val_tgt, interp.evaluate(self.t_eval), rtol=5e-2) + assert paddle.allclose(self.val_tgt, interp.evaluate(self.t_eval), rtol=1e-5) assert paddle.allclose(self.tgt_deri, interp.derivative(self.t_eval), rtol=1e-2) def test_BezierSpline(self): diff --git a/tests/testing_utils.py b/tests/testing_utils.py index 828c91b..0c9e2e3 100644 --- a/tests/testing_utils.py +++ b/tests/testing_utils.py @@ -95,7 +95,7 @@ def _flip(x, dim): t_points = _flip(t_points, 0).clone().detach() sol = _flip(sol, 0).clone().detach() - return f, paddle.assign(sol[0]), t_points, sol + return f, paddle.assign(sol[0]).unsqueeze(0), t_points, sol.unsqueeze(0) if __name__ == "__main__": diff --git a/tests/utils/__init__.py b/tests/utils/__init__.py deleted file mode 100644 index 92aeea2..0000000 --- a/tests/utils/__init__.py +++ /dev/null @@ -1 +0,0 @@ -from .interpolation import *