This commit is contained in:
Alex Blank
2025-05-19 13:59:16 +02:00
parent 426f4d6963
commit c6defa2065
196 changed files with 18625 additions and 1 deletions
@@ -0,0 +1,2 @@
* text=auto
*.sh text eol=lf
@@ -0,0 +1,465 @@
# Created by https://www.toptal.com/developers/gitignore/api/emacs,jetbrains+all,jupyternotebooks,linux,macos,microsoftoffice,notepadpp,python,tortoisegit,vim,visualstudiocode,windows
# Edit at https://www.toptal.com/developers/gitignore?templates=emacs,jetbrains+all,jupyternotebooks,linux,macos,microsoftoffice,notepadpp,python,tortoisegit,vim,visualstudiocode,windows
### Emacs ###
# -*- mode: gitignore; -*-
*~
\#*\#
/.emacs.desktop
/.emacs.desktop.lock
*.elc
auto-save-list
tramp
.\#*
# Org-mode
.org-id-locations
*_archive
# flymake-mode
*_flymake.*
# eshell files
/eshell/history
/eshell/lastdir
# elpa packages
/elpa/
# reftex files
*.rel
# AUCTeX auto folder
/auto/
# cask packages
.cask/
dist/
# Flycheck
flycheck_*.el
# server auth directory
/server/
# projectiles files
.projectile
# directory configuration
.dir-locals.el
# network security
/network-security.data
### JetBrains+all ###
# Covers JetBrains IDEs: IntelliJ, RubyMine, PhpStorm, AppCode, PyCharm, CLion, Android Studio, WebStorm and Rider
# Reference: https://intellij-support.jetbrains.com/hc/en-us/articles/206544839
# User-specific stuff
.idea/**/workspace.xml
.idea/**/tasks.xml
.idea/**/usage.statistics.xml
.idea/**/dictionaries
.idea/**/shelf
# AWS User-specific
.idea/**/aws.xml
# Generated files
.idea/**/contentModel.xml
# Sensitive or high-churn files
.idea/**/dataSources/
.idea/**/dataSources.ids
.idea/**/dataSources.local.xml
.idea/**/sqlDataSources.xml
.idea/**/dynamic.xml
.idea/**/uiDesigner.xml
.idea/**/dbnavigator.xml
# Gradle
.idea/**/gradle.xml
.idea/**/libraries
# Gradle and Maven with auto-import
# When using Gradle or Maven with auto-import, you should exclude module files,
# since they will be recreated, and may cause churn. Uncomment if using
# auto-import.
# .idea/artifacts
# .idea/compiler.xml
# .idea/jarRepositories.xml
# .idea/modules.xml
# .idea/*.iml
# .idea/modules
# *.iml
# *.ipr
# CMake
cmake-build-*/
# Mongo Explorer plugin
.idea/**/mongoSettings.xml
# File-based project format
*.iws
# IntelliJ
out/
# mpeltonen/sbt-idea plugin
.idea_modules/
# JIRA plugin
atlassian-ide-plugin.xml
# Cursive Clojure plugin
.idea/replstate.xml
# SonarLint plugin
.idea/sonarlint/
# Crashlytics plugin (for Android Studio and IntelliJ)
com_crashlytics_export_strings.xml
crashlytics.properties
crashlytics-build.properties
fabric.properties
# Editor-based Rest Client
.idea/httpRequests
# Android studio 3.1+ serialized cache file
.idea/caches/build_file_checksums.ser
### JetBrains+all Patch ###
# Ignore everything but code style settings and run configurations
# that are supposed to be shared within teams.
.idea/*
!.idea/codeStyles
!.idea/runConfigurations
### JupyterNotebooks ###
# gitignore template for Jupyter Notebooks
# website: http://jupyter.org/
.ipynb_checkpoints
*/.ipynb_checkpoints/*
# IPython
profile_default/
ipython_config.py
# Remove previous ipynb_checkpoints
# git rm -r .ipynb_checkpoints/
### Linux ###
# temporary files which can be created if a process still has a handle open of a deleted file
.fuse_hidden*
# KDE directory preferences
.directory
# Linux trash folder which might appear on any partition or disk
.Trash-*
# .nfs files are created when an open file is removed but is still being accessed
.nfs*
### macOS ###
# General
.DS_Store
.AppleDouble
.LSOverride
# Icon must end with two \r
Icon
# Thumbnails
._*
# Files that might appear in the root of a volume
.DocumentRevisions-V100
.fseventsd
.Spotlight-V100
.TemporaryItems
.Trashes
.VolumeIcon.icns
.com.apple.timemachine.donotpresent
# Directories potentially created on remote AFP share
.AppleDB
.AppleDesktop
Network Trash Folder
Temporary Items
.apdisk
### macOS Patch ###
# iCloud generated files
*.icloud
### MicrosoftOffice ###
*.tmp
# Word temporary
~$*.doc*
# Word Auto Backup File
Backup of *.doc*
# Excel temporary
~$*.xls*
# Excel Backup File
*.xlk
# PowerPoint temporary
~$*.ppt*
# Visio autosave temporary files
*.~vsd*
### NotepadPP ###
# Notepad++ backups #
*.bak
### Python ###
# Byte-compiled / optimized / DLL files
__pycache__/
*.py[cod]
*$py.class
# C extensions
*.so
# Distribution / packaging
.Python
build/
develop-eggs/
downloads/
eggs/
.eggs/
lib/
lib64/
parts/
sdist/
var/
wheels/
share/python-wheels/
*.egg-info/
.installed.cfg
*.egg
MANIFEST
# PyInstaller
# Usually these files are written by a python script from a template
# before PyInstaller builds the exe, so as to inject date/other infos into it.
*.manifest
*.spec
# Installer logs
pip-log.txt
pip-delete-this-directory.txt
# Unit test / coverage reports
htmlcov/
.tox/
.nox/
.coverage
.coverage.*
.cache
nosetests.xml
coverage.xml
*.cover
*.py,cover
.hypothesis/
.pytest_cache/
cover/
# Translations
*.mo
*.pot
# Django stuff:
*.log
local_settings.py
db.sqlite3
db.sqlite3-journal
# Flask stuff:
instance/
.webassets-cache
# Scrapy stuff:
.scrapy
# Sphinx documentation
docs/_build/
# PyBuilder
.pybuilder/
target/
# Jupyter Notebook
# IPython
# pyenv
# For a library or package, you might want to ignore these files since the code is
# intended to run in multiple environments; otherwise, check them in:
# .python-version
# pipenv
# According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
# However, in case of collaboration, if having platform-specific dependencies or dependencies
# having no cross-platform support, pipenv may install dependencies that don't work, or not
# install all needed dependencies.
#Pipfile.lock
# poetry
# Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control.
# This is especially recommended for binary packages to ensure reproducibility, and is more
# commonly ignored for libraries.
# https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control
#poetry.lock
# pdm
# Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control.
#pdm.lock
# pdm stores project-wide configurations in .pdm.toml, but it is recommended to not include it
# in version control.
# https://pdm.fming.dev/#use-with-ide
.pdm.toml
# PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm
__pypackages__/
# Celery stuff
celerybeat-schedule
celerybeat.pid
# SageMath parsed files
*.sage.py
# Environments
.env
.venv
env/
venv/
ENV/
env.bak/
venv.bak/
# Spyder project settings
.spyderproject
.spyproject
# Rope project settings
.ropeproject
# mkdocs documentation
/site
# mypy
.mypy_cache/
.dmypy.json
dmypy.json
# Pyre type checker
.pyre/
# pytype static type analyzer
.pytype/
# Cython debug symbols
cython_debug/
# PyCharm
# JetBrains specific template is maintained in a separate JetBrains.gitignore that can
# be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore
# and can be added to the global gitignore or merged into this file. For a more nuclear
# option (not recommended) you can uncomment the following to ignore the entire idea folder.
#.idea/
### Python Patch ###
# Poetry local configuration file - https://python-poetry.org/docs/configuration/#local-configuration
poetry.toml
# ruff
.ruff_cache/
# LSP config files
pyrightconfig.json
### TortoiseGit ###
# Project-level settings
/.tgitconfig
### Vim ###
# Swap
[._]*.s[a-v][a-z]
!*.svg # comment out if you don't need vector files
[._]*.sw[a-p]
[._]s[a-rt-v][a-z]
[._]ss[a-gi-z]
[._]sw[a-p]
# Session
Session.vim
Sessionx.vim
# Temporary
.netrwhist
# Auto-generated tag files
tags
# Persistent undo
[._]*.un~
### VisualStudioCode ###
.vscode/*
!.vscode/settings.json
!.vscode/tasks.json
!.vscode/launch.json
!.vscode/extensions.json
!.vscode/*.code-snippets
# Local History for Visual Studio Code
.history/
# Built Visual Studio Code Extensions
*.vsix
### VisualStudioCode Patch ###
# Ignore all local history of files
.history
.ionide
### Windows ###
# Windows thumbnail cache files
Thumbs.db
Thumbs.db:encryptable
ehthumbs.db
ehthumbs_vista.db
# Dump file
*.stackdump
# Folder config file
[Dd]esktop.ini
# Recycle Bin used on file shares
$RECYCLE.BIN/
# Windows Installer files
*.cab
*.msi
*.msix
*.msm
*.msp
# Windows shortcuts
*.lnk
+201
View File
@@ -0,0 +1,201 @@
Apache License
Version 2.0, January 2004
http://www.apache.org/licenses/
TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
1. Definitions.
"License" shall mean the terms and conditions for use, reproduction,
and distribution as defined by Sections 1 through 9 of this document.
"Licensor" shall mean the copyright owner or entity authorized by
the copyright owner that is granting the License.
"Legal Entity" shall mean the union of the acting entity and all
other entities that control, are controlled by, or are under common
control with that entity. For the purposes of this definition,
"control" means (i) the power, direct or indirect, to cause the
direction or management of such entity, whether by contract or
otherwise, or (ii) ownership of fifty percent (50%) or more of the
outstanding shares, or (iii) beneficial ownership of such entity.
"You" (or "Your") shall mean an individual or Legal Entity
exercising permissions granted by this License.
"Source" form shall mean the preferred form for making modifications,
including but not limited to software source code, documentation
source, and configuration files.
"Object" form shall mean any form resulting from mechanical
transformation or translation of a Source form, including but
not limited to compiled object code, generated documentation,
and conversions to other media types.
"Work" shall mean the work of authorship, whether in Source or
Object form, made available under the License, as indicated by a
copyright notice that is included in or attached to the work
(an example is provided in the Appendix below).
"Derivative Works" shall mean any work, whether in Source or Object
form, that is based on (or derived from) the Work and for which the
editorial revisions, annotations, elaborations, or other modifications
represent, as a whole, an original work of authorship. For the purposes
of this License, Derivative Works shall not include works that remain
separable from, or merely link (or bind by name) to the interfaces of,
the Work and Derivative Works thereof.
"Contribution" shall mean any work of authorship, including
the original version of the Work and any modifications or additions
to that Work or Derivative Works thereof, that is intentionally
submitted to Licensor for inclusion in the Work by the copyright owner
or by an individual or Legal Entity authorized to submit on behalf of
the copyright owner. For the purposes of this definition, "submitted"
means any form of electronic, verbal, or written communication sent
to the Licensor or its representatives, including but not limited to
communication on electronic mailing lists, source code control systems,
and issue tracking systems that are managed by, or on behalf of, the
Licensor for the purpose of discussing and improving the Work, but
excluding communication that is conspicuously marked or otherwise
designated in writing by the copyright owner as "Not a Contribution."
"Contributor" shall mean Licensor and any individual or Legal Entity
on behalf of whom a Contribution has been received by Licensor and
subsequently incorporated within the Work.
2. Grant of Copyright License. Subject to the terms and conditions of
this License, each Contributor hereby grants to You a perpetual,
worldwide, non-exclusive, no-charge, royalty-free, irrevocable
copyright license to reproduce, prepare Derivative Works of,
publicly display, publicly perform, sublicense, and distribute the
Work and such Derivative Works in Source or Object form.
3. Grant of Patent License. Subject to the terms and conditions of
this License, each Contributor hereby grants to You a perpetual,
worldwide, non-exclusive, no-charge, royalty-free, irrevocable
(except as stated in this section) patent license to make, have made,
use, offer to sell, sell, import, and otherwise transfer the Work,
where such license applies only to those patent claims licensable
by such Contributor that are necessarily infringed by their
Contribution(s) alone or by combination of their Contribution(s)
with the Work to which such Contribution(s) was submitted. If You
institute patent litigation against any entity (including a
cross-claim or counterclaim in a lawsuit) alleging that the Work
or a Contribution incorporated within the Work constitutes direct
or contributory patent infringement, then any patent licenses
granted to You under this License for that Work shall terminate
as of the date such litigation is filed.
4. Redistribution. You may reproduce and distribute copies of the
Work or Derivative Works thereof in any medium, with or without
modifications, and in Source or Object form, provided that You
meet the following conditions:
(a) You must give any other recipients of the Work or
Derivative Works a copy of this License; and
(b) You must cause any modified files to carry prominent notices
stating that You changed the files; and
(c) You must retain, in the Source form of any Derivative Works
that You distribute, all copyright, patent, trademark, and
attribution notices from the Source form of the Work,
excluding those notices that do not pertain to any part of
the Derivative Works; and
(d) If the Work includes a "NOTICE" text file as part of its
distribution, then any Derivative Works that You distribute must
include a readable copy of the attribution notices contained
within such NOTICE file, excluding those notices that do not
pertain to any part of the Derivative Works, in at least one
of the following places: within a NOTICE text file distributed
as part of the Derivative Works; within the Source form or
documentation, if provided along with the Derivative Works; or,
within a display generated by the Derivative Works, if and
wherever such third-party notices normally appear. The contents
of the NOTICE file are for informational purposes only and
do not modify the License. You may add Your own attribution
notices within Derivative Works that You distribute, alongside
or as an addendum to the NOTICE text from the Work, provided
that such additional attribution notices cannot be construed
as modifying the License.
You may add Your own copyright statement to Your modifications and
may provide additional or different license terms and conditions
for use, reproduction, or distribution of Your modifications, or
for any such Derivative Works as a whole, provided Your use,
reproduction, and distribution of the Work otherwise complies with
the conditions stated in this License.
5. Submission of Contributions. Unless You explicitly state otherwise,
any Contribution intentionally submitted for inclusion in the Work
by You to the Licensor shall be under the terms and conditions of
this License, without any additional terms or conditions.
Notwithstanding the above, nothing herein shall supersede or modify
the terms of any separate license agreement you may have executed
with Licensor regarding such Contributions.
6. Trademarks. This License does not grant permission to use the trade
names, trademarks, service marks, or product names of the Licensor,
except as required for reasonable and customary use in describing the
origin of the Work and reproducing the content of the NOTICE file.
7. Disclaimer of Warranty. Unless required by applicable law or
agreed to in writing, Licensor provides the Work (and each
Contributor provides its Contributions) on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
implied, including, without limitation, any warranties or conditions
of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
PARTICULAR PURPOSE. You are solely responsible for determining the
appropriateness of using or redistributing the Work and assume any
risks associated with Your exercise of permissions under this License.
8. Limitation of Liability. In no event and under no legal theory,
whether in tort (including negligence), contract, or otherwise,
unless required by applicable law (such as deliberate and grossly
negligent acts) or agreed to in writing, shall any Contributor be
liable to You for damages, including any direct, indirect, special,
incidental, or consequential damages of any character arising as a
result of this License or out of the use or inability to use the
Work (including but not limited to damages for loss of goodwill,
work stoppage, computer failure or malfunction, or any and all
other commercial damages or losses), even if such Contributor
has been advised of the possibility of such damages.
9. Accepting Warranty or Additional Liability. While redistributing
the Work or Derivative Works thereof, You may choose to offer,
and charge a fee for, acceptance of support, warranty, indemnity,
or other liability obligations and/or rights consistent with this
License. However, in accepting such obligations, You may act only
on Your own behalf and on Your sole responsibility, not on behalf
of any other Contributor, and only if You agree to indemnify,
defend, and hold each Contributor harmless for any liability
incurred by, or claims asserted against, such Contributor by reason
of your accepting any such warranty or additional liability.
END OF TERMS AND CONDITIONS
APPENDIX: How to apply the Apache License to your work.
To apply the Apache License to your work, attach the following
boilerplate notice, with the fields enclosed by brackets "[]"
replaced with your own identifying information. (Don't include
the brackets!) The text should be enclosed in the appropriate
comment syntax for the file format. We also recommend that a
file or class name and description of purpose be included on the
same "printed page" as the copyright notice for easier
identification within third-party archives.
Copyright [yyyy] [name of copyright owner]
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
@@ -0,0 +1,14 @@
# from . import timefeatures
# from . import precip
# from . import metr_la
# from .datamodule import DataModule
# from .csv_dataset import *
# from . import woodside
# from . import dow
# from . import synthetic
# from . import image_completion
# from . import copy_task
# from . import cont_copy_task
# from . import m4
# from . import wiki
# from . import monash
@@ -0,0 +1,196 @@
import numpy as np
import pandas as pd
import torch
from torch import nn
import sys
from src.data.datamodule import DataLoaders
from src.data.pred_dataset import *
DSETS = ['ettm1', 'ettm2', 'etth1', 'etth2', 'electricity',
'traffic', 'illness', 'weather', 'exchange'
]
def get_dls(params):
assert params.dset in DSETS, f"Unrecognized dset (`{params.dset}`). Options include: {DSETS}"
if not hasattr(params,'use_time_features'): params.use_time_features = False
if params.dset == 'ettm1':
root_path = '/data/datasets/public/ETDataset/ETT-small/'
size = [params.context_points, 0, params.target_points]
dls = DataLoaders(
datasetCls=Dataset_ETT_minute,
dataset_kwargs={
'root_path': root_path,
'data_path': 'ETTm1.csv',
'features': params.features,
'scale': True,
'size': size,
'use_time_features': params.use_time_features
},
batch_size=params.batch_size,
workers=params.num_workers,
)
elif params.dset == 'ettm2':
root_path = '/data/datasets/public/ETDataset/ETT-small/'
size = [params.context_points, 0, params.target_points]
dls = DataLoaders(
datasetCls=Dataset_ETT_minute,
dataset_kwargs={
'root_path': root_path,
'data_path': 'ETTm2.csv',
'features': params.features,
'scale': True,
'size': size,
'use_time_features': params.use_time_features
},
batch_size=params.batch_size,
workers=params.num_workers,
)
elif params.dset == 'etth1':
root_path = '/data/datasets/public/ETDataset/ETT-small/'
size = [params.context_points, 0, params.target_points]
dls = DataLoaders(
datasetCls=Dataset_ETT_hour,
dataset_kwargs={
'root_path': root_path,
'data_path': 'ETTh1.csv',
'features': params.features,
'scale': True,
'size': size,
'use_time_features': params.use_time_features
},
batch_size=params.batch_size,
workers=params.num_workers,
)
elif params.dset == 'etth2':
root_path = '/data/datasets/public/ETDataset/ETT-small/'
size = [params.context_points, 0, params.target_points]
dls = DataLoaders(
datasetCls=Dataset_ETT_hour,
dataset_kwargs={
'root_path': root_path,
'data_path': 'ETTh2.csv',
'features': params.features,
'scale': True,
'size': size,
'use_time_features': params.use_time_features
},
batch_size=params.batch_size,
workers=params.num_workers,
)
elif params.dset == 'electricity':
root_path = '/data/datasets/public/electricity/'
size = [params.context_points, 0, params.target_points]
dls = DataLoaders(
datasetCls=Dataset_Custom,
dataset_kwargs={
'root_path': root_path,
'data_path': 'electricity.csv',
'features': params.features,
'scale': True,
'size': size,
'use_time_features': params.use_time_features
},
batch_size=params.batch_size,
workers=params.num_workers,
)
elif params.dset == 'traffic':
root_path = '/data/datasets/public/traffic/'
size = [params.context_points, 0, params.target_points]
dls = DataLoaders(
datasetCls=Dataset_Custom,
dataset_kwargs={
'root_path': root_path,
'data_path': 'traffic.csv',
'features': params.features,
'scale': True,
'size': size,
'use_time_features': params.use_time_features
},
batch_size=params.batch_size,
workers=params.num_workers,
)
elif params.dset == 'weather':
root_path = '/data/datasets/public/weather/'
size = [params.context_points, 0, params.target_points]
dls = DataLoaders(
datasetCls=Dataset_Custom,
dataset_kwargs={
'root_path': root_path,
'data_path': 'weather.csv',
'features': params.features,
'scale': True,
'size': size,
'use_time_features': params.use_time_features
},
batch_size=params.batch_size,
workers=params.num_workers,
)
elif params.dset == 'illness':
root_path = '/data/datasets/public/illness/'
size = [params.context_points, 0, params.target_points]
dls = DataLoaders(
datasetCls=Dataset_Custom,
dataset_kwargs={
'root_path': root_path,
'data_path': 'national_illness.csv',
'features': params.features,
'scale': True,
'size': size,
'use_time_features': params.use_time_features
},
batch_size=params.batch_size,
workers=params.num_workers,
)
elif params.dset == 'exchange':
root_path = '/data/datasets/public/exchange_rate/'
size = [params.context_points, 0, params.target_points]
dls = DataLoaders(
datasetCls=Dataset_Custom,
dataset_kwargs={
'root_path': root_path,
'data_path': 'exchange_rate.csv',
'features': params.features,
'scale': True,
'size': size,
'use_time_features': params.use_time_features
},
batch_size=params.batch_size,
workers=params.num_workers,
)
# dataset is assume to have dimension len x nvars
dls.vars, dls.len = dls.train.dataset[0][0].shape[1], params.context_points
dls.c = dls.train.dataset[0][1].shape[0]
return dls
if __name__ == "__main__":
class Params:
dset= 'etth2'
context_points= 384
target_points= 96
batch_size= 64
num_workers= 8
with_ray= False
features='M'
params = Params
dls = get_dls(params)
for i, batch in enumerate(dls.valid):
print(i, len(batch), batch[0].shape, batch[1].shape)
breakpoint()
@@ -0,0 +1,237 @@
import numpy as np
import pandas as pd
import os
import torch
from torch import nn
from src.models.patchTST import PatchTST
from src.learner import Learner, transfer_weights
from src.callback.core import *
from src.callback.tracking import *
from src.callback.patch_mask import *
from src.callback.transforms import *
from src.metrics import *
from src.basics import set_device
from datautils import *
import argparse
parser = argparse.ArgumentParser()
# Pretraining and Finetuning
parser.add_argument('--is_finetune', type=int, default=0, help='do finetuning or not')
parser.add_argument('--is_linear_probe', type=int, default=0, help='if linear_probe: only finetune the last layer')
# Dataset and dataloader
parser.add_argument('--dset_finetune', type=str, default='etth1', help='dataset name')
parser.add_argument('--context_points', type=int, default=512, help='sequence length')
parser.add_argument('--target_points', type=int, default=96, help='forecast horizon')
parser.add_argument('--batch_size', type=int, default=64, help='batch size')
parser.add_argument('--num_workers', type=int, default=0, help='number of workers for DataLoader')
parser.add_argument('--scaler', type=str, default='standard', help='scale the input data')
parser.add_argument('--features', type=str, default='M', help='for multivariate model or univariate model')
# Patch
parser.add_argument('--patch_len', type=int, default=12, help='patch length')
parser.add_argument('--stride', type=int, default=12, help='stride between patch')
# RevIN
parser.add_argument('--revin', type=int, default=1, help='reversible instance normalization')
# Model args
parser.add_argument('--n_layers', type=int, default=3, help='number of Transformer layers')
parser.add_argument('--n_heads', type=int, default=16, help='number of Transformer heads')
parser.add_argument('--d_model', type=int, default=128, help='Transformer d_model')
parser.add_argument('--d_ff', type=int, default=256, help='Tranformer MLP dimension')
parser.add_argument('--dropout', type=float, default=0.2, help='Transformer dropout')
parser.add_argument('--head_dropout', type=float, default=0.2, help='head dropout')
# Optimization args
parser.add_argument('--n_epochs_finetune', type=int, default=20, help='number of finetuning epochs')
parser.add_argument('--lr', type=float, default=1e-4, help='learning rate')
# Pretrained model name
parser.add_argument('--pretrained_model', type=str, default=None, help='pretrained model name')
# model id to keep track of the number of models saved
parser.add_argument('--finetuned_model_id', type=int, default=1, help='id of the saved finetuned model')
parser.add_argument('--model_type', type=str, default='based_model', help='for multivariate model or univariate model')
args = parser.parse_args()
print('args:', args)
args.save_path = 'saved_models/' + args.dset_finetune + '/masked_patchtst/' + args.model_type + '/'
if not os.path.exists(args.save_path): os.makedirs(args.save_path)
# args.save_finetuned_model = '_cw'+str(args.context_points)+'_tw'+str(args.target_points) + '_patch'+str(args.patch_len) + '_stride'+str(args.stride) + '_epochs-finetune' + str(args.n_epochs_finetune) + '_mask' + str(args.mask_ratio) + '_model' + str(args.finetuned_model_id)
suffix_name = '_cw'+str(args.context_points)+'_tw'+str(args.target_points) + '_patch'+str(args.patch_len) + '_stride'+str(args.stride) + '_epochs-finetune' + str(args.n_epochs_finetune) + '_model' + str(args.finetuned_model_id)
if args.is_finetune: args.save_finetuned_model = args.dset_finetune+'_patchtst_finetuned'+suffix_name
elif args.is_linear_probe: args.save_finetuned_model = args.dset_finetune+'_patchtst_linear-probe'+suffix_name
else: args.save_finetuned_model = args.dset_finetune+'_patchtst_finetuned'+suffix_name
# get available GPU devide
set_device()
def get_model(c_in, args, head_type, weight_path=None):
"""
c_in: number of variables
"""
# get number of patches
num_patch = (max(args.context_points, args.patch_len)-args.patch_len) // args.stride + 1
print('number of patches:', num_patch)
# get model
model = PatchTST(c_in=c_in,
target_dim=args.target_points,
patch_len=args.patch_len,
stride=args.stride,
num_patch=num_patch,
n_layers=args.n_layers,
n_heads=args.n_heads,
d_model=args.d_model,
shared_embedding=True,
d_ff=args.d_ff,
dropout=args.dropout,
head_dropout=args.head_dropout,
act='relu',
head_type=head_type,
res_attention=False
)
if weight_path: model = transfer_weights(weight_path, model)
# print out the model size
print('number of model params', sum(p.numel() for p in model.parameters() if p.requires_grad))
return model
def find_lr(head_type):
# get dataloader
dls = get_dls(args)
model = get_model(dls.vars, args, head_type)
# transfer weight
# weight_path = args.save_path + args.pretrained_model + '.pth'
model = transfer_weights(args.pretrained_model, model)
# get loss
loss_func = torch.nn.MSELoss(reduction='mean')
# get callbacks
cbs = [RevInCB(dls.vars)] if args.revin else []
cbs += [PatchCB(patch_len=args.patch_len, stride=args.stride)]
# define learner
learn = Learner(dls, model,
loss_func,
lr=args.lr,
cbs=cbs,
)
# fit the data to the model
suggested_lr = learn.lr_finder()
print('suggested_lr', suggested_lr)
return suggested_lr
def save_recorders(learn):
train_loss = learn.recorder['train_loss']
valid_loss = learn.recorder['valid_loss']
df = pd.DataFrame(data={'train_loss': train_loss, 'valid_loss': valid_loss})
df.to_csv(args.save_path + args.save_finetuned_model + '_losses.csv', float_format='%.6f', index=False)
def finetune_func(lr=args.lr):
print('end-to-end finetuning')
# get dataloader
dls = get_dls(args)
# get model
model = get_model(dls.vars, args, head_type='prediction')
# transfer weight
# weight_path = args.pretrained_model + '.pth'
model = transfer_weights(args.pretrained_model, model)
# get loss
loss_func = torch.nn.MSELoss(reduction='mean')
# get callbacks
cbs = [RevInCB(dls.vars, denorm=True)] if args.revin else []
cbs += [
PatchCB(patch_len=args.patch_len, stride=args.stride),
SaveModelCB(monitor='valid_loss', fname=args.save_finetuned_model, path=args.save_path)
]
# define learner
learn = Learner(dls, model,
loss_func,
lr=lr,
cbs=cbs,
metrics=[mse]
)
# fit the data to the model
#learn.fit_one_cycle(n_epochs=args.n_epochs_finetune, lr_max=lr)
learn.fine_tune(n_epochs=args.n_epochs_finetune, base_lr=lr, freeze_epochs=10)
save_recorders(learn)
def linear_probe_func(lr=args.lr):
print('linear probing')
# get dataloader
dls = get_dls(args)
# get model
model = get_model(dls.vars, args, head_type='prediction')
# transfer weight
# weight_path = args.save_path + args.pretrained_model + '.pth'
model = transfer_weights(args.pretrained_model, model)
# get loss
loss_func = torch.nn.MSELoss(reduction='mean')
# get callbacks
cbs = [RevInCB(dls.vars, denorm=True)] if args.revin else []
cbs += [
PatchCB(patch_len=args.patch_len, stride=args.stride),
SaveModelCB(monitor='valid_loss', fname=args.save_finetuned_model, path=args.save_path)
]
# define learner
learn = Learner(dls, model,
loss_func,
lr=lr,
cbs=cbs,
metrics=[mse]
)
# fit the data to the model
learn.linear_probe(n_epochs=args.n_epochs_finetune, base_lr=lr)
save_recorders(learn)
def test_func(weight_path):
# get dataloader
dls = get_dls(args)
model = get_model(dls.vars, args, head_type='prediction').to('cuda')
# get callbacks
cbs = [RevInCB(dls.vars, denorm=True)] if args.revin else []
cbs += [PatchCB(patch_len=args.patch_len, stride=args.stride)]
learn = Learner(dls, model,cbs=cbs)
out = learn.test(dls.test, weight_path=weight_path+'.pth', scores=[mse,mae]) # out: a list of [pred, targ, score]
print('score:', out[2])
# save results
pd.DataFrame(np.array(out[2]).reshape(1,-1), columns=['mse','mae']).to_csv(args.save_path + args.save_finetuned_model + '_acc.csv', float_format='%.6f', index=False)
return out
if __name__ == '__main__':
if args.is_finetune:
args.dset = args.dset_finetune
# Finetune
suggested_lr = find_lr(head_type='prediction')
finetune_func(suggested_lr)
print('finetune completed')
# Test
out = test_func(args.save_path+args.save_finetuned_model)
print('----------- Complete! -----------')
elif args.is_linear_probe:
args.dset = args.dset_finetune
# Finetune
suggested_lr = find_lr(head_type='prediction')
linear_probe_func(suggested_lr)
print('finetune completed')
# Test
out = test_func(args.save_path+args.save_finetuned_model)
print('----------- Complete! -----------')
else:
args.dset = args.dset_finetune
weight_path = args.save_path+args.dset_finetune+'_patchtst_finetuned'+suffix_name
# Test
out = test_func(weight_path)
print('----------- Complete! -----------')
@@ -0,0 +1,153 @@
import numpy as np
import pandas as pd
import os
import torch
from torch import nn
from src.models.patchTST import PatchTST
from src.learner import Learner, transfer_weights
from src.callback.tracking import *
from src.callback.patch_mask import *
from src.callback.transforms import *
from src.metrics import *
from src.basics import set_device
from datautils import *
import argparse
parser = argparse.ArgumentParser()
# Dataset and dataloader
parser.add_argument('--dset_pretrain', type=str, default='etth1', help='dataset name')
parser.add_argument('--context_points', type=int, default=512, help='sequence length')
parser.add_argument('--target_points', type=int, default=96, help='forecast horizon')
parser.add_argument('--batch_size', type=int, default=64, help='batch size')
parser.add_argument('--num_workers', type=int, default=0, help='number of workers for DataLoader')
parser.add_argument('--scaler', type=str, default='standard', help='scale the input data')
parser.add_argument('--features', type=str, default='M', help='for multivariate model or univariate model')
# Patch
parser.add_argument('--patch_len', type=int, default=12, help='patch length')
parser.add_argument('--stride', type=int, default=12, help='stride between patch')
# RevIN
parser.add_argument('--revin', type=int, default=1, help='reversible instance normalization')
# Model args
parser.add_argument('--n_layers', type=int, default=3, help='number of Transformer layers')
parser.add_argument('--n_heads', type=int, default=16, help='number of Transformer heads')
parser.add_argument('--d_model', type=int, default=128, help='Transformer d_model')
parser.add_argument('--d_ff', type=int, default=512, help='Tranformer MLP dimension')
parser.add_argument('--dropout', type=float, default=0.2, help='Transformer dropout')
parser.add_argument('--head_dropout', type=float, default=0.2, help='head dropout')
# Pretrain mask
parser.add_argument('--mask_ratio', type=float, default=0.4, help='masking ratio for the input')
# Optimization args
parser.add_argument('--n_epochs_pretrain', type=int, default=10, help='number of pre-training epochs')
parser.add_argument('--lr', type=float, default=1e-4, help='learning rate')
# model id to keep track of the number of models saved
parser.add_argument('--pretrained_model_id', type=int, default=1, help='id of the saved pretrained model')
parser.add_argument('--model_type', type=str, default='based_model', help='for multivariate model or univariate model')
args = parser.parse_args()
print('args:', args)
args.save_pretrained_model = 'patchtst_pretrained_cw'+str(args.context_points)+'_patch'+str(args.patch_len) + '_stride'+str(args.stride) + '_epochs-pretrain' + str(args.n_epochs_pretrain) + '_mask' + str(args.mask_ratio) + '_model' + str(args.pretrained_model_id)
args.save_path = 'saved_models/' + args.dset_pretrain + '/masked_patchtst/' + args.model_type + '/'
if not os.path.exists(args.save_path): os.makedirs(args.save_path)
# get available GPU devide
set_device()
def get_model(c_in, args):
"""
c_in: number of variables
"""
# get number of patches
num_patch = (max(args.context_points, args.patch_len)-args.patch_len) // args.stride + 1
print('number of patches:', num_patch)
# get model
model = PatchTST(c_in=c_in,
target_dim=args.target_points,
patch_len=args.patch_len,
stride=args.stride,
num_patch=num_patch,
n_layers=args.n_layers,
n_heads=args.n_heads,
d_model=args.d_model,
shared_embedding=True,
d_ff=args.d_ff,
dropout=args.dropout,
head_dropout=args.head_dropout,
act='relu',
head_type='pretrain',
res_attention=False
)
# print out the model size
print('number of model params', sum(p.numel() for p in model.parameters() if p.requires_grad))
return model
def find_lr():
# get dataloader
dls = get_dls(args)
model = get_model(dls.vars, args)
# get loss
loss_func = torch.nn.MSELoss(reduction='mean')
# get callbacks
cbs = [RevInCB(dls.vars, denorm=False)] if args.revin else []
cbs += [PatchMaskCB(patch_len=args.patch_len, stride=args.stride, mask_ratio=args.mask_ratio)]
# define learner
learn = Learner(dls, model,
loss_func,
lr=args.lr,
cbs=cbs,
)
# fit the data to the model
suggested_lr = learn.lr_finder()
print('suggested_lr', suggested_lr)
return suggested_lr
def pretrain_func(lr=args.lr):
# get dataloader
dls = get_dls(args)
# get model
model = get_model(dls.vars, args)
# get loss
loss_func = torch.nn.MSELoss(reduction='mean')
# get callbacks
cbs = [RevInCB(dls.vars, denorm=False)] if args.revin else []
cbs += [
PatchMaskCB(patch_len=args.patch_len, stride=args.stride, mask_ratio=args.mask_ratio),
SaveModelCB(monitor='valid_loss', fname=args.save_pretrained_model,
path=args.save_path)
]
# define learner
learn = Learner(dls, model,
loss_func,
lr=lr,
cbs=cbs,
#metrics=[mse]
)
# fit the data to the model
learn.fit_one_cycle(n_epochs=args.n_epochs_pretrain, lr_max=lr)
train_loss = learn.recorder['train_loss']
valid_loss = learn.recorder['valid_loss']
df = pd.DataFrame(data={'train_loss': train_loss, 'valid_loss': valid_loss})
df.to_csv(args.save_path + args.save_pretrained_model + '_losses.csv', float_format='%.6f', index=False)
if __name__ == '__main__':
args.dset = args.dset_pretrain
suggested_lr = find_lr()
# Pretrain
pretrain_func(suggested_lr)
print('pretraining completed')
@@ -0,0 +1,162 @@
import numpy as np
import pandas as pd
import os
import torch
from torch import nn
from src.models.patchTST import PatchTST
from src.learner import Learner
from src.callback.core import *
from src.callback.tracking import *
from src.callback.scheduler import *
from src.callback.patch_mask import *
from src.callback.transforms import *
from src.metrics import *
from datautils import get_dls
import argparse
parser = argparse.ArgumentParser()
# Dataset and dataloader
parser.add_argument('--dset', type=str, default='etth1', help='dataset name')
parser.add_argument('--context_points', type=int, default=336, help='sequence length')
parser.add_argument('--target_points', type=int, default=96, help='forecast horizon')
parser.add_argument('--batch_size', type=int, default=64, help='batch size')
parser.add_argument('--num_workers', type=int, default=1, help='number of workers for DataLoader')
parser.add_argument('--scaler', type=str, default='standard', help='scale the input data')
parser.add_argument('--features', type=str, default='M', help='for multivariate model or univariate model')
parser.add_argument('--use_time_features', type=int, default=0, help='whether to use time features or not')
# Patch
parser.add_argument('--patch_len', type=int, default=32, help='patch length')
parser.add_argument('--stride', type=int, default=16, help='stride between patch')
# RevIN
parser.add_argument('--revin', type=int, default=1, help='reversible instance normalization')
# Model args
parser.add_argument('--n_layers', type=int, default=3, help='number of Transformer layers')
parser.add_argument('--n_heads', type=int, default=16, help='number of Transformer heads')
parser.add_argument('--d_model', type=int, default=128, help='Transformer d_model')
parser.add_argument('--d_ff', type=int, default=256, help='Tranformer MLP dimension')
parser.add_argument('--dropout', type=float, default=0.2, help='Transformer dropout')
parser.add_argument('--head_dropout', type=float, default=0, help='head dropout')
# Optimization args
parser.add_argument('--n_epochs', type=int, default=20, help='number of training epochs')
parser.add_argument('--lr', type=float, default=1e-4, help='learning rate')
# model id to keep track of the number of models saved
parser.add_argument('--model_id', type=int, default=1, help='id of the saved model')
parser.add_argument('--model_type', type=str, default='based_model', help='for multivariate model or univariate model')
# training
parser.add_argument('--is_train', type=int, default=1, help='training the model')
args = parser.parse_args()
print('args:', args)
args.save_model_name = 'patchtst_supervised'+'_cw'+str(args.context_points)+'_tw'+str(args.target_points) + '_patch'+str(args.patch_len) + '_stride'+str(args.stride)+'_epochs'+str(args.n_epochs) + '_model' + str(args.model_id)
args.save_path = 'saved_models/' + args.dset + '/patchtst_supervised/' + args.model_type + '/'
if not os.path.exists(args.save_path): os.makedirs(args.save_path)
def get_model(c_in, args):
"""
c_in: number of input variables
"""
# get number of patches
num_patch = (max(args.context_points, args.patch_len)-args.patch_len) // args.stride + 1
print('number of patches:', num_patch)
# get model
model = PatchTST(c_in=c_in,
target_dim=args.target_points,
patch_len=args.patch_len,
stride=args.stride,
num_patch=num_patch,
n_layers=args.n_layers,
n_heads=args.n_heads,
d_model=args.d_model,
shared_embedding=True,
d_ff=args.d_ff,
dropout=args.dropout,
head_dropout=args.head_dropout,
act='relu',
head_type='prediction',
res_attention=False
)
return model
def find_lr():
# get dataloader
dls = get_dls(args)
model = get_model(dls.vars, args)
# get loss
loss_func = torch.nn.MSELoss(reduction='mean')
# get callbacks
cbs = [RevInCB(dls.vars)] if args.revin else []
cbs += [PatchCB(patch_len=args.patch_len, stride=args.stride)]
# define learner
learn = Learner(dls, model, loss_func, cbs=cbs)
# fit the data to the model
return learn.lr_finder()
def train_func(lr=args.lr):
# get dataloader
dls = get_dls(args)
print('in out', dls.vars, dls.c, dls.len)
# get model
model = get_model(dls.vars, args)
# get loss
loss_func = torch.nn.MSELoss(reduction='mean')
# get callbacks
cbs = [RevInCB(dls.vars)] if args.revin else []
cbs += [
PatchCB(patch_len=args.patch_len, stride=args.stride),
SaveModelCB(monitor='valid_loss', fname=args.save_model_name,
path=args.save_path )
]
# define learner
learn = Learner(dls, model,
loss_func,
lr=lr,
cbs=cbs,
metrics=[mse]
)
# fit the data to the model
learn.fit_one_cycle(n_epochs=args.n_epochs, lr_max=lr, pct_start=0.2)
def test_func():
weight_path = args.save_path + args.save_model_name + '.pth'
# get dataloader
dls = get_dls(args)
model = get_model(dls.vars, args)
#model = torch.load(weight_path)
# get callbacks
cbs = [RevInCB(dls.vars)] if args.revin else []
cbs += [PatchCB(patch_len=args.patch_len, stride=args.stride)]
learn = Learner(dls, model,cbs=cbs)
out = learn.test(dls.test, weight_path=weight_path, scores=[mse,mae]) # out: a list of [pred, targ, score_values]
return out
if __name__ == '__main__':
if args.is_train: # training mode
suggested_lr = find_lr()
print('suggested lr:', suggested_lr)
train_func(suggested_lr)
else: # testing mode
out = test_func()
print('score:', out[2])
print('shape:', out[0].shape)
print('----------- Complete! -----------')
@@ -0,0 +1,97 @@
import torch
import collections
from collections import OrderedDict
class GetAttr:
"Inherit from this to have all attr accesses in `self._xtra` passed down to `self.default`"
_default='default'
def _component_attr_filter(self,k):
if k.startswith('__') or k in ('_xtra',self._default): return False
xtra = getattr(self,'_xtra',None)
return xtra is None or k in xtra
def _dir(self):
return [k for k in dir(getattr(self,self._default)) if self._component_attr_filter(k)]
def __getattr__(self, k):
if self._component_attr_filter(k):
attr = getattr(self, self._default, None)
if attr is not None: return getattr(attr,k)
# raise AttributeError(k)
def __dir__(self):
return custom_dir(self,self._dir())
# def __getstate__(self): return self.__dict__
def __setstate__(self,data):
self.__dict__.update(data)
def get_device(use_cuda=True, device_id=None, usage=5):
"Return or set default device; `use_cuda`: None - CUDA if available; True - error if not available; False - CPU"
if not torch.cuda.is_available():
use_cuda = False
else:
if device_id is None:
device_ids = get_available_cuda(usage=usage)
device_id = device_ids[0] # get the first available device
torch.cuda.set_device(device_id)
return torch.device(torch.cuda.current_device()) if use_cuda else torch.device('cpu')
def set_device(usage=5):
"set the device that has usage < default usage "
device_ids = get_available_cuda(usage=usage)
torch.cuda.set_device(device_ids[0]) # get the first available device
def default_device(use_cuda=True):
"Return or set default device; `use_cuda`: None - CUDA if available; True - error if not available; False - CPU"
if not torch.cuda.is_available():
use_cuda = False
return torch.device(torch.cuda.current_device()) if use_cuda else torch.device('cpu')
def get_available_cuda(usage=10):
if not torch.cuda.is_available(): return
# collect available cuda devices, only collect devices that has less that 'usage' percent
device_ids = []
for device in range(torch.cuda.device_count()):
if torch.cuda.utilization(device) < usage: device_ids.append(device)
return device_ids
def to_device(b, device=None, non_blocking=False):
"""
Recursively put `b` on `device`
components of b are torch tensors
"""
if device is None:
device = default_device(use_cuda=True)
if isinstance(b, dict):
return {key: to_device(val, device) for key, val in b.items()}
if isinstance(b, (list, tuple)):
return type(b)(to_device(o, device) for o in b)
return b.to(device, non_blocking=non_blocking)
def to_numpy(b):
"""
Components of b are torch tensors
"""
if isinstance(b, dict):
return {key: to_numpy(val) for key, val in b.items()}
if isinstance(b, (list, tuple)):
return type(b)(to_numpy(o) for o in b)
return b.detach().cpu().numpy()
@@ -0,0 +1,90 @@
__all__ = ['Callback', 'SetupLearnerCB', 'GetPredictionsCB', 'GetTestCB' ]
"""
Callback lists:
> before_fit
- before_epoch
+ before_epoch_train
~ before_batch_train
~ after_batch_train
+ after_epoch_train
+ before_epoch_valid
~ before_batch_valid
~ after_batch_valid
+ after_epoch_valid
- after_epoch
> after_fit
- before_predict
~ before_batch_predict
~ after_batch_predict
- after_predict
"""
from ..basics import *
import torch
DTYPE = torch.float32
class Callback(GetAttr):
_default='learner'
class SetupLearnerCB(Callback):
def __init__(self):
self.device = default_device(use_cuda=True)
def before_batch_train(self): self._to_device()
def before_batch_valid(self): self._to_device()
def before_batch_predict(self): self._to_device()
def before_batch_test(self): self._to_device()
def _to_device(self):
batch = to_device(self.batch, self.device)
if self.n_inp > 1: xb, yb = batch
else: xb, yb = batch, None
self.learner.batch = xb, yb
def before_fit(self):
"Set model to cuda before training"
self.learner.model.to(self.device)
self.learner.device = self.device
class GetPredictionsCB(Callback):
def __init__(self):
super().__init__()
def before_predict(self):
self.preds = []
def after_batch_predict(self):
# append the prediction after each forward batch
self.preds.append(self.pred)
def after_predict(self):
self.preds = torch.concat(self.preds)#.detach().cpu().numpy()
class GetTestCB(Callback):
def __init__(self):
super().__init__()
def before_test(self):
self.preds, self.targets = [], []
def after_batch_test(self):
# append the prediction after each forward batch
self.preds.append(self.pred)
self.targets.append(self.yb)
def after_test(self):
self.preds = torch.concat(self.preds)#.detach().cpu().numpy()
self.targets = torch.concat(self.targets)#.detach().cpu().numpy()
@@ -0,0 +1,173 @@
from .core import Callback
import torch
from torch.utils.data import DistributedSampler, DataLoader, SequentialSampler
from torch.nn.parallel import DistributedDataParallel
from typing import Optional, Dict, Any
import logging
logger = logging.getLogger(__name__)
class DistributedTrainer(Callback):
"Wrap `model` in `DistributedDataParallel` and `dls` in `DistributedDL`"
def __init__(self,
local_rank,
world_size,
sync_bn=True, # Whether to replace all batch norm with `nn.SyncBatchNorm`
**kwargs
):
self.local_rank = local_rank
self.world_size = world_size
self.sync_bn = sync_bn
self.kwargs = kwargs
super().__init__()
def before_fit(self):
self.learner.model = self.prepare_model(
torch.nn.SyncBatchNorm.convert_sync_batchnorm(self.model) if self.sync_bn else self.model,
ddp_kwargs=self.kwargs
)
self.old_train_dl = self.dls.train
self.old_valid_dl = self.dls.valid
self.learner.dls.train = self._wrap_dl(self.dls.train)
self.learner.dls.valid = self._wrap_dl(self.dls.valid)
def _wrap_dl(self, dl):
return dl if isinstance(dl, DistributedDL) else self.prepare_data_loader(dl)
def after_fit(self):
self.learner.model = self.learner.model.module
self.learner.dls.train = self.old_train_dl
self.learner.dls.valid = self.old_valid_dl
def prepare_model(self,
model: torch.nn.Module,
move_to_device: bool = True,
wrap_ddp: bool = True,
ddp_kwargs: Optional[Dict[str, Any]] = None) -> torch.nn.Module:
"""Prepares the model for distributed execution.
Args:
model (torch.nn.Module): A torch model to prepare.
move_to_device (bool): Whether to move the model to the correct
device. If set to False, the model needs to manually be moved
to the correct device.
wrap_ddp (bool): Whether to wrap models in
``DistributedDataParallel``.
ddp_kwargs (Dict[str, Any]): Args to pass into
``DistributedDataParallel`` initialization if ``wrap_ddp`` is
set to True.
"""
ddp_kwargs = ddp_kwargs or {}
rank = self.local_rank
device = torch.device(f"cuda:{rank}")
# device = get_device()
if torch.cuda.is_available():
torch.cuda.set_device(device)
if move_to_device:
logger.info(f"Moving model to device: {device}")
model = model.to(device)
if wrap_ddp and self.world_size > 1:
logger.info("Wrapping provided model in DDP.")
if torch.cuda.is_available():
model = DistributedDataParallel(
model, device_ids=[rank], output_device=rank, **ddp_kwargs)
else:
model = DistributedDataParallel(model, **ddp_kwargs)
return model
def prepare_data_loader(self,
data_loader: torch.utils.data.DataLoader,
add_dist_sampler: bool = True,
move_to_device: bool = True) -> torch.utils.data.DataLoader:
"""
Prepares DataLoader for distributed execution.
This allows you to use the same exact code regardless of number of
workers or the device type being used (CPU, GPU).
Args:
data_loader (torch.utils.data.DataLoader): The DataLoader to
prepare.
add_dist_sampler (bool): Whether to add a DistributedSampler to
the provided DataLoader.
move_to_device (bool): If set, automatically move the data
returned by the data loader to the correct device.
"""
# Only add Distributed Sampler if the following conditions hold:
# 1. More than one training worker is being used.
# 2. A DistributedSampler has not already been added by the user.
# 3. The dataset is not an IterableDataset. Samplers do not worker with
# IterableDatasets.
def with_sampler(loader):
# Automatically set the DistributedSampler
# If using a sampler, the shuffle attribute in the
# DataLoader must be set to False.
# Instead the shuffling is determined by the shuffle attribute
# in the DistributedSampler.
# We identify if shuffling is enabled in the passed in
# DataLoader by seeing if the sampler for the DataLoader is a
# SequentialSampler.
shuffle = not isinstance(loader.sampler, SequentialSampler)
data_loader_args = {
"dataset": loader.dataset,
"batch_size": loader.batch_size,
"shuffle": False,
"num_workers": loader.num_workers,
"collate_fn": loader.collate_fn,
"pin_memory": loader.pin_memory,
"drop_last": loader.drop_last,
"timeout": loader.timeout,
"worker_init_fn": loader.worker_init_fn,
"sampler": DistributedSampler(loader.dataset, shuffle=shuffle)
}
return DataLoader(**data_loader_args)
data_loader = with_sampler(data_loader)
if move_to_device:
if torch.cuda.is_available():
rank = self.local_rank
device = torch.device(f"cuda:{rank}")
else:
device = torch.device("cpu")
data_loader = DistributedDL(data_loader, device)
return data_loader
class DistributedDL(DataLoader):
def __init__(self, base_dataloader: DataLoader, device: torch.device):
self.__dict__.update(getattr(base_dataloader, "__dict__", {}))
self.dataloader = base_dataloader
self.device = device
def _move_to_device(self, item):
def try_move_device(i):
try:
i = i.to(self.device)
except AttributeError:
logger.debug(f"Item {i} cannot be moved to device "
f"{self.device}.")
return i
return tuple(try_move_device(i) for i in item)
def __len__(self):
return len(self.dataloader)
def __iter__(self):
iterator = iter(self.dataloader)
for item in iterator:
yield self._move_to_device(item)
@@ -0,0 +1,176 @@
import torch
from torch import nn
from .core import Callback
# Cell
class PatchCB(Callback):
def __init__(self, patch_len, stride ):
"""
Callback used to perform patching on the batch input data
Args:
patch_len: patch length
stride: stride
"""
self.patch_len = patch_len
self.stride = stride
def before_forward(self): self.set_patch()
def set_patch(self):
"""
take xb from learner and convert to patch: [bs x seq_len x n_vars] -> [bs x num_patch x n_vars x patch_len]
"""
xb_patch, num_patch = create_patch(self.xb, self.patch_len, self.stride) # xb: [bs x seq_len x n_vars]
# learner get the transformed input
self.learner.xb = xb_patch # xb_patch: [bs x num_patch x n_vars x patch_len]
class PatchMaskCB(Callback):
def __init__(self, patch_len, stride, mask_ratio,
mask_when_pred:bool=False):
"""
Callback used to perform the pretext task of reconstruct the original data after a binary mask has been applied.
Args:
patch_len: patch length
stride: stride
mask_ratio: mask ratio
"""
self.patch_len = patch_len
self.stride = stride
self.mask_ratio = mask_ratio
def before_fit(self):
# overwrite the predefined loss function
self.learner.loss_func = self._loss
device = self.learner.device
def before_forward(self): self.patch_masking()
def patch_masking(self):
"""
xb: [bs x seq_len x n_vars] -> [bs x num_patch x n_vars x patch_len]
"""
xb_patch, num_patch = create_patch(self.xb, self.patch_len, self.stride) # xb_patch: [bs x num_patch x n_vars x patch_len]
xb_mask, _, self.mask, _ = random_masking(xb_patch, self.mask_ratio) # xb_mask: [bs x num_patch x n_vars x patch_len]
self.mask = self.mask.bool() # mask: [bs x num_patch x n_vars]
self.learner.xb = xb_mask # learner.xb: masked 4D tensor
self.learner.yb = xb_patch # learner.yb: non-masked 4d tensor
def _loss(self, preds, target):
"""
preds: [bs x num_patch x n_vars x patch_len]
targets: [bs x num_patch x n_vars x patch_len]
"""
loss = (preds - target) ** 2
loss = loss.mean(dim=-1)
loss = (loss * self.mask).sum() / self.mask.sum()
return loss
def create_patch(xb, patch_len, stride):
"""
xb: [bs x seq_len x n_vars]
"""
seq_len = xb.shape[1]
num_patch = (max(seq_len, patch_len)-patch_len) // stride + 1
tgt_len = patch_len + stride*(num_patch-1)
s_begin = seq_len - tgt_len
xb = xb[:, s_begin:, :] # xb: [bs x tgt_len x nvars]
xb = xb.unfold(dimension=1, size=patch_len, step=stride) # xb: [bs x num_patch x n_vars x patch_len]
return xb, num_patch
class Patch(nn.Module):
def __init__(self,seq_len, patch_len, stride):
super().__init__()
self.seq_len = seq_len
self.patch_len = patch_len
self.stride = stride
self.num_patch = (max(seq_len, patch_len)-patch_len) // stride + 1
tgt_len = patch_len + stride*(self.num_patch-1)
self.s_begin = seq_len - tgt_len
def forward(self, x):
"""
x: [bs x seq_len x n_vars]
"""
x = x[:, self.s_begin:, :]
x = x.unfold(dimension=1, size=self.patch_len, step=self.stride) # xb: [bs x num_patch x n_vars x patch_len]
return x
def random_masking(xb, mask_ratio):
# xb: [bs x num_patch x n_vars x patch_len]
bs, L, nvars, D = xb.shape
x = xb.clone()
len_keep = int(L * (1 - mask_ratio))
noise = torch.rand(bs, L, nvars,device=xb.device) # noise in [0, 1], bs x L x nvars
# sort noise for each sample
ids_shuffle = torch.argsort(noise, dim=1) # ascend: small is keep, large is remove
ids_restore = torch.argsort(ids_shuffle, dim=1) # ids_restore: [bs x L x nvars]
# keep the first subset
ids_keep = ids_shuffle[:, :len_keep, :] # ids_keep: [bs x len_keep x nvars]
x_kept = torch.gather(x, dim=1, index=ids_keep.unsqueeze(-1).repeat(1, 1, 1, D)) # x_kept: [bs x len_keep x nvars x patch_len]
# removed x
x_removed = torch.zeros(bs, L-len_keep, nvars, D, device=xb.device) # x_removed: [bs x (L-len_keep) x nvars x patch_len]
x_ = torch.cat([x_kept, x_removed], dim=1) # x_: [bs x L x nvars x patch_len]
# combine the kept part and the removed one
x_masked = torch.gather(x_, dim=1, index=ids_restore.unsqueeze(-1).repeat(1,1,1,D)) # x_masked: [bs x num_patch x nvars x patch_len]
# generate the binary mask: 0 is keep, 1 is remove
mask = torch.ones([bs, L, nvars], device=x.device) # mask: [bs x num_patch x nvars]
mask[:, :len_keep, :] = 0
# unshuffle to get the binary mask
mask = torch.gather(mask, dim=1, index=ids_restore) # [bs x num_patch x nvars]
return x_masked, x_kept, mask, ids_restore
def random_masking_3D(xb, mask_ratio):
# xb: [bs x num_patch x dim]
bs, L, D = xb.shape
x = xb.clone()
len_keep = int(L * (1 - mask_ratio))
noise = torch.rand(bs, L, device=xb.device) # noise in [0, 1], bs x L
# sort noise for each sample
ids_shuffle = torch.argsort(noise, dim=1) # ascend: small is keep, large is remove
ids_restore = torch.argsort(ids_shuffle, dim=1) # ids_restore: [bs x L]
# keep the first subset
ids_keep = ids_shuffle[:, :len_keep] # ids_keep: [bs x len_keep]
x_kept = torch.gather(x, dim=1, index=ids_keep.unsqueeze(-1).repeat(1, 1, D)) # x_kept: [bs x len_keep x dim]
# removed x
x_removed = torch.zeros(bs, L-len_keep, D, device=xb.device) # x_removed: [bs x (L-len_keep) x dim]
x_ = torch.cat([x_kept, x_removed], dim=1) # x_: [bs x L x dim]
# combine the kept part and the removed one
x_masked = torch.gather(x_, dim=1, index=ids_restore.unsqueeze(-1).repeat(1,1,D)) # x_masked: [bs x num_patch x dim]
# generate the binary mask: 0 is keep, 1 is remove
mask = torch.ones([bs, L], device=x.device) # mask: [bs x num_patch]
mask[:, :len_keep] = 0
# unshuffle to get the binary mask
mask = torch.gather(mask, dim=1, index=ids_restore) # [bs x num_patch]
return x_masked, x_kept, mask, ids_restore
if __name__ == "__main__":
bs, L, nvars, D = 2,20,4,5
xb = torch.randn(bs, L, nvars, D)
xb_mask, mask, ids_restore = create_mask(xb, mask_ratio=0.5)
breakpoint()
@@ -0,0 +1,212 @@
__all__ = ['OneCycleLR', 'LRFinderCB', 'LinearLR', 'ExponentialLR']
from cmath import inf
from ..basics import *
from .core import Callback
from torch.optim import lr_scheduler
from torch.optim.lr_scheduler import _LRScheduler
class OneCycleLR(Callback):
def __init__(self, lr_max=None,
total_steps=None,
steps_per_epoch=None,
pct_start=0.3,
anneal_strategy='cos',
cycle_momentum=True,
base_momentum=0.85,
max_momentum=0.95,
div_factor=25.,
final_div_factor=1e4,
three_phase=False,
last_epoch=-1,
verbose=False):
super().__init__()
self.lr_max = lr_max if lr_max else self.lr
self.total_steps, self.steps_per_epoch = total_steps, steps_per_epoch
self.pct_start = pct_start
self.anneal_strategy, self.cycle_momentum = anneal_strategy, cycle_momentum
self.base_momentum, self.max_momentum = base_momentum, max_momentum
self.div_factor, self.final_div_factor = div_factor, final_div_factor
self.three_phase = three_phase
self.last_epoch = last_epoch
self.verbose = verbose
def before_fit(self):
if not self.steps_per_epoch: self.steps_per_epoch = len(self.dls.train)
self.lrs = [] # store lr values
self.scheduler = lr_scheduler.OneCycleLR(optimizer = self.opt,
max_lr = self.lr_max,
total_steps = self.total_steps,
epochs = self.n_epochs,
steps_per_epoch=self.steps_per_epoch,
pct_start=self.pct_start,
anneal_strategy=self.anneal_strategy,
cycle_momentum=self.cycle_momentum,
base_momentum=self.base_momentum,
max_momentum=self.max_momentum,
div_factor=self.div_factor,
final_div_factor=self.final_div_factor,
three_phase=self.three_phase,
last_epoch=self.last_epoch,
verbose=self.verbose
)
def after_batch_train(self):
if self.model.training:
self.scheduler.step()
self.lrs.append( self.scheduler.get_last_lr()[0] )
def after_fit(self):
self.learner.scheduled_lrs = self.lrs
class LRFinderCB(Callback):
def __init__(self, start_lr=1e-7, end_lr=10, num_iter=100, step_mode='exp', beta=0.98, suggestion='valley'):
self.start_lr, self.end_lr = start_lr, end_lr
self.num_iter = num_iter
self.step_mode = step_mode
if beta >= 1: raise ValueError("`num_iter` must be smaller than 1")
else: self.beta = beta
self.suggestion = suggestion
def before_fit(self):
self.losses, self.lrs = [], []
self.best_loss, self.aver_loss = inf, 0
self.train_iter = 0
# save model to load back after fitting
self.temp_path = self.save('current', 'temp/', with_opt=False)
# set base_lr for the optimizer
self.set_lr(self.start_lr)
# check num_iter
if not self.num_iter: self.num_iter = len(self.dls.train)
# if self.num_iter > len(self.dls.train): self.num_iter = len(self.dls.train)
# Initialize the proper learning rate policy
if self.step_mode.lower() == "exp":
self.scheduler = ExponentialLR(self.opt, self.end_lr, self.num_iter)
elif self.step_mode.lower() == "linear":
self.scheduler = LinearLR(self.opt, self.end_lr, self.num_iter)
def after_batch_train(self):
self.train_iter += 1
self.scheduler.step()
self.lrs.append( self.scheduler.get_last_lr()[0] )
# update smooth loss
self.smoothing(self.beta)
if self.smoothed_loss < self.best_loss: self.best_loss = self.smoothed_loss
#Stop if the loss is exploding
if self.smoothed_loss > 4 * self.best_loss:
raise KeyboardInterrupt # stop fit method
if self.train_iter > self.num_iter:
raise KeyboardInterrupt # stop fit method
def smoothing(self, beta):
# Smooth the loss if beta is specified
self.aver_loss = beta * self.aver_loss + (1-beta) *self.loss.detach().item()
self.smoothed_loss = self.aver_loss / (1 - beta**self.train_iter)
self.losses.append(self.smoothed_loss)
def after_fit(self):
# reset the gradients
self.learner.opt.zero_grad()
if self.suggestion == 'valley':
self.suggested_lr = valley(self.lrs, self.losses)
# load back the model at the previous state
self.load(self.temp_path)
def set_lr(self, lrs):
if not isinstance(lrs, list): lrs = [lrs] * len(self.opt.param_groups)
if len(lrs) != len(self.opt.param_groups):
raise ValueError(
"Length of `lrs` is not equal to the number of parameter groups "
+ "in the given optimizer")
# update lr
for param_group, lr in zip(self.opt.param_groups, lrs):
param_group["lr"] = lr
def plot_lr_find(self):
import matplotlib.pyplot as plt
fig, ax = plt.subplots(1,1)
ax.plot(self.lrs, self.losses)
ax.set_ylabel("Loss")
ax.set_xlabel("Learning Rate")
ax.set_xscale('log')
plt.grid()
class LinearLR(_LRScheduler):
"""Linearly increases the learning rate between two boundaries over a number of iterations.
Arguments:
optimizer (torch.optim.Optimizer): wrapped optimizer.
end_lr (float): the final learning rate.
num_iter (int): the number of iterations over which the test occurs.
last_epoch (int, optional): the index of last epoch. Default: -1.
"""
def __init__(self, optimizer, end_lr, num_iter, last_epoch=-1):
self.end_lr = end_lr
if num_iter <= 1: raise ValueError("`num_iter` must be larger than 1")
self.num_iter = num_iter
super(LinearLR, self).__init__(optimizer, last_epoch)
def get_lr(self):
r = (self.last_epoch+1) / (self.num_iter - 1)
return [base_lr + r * (self.end_lr - base_lr) for base_lr in self.base_lrs]
class ExponentialLR(_LRScheduler):
"""Exponentially increases the learning rate between two boundaries over a number of iterations.
Arguments:
optimizer (torch.optim.Optimizer): wrapped optimizer.
end_lr (float): the final learning rate.
num_iter (int): the number of iterations over which the test occurs.
last_epoch (int, optional): the index of last epoch. Default: -1.
"""
def __init__(self, optimizer, end_lr, num_iter, last_epoch=-1):
self.end_lr = end_lr
self.last_epoch = last_epoch
if num_iter <= 1: raise ValueError("`num_iter` must be larger than 1")
self.num_iter = num_iter
super(ExponentialLR, self).__init__(optimizer, last_epoch)
def get_lr(self):
r = (self.last_epoch+1) / (self.num_iter - 1)
return [base_lr * (self.end_lr / base_lr) ** r for base_lr in self.base_lrs]
def valley(lrs:list, losses:list):
"Suggests a learning rate from the longest valley and returns its index"
n = len(losses)
max_start, max_end = 0,0
# find the longest valley
lds = [1]*n
for i in range(1,n):
for j in range(0,i):
if (losses[i] < losses[j]) and (lds[i] < lds[j] + 1):
lds[i] = lds[j] + 1
if lds[max_end] < lds[i]:
max_end = i
max_start = max_end - lds[max_end]
sections = (max_end - max_start) / 3
idx = max_start + int(sections) + int(sections/2)
return float(lrs[idx])
@@ -0,0 +1,284 @@
__all__ = ['TrackTimerCB', 'TrackTrainingCB', 'PrintResultsCB', 'TerminateOnNaNCB',
'TrackerCB', 'SaveModelCB', 'EarlyStoppingCB']
from ..basics import *
from .core import Callback
import torch
import time
import numpy as np
from pathlib import Path
class TrackTimerCB(Callback):
def __init__(self):
super().__init__()
def before_fit(self):
self.learner.epoch_time = None
def before_epoch_train(self):
self.start_time = time.time()
def after_epoch_train(self):
self.learner.epoch_time = self.format_time(time.time() - self.start_time)
def format_time(self, t):
"Format `t` (in seconds) to (h):mm:ss"
t = int(t)
h, m, s = t // 3600, (t // 60) % 60, t % 60
if h != 0:
return f'{h}:{m:02d}:{s:02d}'
else:
return f'{m:02d}:{s:02d}'
class TrackTrainingCB(Callback):
def __init__(self, train_metrics=False, valid_metrics=True):
super().__init__()
self.train_metrics, self.valid_metrics = train_metrics, valid_metrics
def init_cb_(self):
self.setup()
self.initialize_recorder()
if hasattr(self.loss_func, 'reduction'):
self.mean_reduction_ = True if self.loss_func.reduction == 'mean' else False
def before_fit(self):
self.setup()
self.initialize_recorder()
if hasattr(self.loss_func, 'reduction'):
self.mean_reduction_ = True if self.loss_func.reduction == 'mean' else False
def setup(self):
self.valid_loss = False
if self.learner.dls:
if not self.learner.dls.valid: self.valid_metrics = False
else: self.valid_loss = True
if self.metrics:
if not isinstance(self.metrics, list): self.metrics = [self.metrics]
self.metric_names = [func.__name__ for func in self.metrics]
else: self.metrics, self.metric_names = [], []
def initialize_recorder(self):
recorder = {'epoch': [], 'train_loss': []}
if self.valid_loss: recorder['valid_loss'] = []
for name in self.metric_names:
if self.train_metrics: recorder['train_'+name] = []
if self.valid_metrics: recorder['valid_'+name] = []
self.recorder = recorder
self.learner.recorder = recorder
def initialize_batch_recorder(self, with_metrics):
batch_recorder = {'n_samples': [], 'batch_losses': [], 'with_metrics': with_metrics}
self.batch_recorder = batch_recorder
def reset(self):
self.targs, self.preds = [],[]
self.n_samples = 0
self.batch_loss = []
def after_epoch(self):
self.recorder['epoch'].append(self.epoch)
self.learner.recorder = self.recorder
def before_epoch_train(self):
# define storage for batch training loss and metrics
self.initialize_batch_recorder(with_metrics=self.train_metrics)
self.reset()
def before_epoch_valid(self):
# if valid data is available, define storage for batch training loss and metrics
# if self.dls.valid: self.initialize_batch_recorder(with_metrics=self.valid_metrics)
self.initialize_batch_recorder(with_metrics=self.valid_metrics)
self.reset()
def after_epoch_train(self):
values = self.compute_scores()
# save training loss after one epoch
self.recorder['train_loss'].append( values['loss'] )
# save metrics after one epoch
if self.train_metrics:
for name, func in zip(self.metric_names, self.metrics):
self.recorder['train_'+name].append( values[name] )
def after_epoch_valid(self):
# if there is no valid data, don't store
if not self.learner.dls.valid: return
values = self.compute_scores()
# save training loss after one epoch
self.recorder['valid_loss'].append( values['loss'] )
# save metrics after one epoch
if self.valid_metrics:
for name, func in zip(self.metric_names, self.metrics):
self.recorder['valid_'+name].append( values[name] )
def after_batch_train(self): self.accumulate() # save batch recorder
def after_batch_valid(self): self.accumulate()
def accumulate(self ):
xb, yb = self.batch
bs = len(xb)
self.batch_recorder['n_samples'].append(bs)
# get batch loss
loss = self.loss.detach()*bs if self.mean_reduction_ else self.loss.detach()
self.batch_recorder['batch_losses'].append(loss)
if yb is None: self.batch_recorder['with_metrics'] = False
if len(self.metrics) == 0: self.batch_recorder['with_metrics'] = False
# accumulate prediction and target
if self.batch_recorder['with_metrics']:
self.preds.append(self.pred.detach().cpu())
self.targs.append(yb.detach().cpu())
def compute_scores(self):
"calculate losses and metrics after each epoch"
values = {}
# calculate loss after each epoch
n = sum(self.batch_recorder['n_samples']) # get total number of samples
values['loss'] = sum(self.batch_recorder['batch_losses']).item()/n # averaging
# calculate metrics if available after each epoch
if len(self.preds) == 0: return values
self.preds = torch.cat(self.preds)
self.targs = torch.cat(self.targs)
for func in self.metrics:
# values[func.__name__] = func(self.targs, self.preds)
values[func.__name__] = func(self.targs, self.preds)
return values
class TerminateOnNaNCB(Callback):
" A callback to stop the training if loss is NaN"
def after_batch_train(self):
if torch.isinf(self.loss) or torch.isnan(self.loss): raise KeyboardInterrupt
class PrintResultsCB(Callback):
def __init__(self):
super().__init__()
def get_header(self, recorder):
"recorder is a dictionary"
header = list(recorder.keys())
return header+['time']
def before_fit(self):
if self.run_finder: return # don't print if lr_finder is called
if not hasattr(self.learner, 'recorder'): return # don't print if there is no recorder
header = self.get_header(self.learner.recorder)
self.print_header = '{:>15s}'*len(header)
self.print_value = '{:>15d}' + '{:>15.6f}'*(len(header)-2) + '{:>15}'
print(self.print_header.format(*header))
def after_epoch(self):
if self.run_finder: return # don't print if lr_finder is called
if not hasattr(self.learner, 'recorder'): return # don't print if there is no recorder
epoch_logs = []
for key in self.learner.recorder:
value=self.learner.recorder[key][-1] if self.learner.recorder[key] else None
epoch_logs += [value]
if self.learner.epoch_time: epoch_logs.append(self.learner.epoch_time)
# print('epoch_logs', epoch_logs)
print(self.print_value.format(*epoch_logs))
class TrackerCB(Callback):
def __init__(self, monitor='train_loss', comp=None, min_delta=0.):
super().__init__()
if comp is None: comp = np.less if 'loss' in monitor or 'error' in monitor else np.greater
if comp == np.less: min_delta *= -1
self.monitor, self.comp, self.min_delta = monitor, comp, min_delta
def before_fit(self):
if self.run_finder: return
if self.best is None: self.best = float('inf') if self.comp == np.less else -float('inf')
self.monitor_names = list(self.learner.recorder.keys())
assert self.monitor in self.monitor_names
def after_epoch(self):
if self.run_finder: return
val = self.learner.recorder[self.monitor][-1]
if self.comp(val - self.min_delta, self.best): self.best, self.new_best = val,True
else: self.new_best = False
class SaveModelCB(TrackerCB):
def __init__(self, monitor='train_loss', comp=None, min_delta=0.,
every_epoch=False, fname='model', path=None, with_opt=False, save_process_id=0, global_rank=None):
super().__init__(monitor=monitor, comp=comp, min_delta=min_delta)
self.every_epoch = every_epoch
self.last_saved_path = None
self.path, self.fname = path, fname
self.with_opt = with_opt
self.save_process_id = save_process_id
# Identify the worker that saves the model to a file: check if the process' global_rank == save_process_id
# If running locally using either a cpu/gpu without using DDP -> set save_process_id = global_rank
# Else if running in DDP mode but user doesn't specify global_rank -> global_rank = current_device
# (local_rank 0 from each node will save the model)
# Else if user provides the global_rank -> use the global_rank to check
if global_rank:
self.global_rank = int(global_rank)
else:
if torch.cuda.is_available():
self.global_rank = torch.cuda.current_device()
if not torch.distributed.is_initialized():
self.save_process_id = self.global_rank
else:
self.global_rank = 0
def _save(self, fname, path):
if self.global_rank == self.save_process_id:
self.last_saved_path = self.learner.save(fname, path, with_opt=self.with_opt)
def after_epoch(self):
if self.every_epoch:
if ((self.epoch%self.every_epoch) == 0) or (self.epoch==self.n_epochs-1):
self._save(f'{self.fname}_{self.epoch}', self.path)
else:
super().after_epoch()
if self.new_best:
print(f'Better model found at epoch {self.epoch} with {self.monitor} value: {self.best}.')
self._save(f'{self.fname}', self.path)
def after_fit(self):
if self.run_finder: return
if not self.every_epoch and self.global_rank == self.save_process_id:
self.learner.load(self.last_saved_path, with_opt=self.with_opt)
class EarlyStoppingCB(TrackerCB):
def __init__(self, monitor='train_loss', comp=None, min_delta=0,
patient=5):
super().__init__(monitor=monitor, comp=comp, min_delta=min_delta)
self.patient = patient
def before_fit(self):
# set the impatient level
self.impatient_level = 0
super().before_fit()
def after_epoch(self):
super().after_epoch()
if self.new_best: self.impatient_level = 0 # reset the impatience
else:
self.impatient_level += 1
if self.impatient_level > self.patient:
print(f'No improvement since epoch {self.epoch-self.impatient_level}: early stopping')
raise KeyboardInterrupt
@@ -0,0 +1,39 @@
import torch
import torch.nn as nn
from .core import Callback
from src.models.layers.revin import RevIN
class RevInCB(Callback):
def __init__(self, num_features: int, eps=1e-5,
affine:bool=False, denorm:bool=True):
"""
:param num_features: the number of features or channels
:param eps: a value added for numerical stability
:param affine: if True, RevIN has learnable affine parameters
:param denorm: if True, the output will be de-normalized
This callback only works with affine=False.
if affine=True, the learnable affine_weights and affine_bias are not learnt
"""
super().__init__()
self.num_features = num_features
self.eps = eps
self.affine = affine
self.denorm = denorm
self.revin = RevIN(num_features, eps, affine)
def before_forward(self): self.revin_norm()
def after_forward(self):
if self.denorm: self.revin_denorm()
def revin_norm(self):
xb_revin = self.revin(self.xb, 'norm') # xb_revin: [bs x seq_len x nvars]
self.learner.xb = xb_revin
def revin_denorm(self):
pred = self.revin(self.pred, 'denorm') # pred: [bs x target_window x nvars]
self.learner.pred = pred
@@ -0,0 +1,80 @@
import warnings
import torch
from torch.utils.data import DataLoader
class DataLoaders:
def __init__(
self,
datasetCls,
dataset_kwargs: dict,
batch_size: int,
workers: int=0,
collate_fn=None,
shuffle_train = True,
shuffle_val = False
):
super().__init__()
self.datasetCls = datasetCls
self.batch_size = batch_size
if "split" in dataset_kwargs.keys():
del dataset_kwargs["split"]
self.dataset_kwargs = dataset_kwargs
self.workers = workers
self.collate_fn = collate_fn
self.shuffle_train, self.shuffle_val = shuffle_train, shuffle_val
self.train = self.train_dataloader()
self.valid = self.val_dataloader()
self.test = self.test_dataloader()
def train_dataloader(self):
return self._make_dloader("train", shuffle=self.shuffle_train)
def val_dataloader(self):
return self._make_dloader("val", shuffle=self.shuffle_val)
def test_dataloader(self):
return self._make_dloader("test", shuffle=False)
def _make_dloader(self, split, shuffle=False):
dataset = self.datasetCls(**self.dataset_kwargs, split=split)
if len(dataset) == 0: return None
return DataLoader(
dataset,
shuffle=shuffle,
batch_size=self.batch_size,
num_workers=self.workers,
collate_fn=self.collate_fn,
)
@classmethod
def add_cli(self, parser):
parser.add_argument("--batch_size", type=int, default=128)
parser.add_argument(
"--workers",
type=int,
default=6,
help="number of parallel workers for pytorch dataloader",
)
def add_dl(self, test_data, batch_size=None, **kwargs):
# check of test_data is already a DataLoader
from ray.train.torch import _WrappedDataLoader
if isinstance(test_data, DataLoader) or isinstance(test_data, _WrappedDataLoader):
return test_data
# get batch_size if not defined
if batch_size is None: batch_size=self.batch_size
# check if test_data is Dataset, if not, wrap Dataset
if not isinstance(test_data, Dataset):
test_data = self.train.dataset.new(test_data)
# create a new DataLoader from Dataset
test_data = self.train.new(test_data, batch_size, **kwargs)
return test_data
@@ -0,0 +1,416 @@
import os
import numpy as np
import pandas as pd
import os
import torch
from torch.utils.data import Dataset, DataLoader
from sklearn.preprocessing import StandardScaler
from src.data.timefeatures import time_features
import warnings
warnings.filterwarnings('ignore')
class Dataset_ETT_hour(Dataset):
def __init__(self, root_path, split='train', size=None,
features='S', data_path='ETTh1.csv',
target='OT', scale=True, timeenc=0, freq='h',
use_time_features=False
):
# size [seq_len, label_len, pred_len]
# info
if size == None:
self.seq_len = 24 * 4 * 4
self.label_len = 24 * 4
self.pred_len = 24 * 4
else:
self.seq_len = size[0]
self.label_len = size[1]
self.pred_len = size[2]
# init
assert split in ['train', 'test', 'val']
type_map = {'train': 0, 'val': 1, 'test': 2}
self.set_type = type_map[split]
self.features = features
self.target = target
self.scale = scale
self.timeenc = timeenc
self.freq = freq
self.use_time_features = use_time_features
self.root_path = root_path
self.data_path = data_path
self.__read_data__()
def __read_data__(self):
self.scaler = StandardScaler()
df_raw = pd.read_csv(os.path.join(self.root_path,
self.data_path))
border1s = [0, 12 * 30 * 24 - self.seq_len, 12 * 30 * 24 + 4 * 30 * 24 - self.seq_len]
border2s = [12 * 30 * 24, 12 * 30 * 24 + 4 * 30 * 24, 12 * 30 * 24 + 8 * 30 * 24]
border1 = border1s[self.set_type]
border2 = border2s[self.set_type]
if self.features == 'M' or self.features == 'MS':
cols_data = df_raw.columns[1:]
df_data = df_raw[cols_data]
elif self.features == 'S':
df_data = df_raw[[self.target]]
if self.scale:
train_data = df_data[border1s[0]:border2s[0]]
self.scaler.fit(train_data.values)
data = self.scaler.transform(df_data.values)
else:
data = df_data.values
df_stamp = df_raw[['date']][border1:border2]
df_stamp['date'] = pd.to_datetime(df_stamp.date)
if self.timeenc == 0:
df_stamp['month'] = df_stamp.date.apply(lambda row: row.month, 1)
df_stamp['day'] = df_stamp.date.apply(lambda row: row.day, 1)
df_stamp['weekday'] = df_stamp.date.apply(lambda row: row.weekday(), 1)
df_stamp['hour'] = df_stamp.date.apply(lambda row: row.hour, 1)
data_stamp = df_stamp.drop(['date'], axis=1).values
elif self.timeenc == 1:
data_stamp = time_features(pd.to_datetime(df_stamp['date'].values), freq=self.freq)
data_stamp = data_stamp.transpose(1, 0)
self.data_x = data[border1:border2]
self.data_y = data[border1:border2]
self.data_stamp = data_stamp
def __getitem__(self, index):
s_begin = index
s_end = s_begin + self.seq_len
r_begin = s_end - self.label_len
r_end = r_begin + self.label_len + self.pred_len
seq_x = self.data_x[s_begin:s_end]
seq_y = self.data_y[r_begin:r_end]
seq_x_mark = self.data_stamp[s_begin:s_end]
seq_y_mark = self.data_stamp[r_begin:r_end]
if self.use_time_features: return _torch(seq_x, seq_y, seq_x_mark, seq_y_mark)
else: return _torch(seq_x, seq_y)
def __len__(self):
return len(self.data_x) - self.seq_len - self.pred_len + 1
def inverse_transform(self, data):
return self.scaler.inverse_transform(data)
class Dataset_ETT_minute(Dataset):
def __init__(self, root_path, split='train', size=None,
features='S', data_path='ETTm1.csv',
target='OT', scale=True, timeenc=0, freq='t',
use_time_features=False
):
# size [seq_len, label_len, pred_len]
# info
if size == None:
self.seq_len = 24 * 4 * 4
self.label_len = 24 * 4
self.pred_len = 24 * 4
else:
self.seq_len = size[0]
self.label_len = size[1]
self.pred_len = size[2]
# init
assert split in ['train', 'test', 'val']
type_map = {'train': 0, 'val': 1, 'test': 2}
self.set_type = type_map[split]
self.features = features
self.target = target
self.scale = scale
self.timeenc = timeenc
self.freq = freq
self.use_time_features = use_time_features
self.root_path = root_path
self.data_path = data_path
self.__read_data__()
def __read_data__(self):
self.scaler = StandardScaler()
df_raw = pd.read_csv(os.path.join(self.root_path,
self.data_path))
border1s = [0, 12 * 30 * 24 * 4 - self.seq_len, 12 * 30 * 24 * 4 + 4 * 30 * 24 * 4 - self.seq_len]
border2s = [12 * 30 * 24 * 4, 12 * 30 * 24 * 4 + 4 * 30 * 24 * 4, 12 * 30 * 24 * 4 + 8 * 30 * 24 * 4]
border1 = border1s[self.set_type]
border2 = border2s[self.set_type]
if self.features == 'M' or self.features == 'MS':
cols_data = df_raw.columns[1:]
df_data = df_raw[cols_data]
elif self.features == 'S':
df_data = df_raw[[self.target]]
if self.scale:
train_data = df_data[border1s[0]:border2s[0]]
self.scaler.fit(train_data.values)
data = self.scaler.transform(df_data.values)
else:
data = df_data.values
df_stamp = df_raw[['date']][border1:border2]
df_stamp['date'] = pd.to_datetime(df_stamp.date)
if self.timeenc == 0:
df_stamp['month'] = df_stamp.date.apply(lambda row: row.month, 1)
df_stamp['day'] = df_stamp.date.apply(lambda row: row.day, 1)
df_stamp['weekday'] = df_stamp.date.apply(lambda row: row.weekday(), 1)
df_stamp['hour'] = df_stamp.date.apply(lambda row: row.hour, 1)
df_stamp['minute'] = df_stamp.date.apply(lambda row: row.minute, 1)
df_stamp['minute'] = df_stamp.minute.map(lambda x: x // 15)
data_stamp = df_stamp.drop(['date'], axis=1).values
elif self.timeenc == 1:
data_stamp = time_features(pd.to_datetime(df_stamp['date'].values), freq=self.freq)
data_stamp = data_stamp.transpose(1, 0)
self.data_x = data[border1:border2]
self.data_y = data[border1:border2]
self.data_stamp = data_stamp
def __getitem__(self, index):
s_begin = index
s_end = s_begin + self.seq_len
r_begin = s_end - self.label_len
r_end = r_begin + self.label_len + self.pred_len
seq_x = self.data_x[s_begin:s_end]
seq_y = self.data_y[r_begin:r_end]
seq_x_mark = self.data_stamp[s_begin:s_end]
seq_y_mark = self.data_stamp[r_begin:r_end]
if self.use_time_features: return _torch(seq_x, seq_y, seq_x_mark, seq_y_mark)
else: return _torch(seq_x, seq_y)
def __len__(self):
return len(self.data_x) - self.seq_len - self.pred_len + 1
def inverse_transform(self, data):
return self.scaler.inverse_transform(data)
class Dataset_Custom(Dataset):
def __init__(self, root_path, split='train', size=None,
features='S', data_path='ETTh1.csv',
target='OT', scale=True, timeenc=0, freq='h',
time_col_name='date', use_time_features=False,
train_split=0.7, test_split=0.2
):
# size [seq_len, label_len, pred_len]
# info
if size == None:
self.seq_len = 24 * 4 * 4
self.label_len = 24 * 4
self.pred_len = 24 * 4
else:
self.seq_len = size[0]
self.label_len = size[1]
self.pred_len = size[2]
# init
assert split in ['train', 'test', 'val']
type_map = {'train': 0, 'val': 1, 'test': 2}
self.set_type = type_map[split]
self.features = features
self.target = target
self.scale = scale
self.timeenc = timeenc
self.freq = freq
self.time_col_name = time_col_name
self.use_time_features = use_time_features
# train test ratio
self.train_split, self.test_split = train_split, test_split
self.root_path = root_path
self.data_path = data_path
self.__read_data__()
def __read_data__(self):
self.scaler = StandardScaler()
df_raw = pd.read_csv(os.path.join(self.root_path,
self.data_path))
'''
df_raw.columns: [time_col_name, ...(other features), target feature]
'''
cols = list(df_raw.columns)
#cols.remove(self.target) if self.target
#cols.remove(self.time_col_name)
#df_raw = df_raw[[self.time_col_name] + cols + [self.target]]
num_train = int(len(df_raw) * self.train_split)
num_test = int(len(df_raw) * self.test_split)
num_vali = len(df_raw) - num_train - num_test
border1s = [0, num_train - self.seq_len, len(df_raw) - num_test - self.seq_len]
border2s = [num_train, num_train + num_vali, len(df_raw)]
border1 = border1s[self.set_type]
border2 = border2s[self.set_type]
if self.features == 'M' or self.features == 'MS':
cols_data = df_raw.columns[1:]
df_data = df_raw[cols_data]
elif self.features == 'S':
df_data = df_raw[[self.target]]
if self.scale:
train_data = df_data[border1s[0]:border2s[0]]
self.scaler.fit(train_data.values)
data = self.scaler.transform(df_data.values)
else:
data = df_data.values
df_stamp = df_raw[[self.time_col_name]][border1:border2]
df_stamp[self.time_col_name] = pd.to_datetime(df_stamp[self.time_col_name])
if self.timeenc == 0:
df_stamp['month'] = df_stamp[self.time_col_name].apply(lambda row: row.month, 1)
df_stamp['day'] = df_stamp[self.time_col_name].apply(lambda row: row.day, 1)
df_stamp['weekday'] = df_stamp[self.time_col_name].apply(lambda row: row.weekday(), 1)
df_stamp['hour'] = df_stamp[self.time_col_name].apply(lambda row: row.hour, 1)
data_stamp = df_stamp.drop([self.time_col_name], axis=1).values
elif self.timeenc == 1:
data_stamp = time_features(pd.to_datetime(df_stamp[self.time_col_name].values), freq=self.freq)
data_stamp = data_stamp.transpose(1, 0)
self.data_x = data[border1:border2]
self.data_y = data[border1:border2]
self.data_stamp = data_stamp
def __getitem__(self, index):
s_begin = index
s_end = s_begin + self.seq_len
r_begin = s_end - self.label_len
r_end = r_begin + self.label_len + self.pred_len
seq_x = self.data_x[s_begin:s_end]
seq_y = self.data_y[r_begin:r_end]
seq_x_mark = self.data_stamp[s_begin:s_end]
seq_y_mark = self.data_stamp[r_begin:r_end]
if self.use_time_features: return _torch(seq_x, seq_y, seq_x_mark, seq_y_mark)
else: return _torch(seq_x, seq_y)
def __len__(self):
return len(self.data_x) - self.seq_len - self.pred_len + 1
def inverse_transform(self, data):
return self.scaler.inverse_transform(data)
class Dataset_Pred(Dataset):
def __init__(self, root_path, split='pred', size=None,
features='S', data_path='ETTh1.csv',
target='OT', scale=True, inverse=False, timeenc=0, freq='15min', cols=None):
# size [seq_len, label_len, pred_len]
# info
if size == None:
self.seq_len = 24 * 4 * 4
self.label_len = 24 * 4
self.pred_len = 24 * 4
else:
self.seq_len = size[0]
self.label_len = size[1]
self.pred_len = size[2]
# init
assert split in ['pred']
self.features = features
self.target = target
self.scale = scale
self.inverse = inverse
self.timeenc = timeenc
self.freq = freq
self.cols = cols
self.root_path = root_path
self.data_path = data_path
self.__read_data__()
def __read_data__(self):
self.scaler = StandardScaler()
df_raw = pd.read_csv(os.path.join(self.root_path,
self.data_path))
'''
df_raw.columns: ['date', ...(other features), target feature]
'''
if self.cols:
cols = self.cols.copy()
cols.remove(self.target)
else:
cols = list(df_raw.columns)
cols.remove(self.target)
cols.remove('date')
df_raw = df_raw[['date'] + cols + [self.target]]
border1 = len(df_raw) - self.seq_len
border2 = len(df_raw)
if self.features == 'M' or self.features == 'MS':
cols_data = df_raw.columns[1:]
df_data = df_raw[cols_data]
elif self.features == 'S':
df_data = df_raw[[self.target]]
if self.scale:
self.scaler.fit(df_data.values)
data = self.scaler.transform(df_data.values)
else:
data = df_data.values
tmp_stamp = df_raw[['date']][border1:border2]
tmp_stamp['date'] = pd.to_datetime(tmp_stamp.date)
pred_dates = pd.date_range(tmp_stamp.date.values[-1], periods=self.pred_len + 1, freq=self.freq)
df_stamp = pd.DataFrame(columns=['date'])
df_stamp.date = list(tmp_stamp.date.values) + list(pred_dates[1:])
if self.timeenc == 0:
df_stamp['month'] = df_stamp.date.apply(lambda row: row.month, 1)
df_stamp['day'] = df_stamp.date.apply(lambda row: row.day, 1)
df_stamp['weekday'] = df_stamp.date.apply(lambda row: row.weekday(), 1)
df_stamp['hour'] = df_stamp.date.apply(lambda row: row.hour, 1)
df_stamp['minute'] = df_stamp.date.apply(lambda row: row.minute, 1)
df_stamp['minute'] = df_stamp.minute.map(lambda x: x // 15)
data_stamp = df_stamp.drop(['date'], axis=1).values
elif self.timeenc == 1:
data_stamp = time_features(pd.to_datetime(df_stamp['date'].values), freq=self.freq)
data_stamp = data_stamp.transpose(1, 0)
self.data_x = data[border1:border2]
if self.inverse:
self.data_y = df_data.values[border1:border2]
else:
self.data_y = data[border1:border2]
self.data_stamp = data_stamp
def __getitem__(self, index):
s_begin = index
s_end = s_begin + self.seq_len
r_begin = s_end - self.label_len
r_end = r_begin + self.label_len + self.pred_len
seq_x = self.data_x[s_begin:s_end]
if self.inverse:
seq_y = self.data_x[r_begin:r_begin + self.label_len]
else:
seq_y = self.data_y[r_begin:r_begin + self.label_len]
seq_x_mark = self.data_stamp[s_begin:s_end]
seq_y_mark = self.data_stamp[r_begin:r_end]
return seq_x, seq_y, seq_x_mark, seq_y_mark
def __len__(self):
return len(self.data_x) - self.seq_len + 1
def inverse_transform(self, data):
return self.scaler.inverse_transform(data)
def _torch(*dfs):
return tuple(torch.from_numpy(x).float() for x in dfs)
@@ -0,0 +1,134 @@
from typing import List
import numpy as np
import pandas as pd
from pandas.tseries import offsets
from pandas.tseries.frequencies import to_offset
class TimeFeature:
def __init__(self):
pass
def __call__(self, index: pd.DatetimeIndex) -> np.ndarray:
pass
def __repr__(self):
return self.__class__.__name__ + "()"
class SecondOfMinute(TimeFeature):
"""Minute of hour encoded as value between [-0.5, 0.5]"""
def __call__(self, index: pd.DatetimeIndex) -> np.ndarray:
return index.second / 59.0 - 0.5
class MinuteOfHour(TimeFeature):
"""Minute of hour encoded as value between [-0.5, 0.5]"""
def __call__(self, index: pd.DatetimeIndex) -> np.ndarray:
return index.minute / 59.0 - 0.5
class HourOfDay(TimeFeature):
"""Hour of day encoded as value between [-0.5, 0.5]"""
def __call__(self, index: pd.DatetimeIndex) -> np.ndarray:
return index.hour / 23.0 - 0.5
class DayOfWeek(TimeFeature):
"""Hour of day encoded as value between [-0.5, 0.5]"""
def __call__(self, index: pd.DatetimeIndex) -> np.ndarray:
return index.dayofweek / 6.0 - 0.5
class DayOfMonth(TimeFeature):
"""Day of month encoded as value between [-0.5, 0.5]"""
def __call__(self, index: pd.DatetimeIndex) -> np.ndarray:
return (index.day - 1) / 30.0 - 0.5
class DayOfYear(TimeFeature):
"""Day of year encoded as value between [-0.5, 0.5]"""
def __call__(self, index: pd.DatetimeIndex) -> np.ndarray:
return (index.dayofyear - 1) / 365.0 - 0.5
class MonthOfYear(TimeFeature):
"""Month of year encoded as value between [-0.5, 0.5]"""
def __call__(self, index: pd.DatetimeIndex) -> np.ndarray:
return (index.month - 1) / 11.0 - 0.5
class WeekOfYear(TimeFeature):
"""Week of year encoded as value between [-0.5, 0.5]"""
def __call__(self, index: pd.DatetimeIndex) -> np.ndarray:
return (index.isocalendar().week - 1) / 52.0 - 0.5
def time_features_from_frequency_str(freq_str: str) -> List[TimeFeature]:
"""
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.YearEnd: [],
offsets.QuarterEnd: [MonthOfYear],
offsets.MonthEnd: [MonthOfYear],
offsets.Week: [DayOfMonth, WeekOfYear],
offsets.Day: [DayOfWeek, DayOfMonth, DayOfYear],
offsets.BusinessDay: [DayOfWeek, DayOfMonth, DayOfYear],
offsets.Hour: [HourOfDay, DayOfWeek, DayOfMonth, DayOfYear],
offsets.Minute: [
MinuteOfHour,
HourOfDay,
DayOfWeek,
DayOfMonth,
DayOfYear,
],
offsets.Second: [
SecondOfMinute,
MinuteOfHour,
HourOfDay,
DayOfWeek,
DayOfMonth,
DayOfYear,
],
}
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]
supported_freq_msg = f"""
Unsupported frequency {freq_str}
The following frequencies are supported:
Y - yearly
alias: A
M - monthly
W - weekly
D - daily
B - business days
H - hourly
T - minutely
alias: min
S - secondly
"""
raise RuntimeError(supported_freq_msg)
def time_features(dates, freq='h'):
return np.vstack([feat(dates) for feat in time_features_from_frequency_str(freq)])
@@ -0,0 +1,517 @@
from typing import List
import torch
from torch.optim import Adam
from torch import nn
from torch.nn.parallel import DistributedDataParallel
from .basics import *
from .callback.core import *
from .callback.tracking import *
from .callback.scheduler import *
from .callback.distributed import *
from .utils import *
from pathlib import Path
from tqdm import tqdm
import numpy as np
from sklearn.base import BaseEstimator
from unittest.mock import patch
class Learner(GetAttr):
def __init__(self, dls, model,
loss_func=None,
lr=1e-3,
cbs=None,
metrics=None,
opt_func=Adam,
**kwargs):
self.model, self.dls, self.loss_func, self.lr = model, dls, loss_func, lr
self.opt_func = opt_func
#self.opt = self.opt_func(self.model.parameters(), self.lr)
self.set_opt()
self.metrics = metrics
self.n_inp = 2
# self.n_inp = self.dls.train.dataset.n_inp if self.dls else 0
# Initialize callbacks
if cbs and not isinstance(cbs, List): cbs = [cbs]
self.initialize_callbacks(cbs)
# Indicator of running lr_finder
self.run_finder = False
def set_opt(self):
if self.model:
self.opt = self.opt_func(self.model.parameters(), self.lr)
else: self.opt = None
def default_callback(self):
"get a set of default callbacks"
default_cbs = [ SetupLearnerCB(), TrackTimerCB(),
TrackTrainingCB(train_metrics=False, valid_metrics=True)]
return default_cbs
def initialize_callbacks(self, cbs):
default_cbs = self.default_callback()
self.cbs = update_callbacks(cbs, default_cbs) if cbs else default_cbs
# add print CB
self.cbs += [PrintResultsCB()]
for cb in self.cbs: cb.learner = self
self('init_cb')
def add_callback(self, cb):
if not cb: return
cb.learner = self
self.cbs = update_callback(cb, self.cbs)
def add_callbacks(self, cbs):
if not isinstance(cbs, list): cbs = [cbs]
for cb in cbs: self.add_callback(cb)
def remove_callback(self, cb):
cb.learn = None
self.cbs, removed_cb = remove_callback(cb, self.cbs)
return removed_cb
def remove_callbacks(self, cb_list):
for cb in cb_list: self.remove_callback(cb)
def fit(self, n_epochs, lr=None, cbs=None, do_valid=True):
" fit the model "
self.n_epochs = n_epochs
if not self.dls.valid: do_valid = False
if cbs: self.add_callbacks(cbs)
if lr: self.opt = self.opt_func(self.model.parameters(), lr)
self('before_fit')
try:
for self.epoch in range(n_epochs):
self('before_epoch')
self.one_epoch(train=True)
# if self.dls.valid:
if do_valid: self.one_epoch(train=False)
self('after_epoch')
except KeyboardInterrupt: pass
self('after_fit')
def fit_one_cycle(self, n_epochs, lr_max=None, pct_start=0.3):
self.n_epochs = n_epochs
self.lr_max = lr_max if lr_max else self.lr
cb = OneCycleLR(lr_max=self.lr_max, pct_start=pct_start)
self.fit(self.n_epochs, cbs=cb)
def one_epoch(self, train):
self.epoch_train() if train else self.epoch_validate()
def epoch_train(self):
self('before_epoch_train')
self.model.train()
self.dl = self.dls.train
self.all_batches('train')
self('after_epoch_train')
def epoch_validate(self, dl=None):
self('before_epoch_valid')
# model at evaluation mode
self.model.eval()
self.dl = dl if dl else self.dls.valid
if self.dl:
with torch.no_grad(): self.all_batches('valid')
self('after_epoch_valid')
def all_batches(self, type_):
# for self.num,self.batch in enumerate(progress_bar(dl, leave=False)):
for num, batch in enumerate(self.dl):
self.iter, self.batch = num, batch
if type_ == 'train': self.batch_train()
elif type_ == 'valid': self.batch_validate()
elif type_ == 'predict': self.batch_predict()
elif type_ == 'test': self.batch_test()
def batch_train(self):
self('before_batch_train')
self._do_batch_train()
self('after_batch_train')
def batch_validate(self):
self('before_batch_valid')
self._do_batch_validate()
self('after_batch_valid')
def batch_predict(self):
self('before_batch_predict')
self._do_batch_predict()
self('after_batch_predict')
def batch_test(self):
self('before_batch_test')
self._do_batch_test()
self('after_batch_test')
def _do_batch_train(self):
# forward + get loss + backward + optimize
self.pred, self.loss = self.train_step(self.batch)
# zero the parameter gradients
self.opt.zero_grad()
# gradient
self.loss.backward()
# update weights
self.opt.step()
def train_step(self, batch):
# get the inputs
self.xb, self.yb = batch
# forward
pred = self.model_forward()
# compute loss
loss = self.loss_func(pred, self.yb)
return pred, loss
def model_forward(self):
self('before_forward')
self.pred = self.model(self.xb)
self('after_forward')
return self.pred
def _do_batch_validate(self):
# forward + calculate loss
self.pred, self.loss = self.valid_step(self.batch)
def valid_step(self, batch):
# get the inputs
self.xb, self.yb = batch
# forward
pred = self.model_forward()
# compute loss
loss = self.loss_func(pred, self.yb)
return pred, loss
def _do_batch_predict(self):
self.pred = self.predict_step(self.batch)
def predict_step(self, batch):
# get the inputs
self.xb, self.yb = batch
# forward
pred = self.model_forward()
return pred
def _do_batch_test(self):
self.pred, self.yb = self.test_step(self.batch)
def test_step(self, batch):
# get the inputs
self.xb, self.yb = batch
# forward
pred = self.model_forward()
return pred, self.yb
def _predict(self, dl=None):
# self('before_validate')
self('before_predict')
if dl is None: return
self.dl = dl
self.n_inp = dl.dataset.n_inp
self.model.eval() # model at evaluation mode
with torch.no_grad(): self.all_batches('predict')
self('after_predict')
def predict(self, test_data, weight_path=None, Dataset=None, Dataloader=None, batch_size=None):
"""_summary_
Args:
test_data can be a tensor, numpy array, dataset or dataloader
Returns:
_type_: _description_
"""
if weight_path is not None: self.load(weight_path)
cb = GetPredictionsCB()
self.add_callback(cb)
test_dl = self._prepare_data(test_data, Dataset, Dataloader, batch_size)
self._predict(test_dl)
self.preds = cb.preds
return to_numpy(self.preds)
def test(self, dl, weight_path=None, scores=None):
"""_summary_
Args:
test_data can be a tensor, numpy array, dataset or dataloader
Returns:
_type_: _description_
"""
if dl is None: return
else: self.dl = dl
if weight_path is not None: self.load(weight_path)
cb = GetTestCB()
self.add_callback(cb)
self('before_test')
self.model.eval()
with torch.no_grad(): self.all_batches('test')
self('after_test')
self.preds, self.targets = to_numpy([cb.preds, cb.targets])
# calculate scores
if scores:
s_vals = [score(cb.targets, cb.preds).to('cpu').numpy() for score in list(scores)]
return self.preds, self.targets, s_vals
else: return self.preds, self.targets
def _prepare_data(self, test_data, Dataset=None, Dataloader=None, batch_size=None):
if test_data is None: return test_data
if Dataset and Dataloader:
test_dset = Dataset(test_data)
if not batch_size: batch_size=16
test_dl = Dataloader(test_dset, batch_size)
else:
if self.dls:
# add test_data to the dataloader defined in the dls.train
test_dl = self.dls.add_dl(test_data, batch_size=batch_size)
else: test_dl = test_data # assume test_data is already a form of dataloader
return test_dl
def get_layer_output(self, inp, layers=None, unwrap=False):
"""
Args:
inp: can be numpy array, torch tensor or dataloader
"""
self.model.eval()
device = next(self.model.parameters()).device
if isinstance(inp, np.ndarray): inp = torch.Tensor(inp).to(device)
if isinstance(inp, torch.Tensor): inp = inp.to(device)
return get_layer_output(inp, model=self.model, layers=layers, unwrap=unwrap)
def fine_tune(self, n_epochs, base_lr=None, freeze_epochs=1, pct_start=0.3):
"""
fintune the pretrained model. First the entire model is freezed, only head is trained
up to a freeze_epochs number. Then the model is unfreezed and the entire model is trained
"""
assert (n_epochs>0)|(freeze_epochs>0), "Either n_epochs or freeze_epochs has to be > 0"
if not base_lr: base_lr = self.lr
# Finetune the head of freeze_epochs > 0:
if freeze_epochs > 0:
print('Finetune the head')
self.freeze()
self.fit_one_cycle(freeze_epochs, lr_max=base_lr, pct_start=pct_start)
# Finetune the entire network if n_epochs > 0
if n_epochs > 0:
print('Finetune the entire network')
self.unfreeze()
self.fit_one_cycle(n_epochs, lr_max=base_lr/2, pct_start=pct_start)
def linear_probe(self, n_epochs, base_lr=None, pct_start=0.3):
"""
linear probing the pretrained model. The model is freeze except the head during finetuning
"""
assert (n_epochs>0), "n_epochs has to be > 0"
if not base_lr: base_lr = self.lr
print('Finetune the head')
self.freeze()
self.fit_one_cycle(n_epochs, lr_max=base_lr, pct_start=pct_start)
def lr_finder(self, start_lr=1e-7, end_lr=10, num_iter=100, step_mode='exp', show_plot=True, suggestion='valley'):
"""
find the learning rate
"""
n_epochs = num_iter//len(self.dls.train) + 1
# indicator of lr_finder method is applied
self.run_finder = True
# add LRFinderCB to callback list and will remove later
cb = LRFinderCB(start_lr, end_lr, num_iter, step_mode, suggestion=suggestion)
# fit
self.fit(n_epochs=n_epochs, cbs=cb, do_valid=False)
# should remove LRFinderCB callback after fitting
self.remove_callback(cb)
self.run_finder = False
if show_plot: cb.plot_lr_find()
if suggestion: return cb.suggested_lr
def freeze(self):
"""
freeze the model head
require the model to have head attribute
"""
if hasattr(get_model(self.model), 'head'):
# print('model head is available')
for param in get_model(self.model).parameters(): param.requires_grad = False
for param in get_model(self.model).head.parameters(): param.requires_grad = True
# print('model is frozen except the head')
def unfreeze(self):
for param in get_model(self.model).parameters(): param.requires_grad = True
def __call__(self, name):
for cb in self.cbs:
attr = getattr(cb, name)
if attr is not None: attr()
def save(self, fname, path, **kwargs):
"""
Save model and optimizer state (if `with_opt`) to `self.path/file`
"""
fname = join_path_file(fname, path, ext='.pth')
save_model(fname, self.model, getattr(self,'opt',None), **kwargs)
return fname
def load(self, fname, with_opt=False, device='cuda', strict=True, **kwargs):
"""
load the model
"""
if not torch.cuda.is_available():
device = "cpu"
load_model(fname, self.model, self.opt, with_opt, device=device, strict=strict)
def get_params(self, deep=True, **kwargs):
params = BaseEstimator.get_params(self, deep=deep, **kwargs)
return params
def _get_param_names(self):
return (k for k in self.__dict__ if not k.endswith('_'))
def set_params(self, **kwargs):
params = {}
for key, val in kwargs.items():
params[key] = val
BaseEstimator.set_params(self, **params)
def to_distributed(self,
sync_bn=True, # Whether to replace all batch norm with `nn.SyncBatchNorm`
**kwargs
):
local_rank = int(os.environ.get('LOCAL_RANK'))
world_size = int(os.environ.get('WORLD_SIZE'))
rank = int(os.environ.get('RANK'))
print('Process {} (out of {})'.format(
rank, torch.distributed.get_world_size()))
self.add_callback(DistributedTrainer(local_rank=local_rank, world_size=world_size, sync_bn=sync_bn, **kwargs))
return self
def save_model(path, model, opt, with_opt=True, pickle_protocol=2):
"Save `model` to `file` along with `opt` (if available, and if `with_opt`)"
if opt is None: with_opt=False
state = get_model(model).state_dict()
if with_opt: state = {'model': state, 'opt':opt.state_dict()}
torch.save(state, path, pickle_protocol=pickle_protocol)
def load_model(path, model, opt=None, with_opt=False, device='cpu', strict=True):
" load the saved model "
state = torch.load(path, map_location=device)
if not opt: with_opt=False
model_state = state['model'] if with_opt else state
get_model(model).load_state_dict(model_state, strict=strict)
if with_opt: opt.load_state_dict(state['opt'])
model = model.to(device)
def join_path_file(file, path, ext=''):
"Return `path/file` if file is a string or a `Path`, file otherwise"
if not isinstance(file, (str, Path)): return file
if not isinstance(path, Path): path = Path(path)
path.mkdir(parents=True, exist_ok=True)
return path/f'{file}{ext}'
def get_model(model):
"Return the model maybe wrapped inside `model`."
return model.module if isinstance(model, (DistributedDataParallel, nn.DataParallel)) else model
def transfer_weights(weights_path, model, exclude_head=True, device='cpu'):
# state_dict = model.state_dict()
new_state_dict = torch.load(weights_path, map_location=device)
#print('new_state_dict',new_state_dict)
matched_layers = 0
unmatched_layers = []
for name, param in model.state_dict().items():
if exclude_head and 'head' in name: continue
if name in new_state_dict:
matched_layers += 1
input_param = new_state_dict[name]
if input_param.shape == param.shape: param.copy_(input_param)
else: unmatched_layers.append(name)
else:
unmatched_layers.append(name)
pass # these are weights that weren't in the original model, such as a new head
if matched_layers == 0: raise Exception("No shared weight names were found between the models")
else:
if len(unmatched_layers) > 0:
print(f'check unmatched_layers: {unmatched_layers}')
else:
print(f"weights from {weights_path} successfully transferred!\n")
model = model.to(device)
return model
def update_callback(cb, list_cbs):
for cb_ in list_cbs:
if type(cb_) == type(cb): list_cbs.remove(cb_)
list_cbs += [cb]
return list_cbs
def update_callbacks(list_cbs, default_cbs):
for cb in list_cbs: default_cbs = update_callback(cb, default_cbs)
return default_cbs
def remove_callback(cb, list_cbs):
for cb_ in list_cbs:
if type(cb_) == type(cb):
list_cbs.remove(cb_)
break
return list_cbs, cb_
def get_layer_output(inp, model, layers=None, unwrap=False):
"""
layers is a list of module names
"""
orig_model = model
if unwrap: model = unwrap_model(model)
if not layers: layers = list(dict(model.named_children()).keys())
if not isinstance(layers, list): layers = [layers]
activation = {}
def getActivation(name):
# the hook signature
def hook(model, input, output):
activation[name] = output.detach().cpu().numpy()
return hook
# register forward hooks on the layers of choice
h_list = [getattr(model, layer).register_forward_hook(getActivation(layer)) for layer in layers]
model.eval()
out = orig_model(inp)
for h in h_list: h.remove()
return activation
@@ -0,0 +1,21 @@
import torch
from torch import Tensor
import torch.nn.functional as F
def mse(y_true, y_pred):
return F.mse_loss(y_true, y_pred, reduction='mean')
def rmse(y_true, y_pred):
return torch.sqrt(F.mse_loss(y_true, y_pred, reduction='mean'))
def mae(y_true, y_pred):
return F.l1_loss(y_true, y_pred, reduction='mean')
def r2_score(y_true, y_pred):
from sklearn.metrics import r2_score
return r2_score(y_true, y_pred)
def mape(y_true, y_pred):
from sklearn.metrics import mean_absolute_percentage_error
return mean_absolute_percentage_error(y_true, y_pred)
@@ -0,0 +1,115 @@
import torch
from torch import nn
from torch import Tensor
import torch.nn.functional as F
from typing import Callable, Optional
class MultiheadAttention(nn.Module):
def __init__(self, d_model, n_heads, d_k=None, d_v=None, res_attention=False, attn_dropout=0., proj_dropout=0., qkv_bias=True, lsa=False):
"""Multi Head Attention Layer
Input shape:
Q: [batch_size (bs) x max_q_len x d_model]
K, V: [batch_size (bs) x q_len x d_model]
mask: [q_len x q_len]
"""
super().__init__()
d_k = d_model // n_heads if d_k is None else d_k
d_v = d_model // n_heads if d_v is None else d_v
self.n_heads, self.d_k, self.d_v = n_heads, d_k, d_v
self.W_Q = nn.Linear(d_model, d_k * n_heads, bias=qkv_bias)
self.W_K = nn.Linear(d_model, d_k * n_heads, bias=qkv_bias)
self.W_V = nn.Linear(d_model, d_v * n_heads, bias=qkv_bias)
# Scaled Dot-Product Attention (multiple heads)
self.res_attention = res_attention
self.sdp_attn = ScaledDotProductAttention(d_model, n_heads, attn_dropout=attn_dropout, res_attention=self.res_attention, lsa=lsa)
# Poject output
self.to_out = nn.Sequential(nn.Linear(n_heads * d_v, d_model), nn.Dropout(proj_dropout))
def forward(self, Q:Tensor, K:Optional[Tensor]=None, V:Optional[Tensor]=None, prev:Optional[Tensor]=None,
key_padding_mask:Optional[Tensor]=None, attn_mask:Optional[Tensor]=None):
bs = Q.size(0)
if K is None: K = Q
if V is None: V = Q
# Linear (+ split in multiple heads)
q_s = self.W_Q(Q).view(bs, -1, self.n_heads, self.d_k).transpose(1,2) # q_s : [bs x n_heads x max_q_len x d_k]
k_s = self.W_K(K).view(bs, -1, self.n_heads, self.d_k).permute(0,2,3,1) # k_s : [bs x n_heads x d_k x q_len] - transpose(1,2) + transpose(2,3)
v_s = self.W_V(V).view(bs, -1, self.n_heads, self.d_v).transpose(1,2) # v_s : [bs x n_heads x q_len x d_v]
# Apply Scaled Dot-Product Attention (multiple heads)
if self.res_attention:
output, attn_weights, attn_scores = self.sdp_attn(q_s, k_s, v_s, prev=prev, key_padding_mask=key_padding_mask, attn_mask=attn_mask)
else:
output, attn_weights = self.sdp_attn(q_s, k_s, v_s, key_padding_mask=key_padding_mask, attn_mask=attn_mask)
# output: [bs x n_heads x q_len x d_v], attn: [bs x n_heads x q_len x q_len], scores: [bs x n_heads x max_q_len x q_len]
# back to the original inputs dimensions
output = output.transpose(1, 2).contiguous().view(bs, -1, self.n_heads * self.d_v) # output: [bs x q_len x n_heads * d_v]
output = self.to_out(output)
if self.res_attention: return output, attn_weights, attn_scores
else: return output, attn_weights
class ScaledDotProductAttention(nn.Module):
r"""Scaled Dot-Product Attention module (Attention is all you need by Vaswani et al., 2017) with optional residual attention from previous layer
(Realformer: Transformer likes residual attention by He et al, 2020) and locality self sttention (Vision Transformer for Small-Size Datasets
by Lee et al, 2021)"""
def __init__(self, d_model, n_heads, attn_dropout=0., res_attention=False, lsa=False):
super().__init__()
self.attn_dropout = nn.Dropout(attn_dropout)
self.res_attention = res_attention
head_dim = d_model // n_heads
self.scale = nn.Parameter(torch.tensor(head_dim ** -0.5), requires_grad=lsa)
self.lsa = lsa
def forward(self, q:Tensor, k:Tensor, v:Tensor, prev:Optional[Tensor]=None, key_padding_mask:Optional[Tensor]=None, attn_mask:Optional[Tensor]=None):
'''
Input shape:
q : [bs x n_heads x max_q_len x d_k]
k : [bs x n_heads x d_k x seq_len]
v : [bs x n_heads x seq_len x d_v]
prev : [bs x n_heads x q_len x seq_len]
key_padding_mask: [bs x seq_len]
attn_mask : [1 x seq_len x seq_len]
Output shape:
output: [bs x n_heads x q_len x d_v]
attn : [bs x n_heads x q_len x seq_len]
scores : [bs x n_heads x q_len x seq_len]
'''
# Scaled MatMul (q, k) - similarity scores for all pairs of positions in an input sequence
attn_scores = torch.matmul(q, k) * self.scale # attn_scores : [bs x n_heads x max_q_len x q_len]
# Add pre-softmax attention scores from the previous layer (optional)
if prev is not None: attn_scores = attn_scores + prev
# Attention mask (optional)
if attn_mask is not None: # attn_mask with shape [q_len x seq_len] - only used when q_len == seq_len
if attn_mask.dtype == torch.bool:
attn_scores.masked_fill_(attn_mask, -np.inf)
else:
attn_scores += attn_mask
# Key padding mask (optional)
if key_padding_mask is not None: # mask with shape [bs x q_len] (only when max_w_len == q_len)
attn_scores.masked_fill_(key_padding_mask.unsqueeze(1).unsqueeze(2), -np.inf)
# normalize the attention weights
attn_weights = F.softmax(attn_scores, dim=-1) # attn_weights : [bs x n_heads x max_q_len x q_len]
attn_weights = self.attn_dropout(attn_weights)
# compute the new values given the attention weights
output = torch.matmul(attn_weights, v) # output: [bs x n_heads x max_q_len x d_v]
if self.res_attention: return output, attn_weights, attn_scores
else: return output, attn_weights
@@ -0,0 +1,47 @@
__all__ = ['Transpose', 'LinBnDrop', 'SigmoidRange', 'sigmoid_range', 'get_activation_fn']
import torch
from torch import nn
class Transpose(nn.Module):
def __init__(self, *dims, contiguous=False):
super().__init__()
self.dims, self.contiguous = dims, contiguous
def forward(self, x):
if self.contiguous: return x.transpose(*self.dims).contiguous()
else: return x.transpose(*self.dims)
class SigmoidRange(nn.Module):
def __init__(self, low, high):
super().__init__()
self.low, self.high = low, high
# self.low, self.high = ranges
def forward(self, x):
# return sigmoid_range(x, self.low, self.high)
return torch.sigmoid(x) * (self.high - self.low) + self.low
class LinBnDrop(nn.Sequential):
"Module grouping `BatchNorm1d`, `Dropout` and `Linear` layers"
def __init__(self, n_in, n_out, bn=True, p=0., act=None, lin_first=False):
layers = [nn.BatchNorm2d(n_out if lin_first else n_in, ndim=1)] if bn else []
if p != 0: layers.append(nn.Dropout(p))
lin = [nn.Linear(n_in, n_out, bias=not bn)]
if act is not None: lin.append(act)
layers = lin+layers if lin_first else layers+lin
super().__init__(*layers)
def sigmoid_range(x, low, high):
"Sigmoid function with range `(low, high)`"
return torch.sigmoid(x) * (high - low) + low
def get_activation_fn(activation):
if callable(activation): return activation()
elif activation.lower() == "relu": return nn.ReLU()
elif activation.lower() == "gelu": return nn.GELU()
raise ValueError(f'{activation} is not available. You can use "relu", "gelu", or a callable')
@@ -0,0 +1,103 @@
import torch
from torch import nn
class LinearRegressionHead(nn.Module):
def __init__(self, n_vars, d_model, output_dim, head_dropout, y_range=None):
super().__init__()
self.y_range = y_range
self.flatten = nn.Flatten(start_dim=1)
self.dropout = nn.Dropout(head_dropout)
self.linear = nn.Linear(n_vars*d_model, output_dim)
def forward(self, x):
"""
x: [bs x nvars x d_model x num_patch]
output: [bs x output_dim]
"""
x = x[:,:,:,-1] # only consider the last item in the sequence, x: bs x nvars x d_model
x = self.flatten(x) # x: bs x nvars * d_model
x = self.dropout(x)
y = self.linear(x) # y: bs x output_dim
if self.y_range: y = SigmoidRange(*self.y_range)(y)
return y
class LinearClassificationHead(nn.Module):
def __init__(self, n_vars, d_model, n_classes, head_dropout):
super().__init__()
self.flatten = nn.Flatten(start_dim=1)
self.dropout = nn.Dropout(head_dropout)
self.linear = nn.Linear(n_vars*d_model, n_classes)
def forward(self, x):
"""
x: [bs x nvars x d_model x num_patch]
output: [bs x n_classes]
"""
x = x[:,:,:,-1] # only consider the last item in the sequence, x: bs x nvars x d_model
x = self.flatten(x) # x: bs x nvars * d_model
x = self.dropout(x)
y = self.linear(x) # y: bs x n_classes
return y
class LinearPredictionHead(nn.Module):
def __init__(self, individual, n_vars, d_model, num_patch, forecast_len, head_dropout=0):
super().__init__()
self.individual = individual
self.n_vars = n_vars
head_dim = d_model*num_patch
if self.individual:
self.linears = nn.ModuleList()
self.dropouts = nn.ModuleList()
self.flattens = nn.ModuleList()
for i in range(self.n_vars):
self.flattens.append(nn.Flatten(start_dim=-2))
self.linears.append(nn.Linear(head_dim, forecast_len))
self.dropouts.append(nn.Dropout(head_dropout))
else:
self.flatten = nn.Flatten(start_dim=-2)
self.linear = nn.Linear(head_dim, forecast_len)
self.dropout = nn.Dropout(head_dropout)
def forward(self, x):
"""
x: [bs x nvars x d_model x num_patch]
output: [bs x forecast_len x nvars]
"""
if self.individual:
x_out = []
for i in range(self.n_vars):
z = self.flattens[i](x[:,i,:,:]) # z: [bs x d_model * num_patch]
z = self.linears[i](z) # z: [bs x forecast_len]
z = self.dropouts[i](z)
x_out.append(z)
x = torch.stack(x_out, dim=1) # x: [bs x nvars x forecast_len]
else:
x = self.flatten(x)
x = self.dropout(x)
x = self.linear(x)
return x.transpose(2,1) # [bs x forecast_len x nvars]
class LinearPretrainHead(nn.Module):
def __init__(self, d_model, patch_len, dropout):
super().__init__()
self.dropout = nn.Dropout(dropout)
self.linear = nn.Linear(d_model, patch_len)
def forward(self, x):
"""
x: tensor [bs x nvars x d_model x num_patch]
output: tensor [bs x nvars x num_patch x patch_len]
"""
x = x.transpose(2,3) # [bs x nvars x num_patch x d_model]
x = self.linear( self.dropout(x) ) # [bs x nvars x num_patch x patch_len]
x = x.permute(0,2,1,3) # [bs x num_patch x nvars x patch_len]
return x
@@ -0,0 +1,46 @@
__all__ = ['PositionalEncoding', 'SinCosPosEncoding', 'positional_encoding']
# Cell
import torch
from torch import nn
import math
# Cell
def PositionalEncoding(q_len, d_model, normalize=True):
pe = torch.zeros(q_len, d_model)
position = torch.arange(0, q_len).unsqueeze(1)
div_term = torch.exp(torch.arange(0, d_model, 2) * -(math.log(10000.0) / d_model))
pe[:, 0::2] = torch.sin(position * div_term)
pe[:, 1::2] = torch.cos(position * div_term)
if normalize:
pe = pe - pe.mean()
pe = pe / (pe.std() * 10)
return pe
SinCosPosEncoding = PositionalEncoding
def positional_encoding(pe, learn_pe, q_len, d_model):
# Positional encoding
if pe == None:
W_pos = torch.empty((q_len, d_model)) # pe = None and learn_pe = False can be used to measure impact of pe
nn.init.uniform_(W_pos, -0.02, 0.02)
learn_pe = False
elif pe == 'zero':
W_pos = torch.empty((q_len, 1))
nn.init.uniform_(W_pos, -0.02, 0.02)
elif pe == 'zeros':
W_pos = torch.empty((q_len, d_model))
nn.init.uniform_(W_pos, -0.02, 0.02)
elif pe == 'normal' or pe == 'gauss':
W_pos = torch.zeros((q_len, 1))
torch.nn.init.normal_(W_pos, mean=0.0, std=0.1)
elif pe == 'uniform':
W_pos = torch.zeros((q_len, 1))
nn.init.uniform_(W_pos, a=0.0, b=0.1)
elif pe == 'sincos': W_pos = PositionalEncoding(q_len, d_model, normalize=True)
else: raise ValueError(f"{pe} is not a valid pe (positional encoder. Available types: 'gauss'=='normal', \
'zeros', 'zero', uniform', 'sincos', None.)")
return nn.Parameter(W_pos, requires_grad=learn_pe)
@@ -0,0 +1,51 @@
import torch
from torch import nn
class RevIN(nn.Module):
def __init__(self, num_features: int, eps=1e-5, affine=True):
"""
:param num_features: the number of features or channels
:param eps: a value added for numerical stability
:param affine: if True, RevIN has learnable affine parameters
"""
super(RevIN, self).__init__()
self.num_features = num_features
self.eps = eps
self.affine = affine
if self.affine:
self._init_params()
def forward(self, x, mode:str):
if mode == 'norm':
self._get_statistics(x)
x = self._normalize(x)
elif mode == 'denorm':
x = self._denormalize(x)
else: raise NotImplementedError
return x
def _init_params(self):
# initialize RevIN params: (C,)
self.affine_weight = nn.Parameter(torch.ones(self.num_features))
self.affine_bias = nn.Parameter(torch.zeros(self.num_features))
def _get_statistics(self, x):
dim2reduce = tuple(range(1, x.ndim-1))
self.mean = torch.mean(x, dim=dim2reduce, keepdim=True).detach()
self.stdev = torch.sqrt(torch.var(x, dim=dim2reduce, keepdim=True, unbiased=False) + self.eps).detach()
def _normalize(self, x):
x = x - self.mean
x = x / self.stdev
if self.affine:
x = x * self.affine_weight
x = x + self.affine_bias
return x
def _denormalize(self, x):
if self.affine:
x = x - self.affine_bias
x = x / (self.affine_weight + self.eps*self.eps)
x = x * self.stdev
x = x + self.mean
return x
@@ -0,0 +1,333 @@
__all__ = ['PatchTST']
# Cell
from typing import Callable, Optional
import torch
from torch import nn
from torch import Tensor
import torch.nn.functional as F
import numpy as np
from collections import OrderedDict
from ..models.layers.pos_encoding import *
from ..models.layers.basics import *
from ..models.layers.attention import *
# Cell
class PatchTST(nn.Module):
"""
Output dimension:
[bs x target_dim x nvars] for prediction
[bs x target_dim] for regression
[bs x target_dim] for classification
[bs x num_patch x n_vars x patch_len] for pretrain
"""
def __init__(self, c_in:int, target_dim:int, patch_len:int, stride:int, num_patch:int,
n_layers:int=3, d_model=128, n_heads=16, shared_embedding=True, d_ff:int=256,
norm:str='BatchNorm', attn_dropout:float=0., dropout:float=0., act:str="gelu",
res_attention:bool=True, pre_norm:bool=False, store_attn:bool=False,
pe:str='zeros', learn_pe:bool=True, head_dropout = 0,
head_type = "prediction", individual = False,
y_range:Optional[tuple]=None, verbose:bool=False, **kwargs):
super().__init__()
assert head_type in ['pretrain', 'prediction', 'regression', 'classification'], 'head type should be either pretrain, prediction, or regression'
# Backbone
self.backbone = PatchTSTEncoder(c_in, num_patch=num_patch, patch_len=patch_len,
n_layers=n_layers, d_model=d_model, n_heads=n_heads,
shared_embedding=shared_embedding, d_ff=d_ff,
attn_dropout=attn_dropout, dropout=dropout, act=act,
res_attention=res_attention, pre_norm=pre_norm, store_attn=store_attn,
pe=pe, learn_pe=learn_pe, verbose=verbose, **kwargs)
# Head
self.n_vars = c_in
self.head_type = head_type
if head_type == "pretrain":
self.head = PretrainHead(d_model, patch_len, head_dropout) # custom head passed as a partial func with all its kwargs
elif head_type == "prediction":
self.head = PredictionHead(individual, self.n_vars, d_model, num_patch, target_dim, head_dropout)
elif head_type == "regression":
self.head = RegressionHead(self.n_vars, d_model, target_dim, head_dropout, y_range)
elif head_type == "classification":
self.head = ClassificationHead(self.n_vars, d_model, target_dim, head_dropout)
def forward(self, z):
"""
z: tensor [bs x num_patch x n_vars x patch_len]
"""
z = self.backbone(z) # z: [bs x nvars x d_model x num_patch]
z = self.head(z)
# z: [bs x target_dim x nvars] for prediction
# [bs x target_dim] for regression
# [bs x target_dim] for classification
# [bs x num_patch x n_vars x patch_len] for pretrain
return z
class RegressionHead(nn.Module):
def __init__(self, n_vars, d_model, output_dim, head_dropout, y_range=None):
super().__init__()
self.y_range = y_range
self.flatten = nn.Flatten(start_dim=1)
self.dropout = nn.Dropout(head_dropout)
self.linear = nn.Linear(n_vars*d_model, output_dim)
def forward(self, x):
"""
x: [bs x nvars x d_model x num_patch]
output: [bs x output_dim]
"""
x = x[:,:,:,-1] # only consider the last item in the sequence, x: bs x nvars x d_model
x = self.flatten(x) # x: bs x nvars * d_model
x = self.dropout(x)
y = self.linear(x) # y: bs x output_dim
if self.y_range: y = SigmoidRange(*self.y_range)(y)
return y
class ClassificationHead(nn.Module):
def __init__(self, n_vars, d_model, n_classes, head_dropout):
super().__init__()
self.flatten = nn.Flatten(start_dim=1)
self.dropout = nn.Dropout(head_dropout)
self.linear = nn.Linear(n_vars*d_model, n_classes)
def forward(self, x):
"""
x: [bs x nvars x d_model x num_patch]
output: [bs x n_classes]
"""
x = x[:,:,:,-1] # only consider the last item in the sequence, x: bs x nvars x d_model
x = self.flatten(x) # x: bs x nvars * d_model
x = self.dropout(x)
y = self.linear(x) # y: bs x n_classes
return y
class PredictionHead(nn.Module):
def __init__(self, individual, n_vars, d_model, num_patch, forecast_len, head_dropout=0, flatten=False):
super().__init__()
self.individual = individual
self.n_vars = n_vars
self.flatten = flatten
head_dim = d_model*num_patch
if self.individual:
self.linears = nn.ModuleList()
self.dropouts = nn.ModuleList()
self.flattens = nn.ModuleList()
for i in range(self.n_vars):
self.flattens.append(nn.Flatten(start_dim=-2))
self.linears.append(nn.Linear(head_dim, forecast_len))
self.dropouts.append(nn.Dropout(head_dropout))
else:
self.flatten = nn.Flatten(start_dim=-2)
self.linear = nn.Linear(head_dim, forecast_len)
self.dropout = nn.Dropout(head_dropout)
def forward(self, x):
"""
x: [bs x nvars x d_model x num_patch]
output: [bs x forecast_len x nvars]
"""
if self.individual:
x_out = []
for i in range(self.n_vars):
z = self.flattens[i](x[:,i,:,:]) # z: [bs x d_model * num_patch]
z = self.linears[i](z) # z: [bs x forecast_len]
z = self.dropouts[i](z)
x_out.append(z)
x = torch.stack(x_out, dim=1) # x: [bs x nvars x forecast_len]
else:
x = self.flatten(x) # x: [bs x nvars x (d_model * num_patch)]
x = self.dropout(x)
x = self.linear(x) # x: [bs x nvars x forecast_len]
return x.transpose(2,1) # [bs x forecast_len x nvars]
class PretrainHead(nn.Module):
def __init__(self, d_model, patch_len, dropout):
super().__init__()
self.dropout = nn.Dropout(dropout)
self.linear = nn.Linear(d_model, patch_len)
def forward(self, x):
"""
x: tensor [bs x nvars x d_model x num_patch]
output: tensor [bs x nvars x num_patch x patch_len]
"""
x = x.transpose(2,3) # [bs x nvars x num_patch x d_model]
x = self.linear( self.dropout(x) ) # [bs x nvars x num_patch x patch_len]
x = x.permute(0,2,1,3) # [bs x num_patch x nvars x patch_len]
return x
class PatchTSTEncoder(nn.Module):
def __init__(self, c_in, num_patch, patch_len,
n_layers=3, d_model=128, n_heads=16, shared_embedding=True,
d_ff=256, norm='BatchNorm', attn_dropout=0., dropout=0., act="gelu", store_attn=False,
res_attention=True, pre_norm=False,
pe='zeros', learn_pe=True, verbose=False, **kwargs):
super().__init__()
self.n_vars = c_in
self.num_patch = num_patch
self.patch_len = patch_len
self.d_model = d_model
self.shared_embedding = shared_embedding
# Input encoding: projection of feature vectors onto a d-dim vector space
if not shared_embedding:
self.W_P = nn.ModuleList()
for _ in range(self.n_vars): self.W_P.append(nn.Linear(patch_len, d_model))
else:
self.W_P = nn.Linear(patch_len, d_model)
# Positional encoding
self.W_pos = positional_encoding(pe, learn_pe, num_patch, d_model)
# Residual dropout
self.dropout = nn.Dropout(dropout)
# Encoder
self.encoder = TSTEncoder(d_model, n_heads, d_ff=d_ff, norm=norm, attn_dropout=attn_dropout, dropout=dropout,
pre_norm=pre_norm, activation=act, res_attention=res_attention, n_layers=n_layers,
store_attn=store_attn)
def forward(self, x) -> Tensor:
"""
x: tensor [bs x num_patch x nvars x patch_len]
"""
bs, num_patch, n_vars, patch_len = x.shape
# Input encoding
if not self.shared_embedding:
x_out = []
for i in range(n_vars):
z = self.W_P[i](x[:,:,i,:])
x_out.append(z)
x = torch.stack(x_out, dim=2)
else:
x = self.W_P(x) # x: [bs x num_patch x nvars x d_model]
x = x.transpose(1,2) # x: [bs x nvars x num_patch x d_model]
u = torch.reshape(x, (bs*n_vars, num_patch, self.d_model) ) # u: [bs * nvars x num_patch x d_model]
u = self.dropout(u + self.W_pos) # u: [bs * nvars x num_patch x d_model]
# Encoder
z = self.encoder(u) # z: [bs * nvars x num_patch x d_model]
z = torch.reshape(z, (-1,n_vars, num_patch, self.d_model)) # z: [bs x nvars x num_patch x d_model]
z = z.permute(0,1,3,2) # z: [bs x nvars x d_model x num_patch]
return z
# Cell
class TSTEncoder(nn.Module):
def __init__(self, d_model, n_heads, d_ff=None,
norm='BatchNorm', attn_dropout=0., dropout=0., activation='gelu',
res_attention=False, n_layers=1, pre_norm=False, store_attn=False):
super().__init__()
self.layers = nn.ModuleList([TSTEncoderLayer(d_model, n_heads=n_heads, d_ff=d_ff, norm=norm,
attn_dropout=attn_dropout, dropout=dropout,
activation=activation, res_attention=res_attention,
pre_norm=pre_norm, store_attn=store_attn) for i in range(n_layers)])
self.res_attention = res_attention
def forward(self, src:Tensor):
"""
src: tensor [bs x q_len x d_model]
"""
output = src
scores = None
if self.res_attention:
for mod in self.layers: output, scores = mod(output, prev=scores)
return output
else:
for mod in self.layers: output = mod(output)
return output
class TSTEncoderLayer(nn.Module):
def __init__(self, d_model, n_heads, d_ff=256, store_attn=False,
norm='BatchNorm', attn_dropout=0, dropout=0., bias=True,
activation="gelu", res_attention=False, pre_norm=False):
super().__init__()
assert not d_model%n_heads, f"d_model ({d_model}) must be divisible by n_heads ({n_heads})"
d_k = d_model // n_heads
d_v = d_model // n_heads
# Multi-Head attention
self.res_attention = res_attention
self.self_attn = MultiheadAttention(d_model, n_heads, d_k, d_v, attn_dropout=attn_dropout, proj_dropout=dropout, res_attention=res_attention)
# Add & Norm
self.dropout_attn = nn.Dropout(dropout)
if "batch" in norm.lower():
self.norm_attn = nn.Sequential(Transpose(1,2), nn.BatchNorm1d(d_model), Transpose(1,2))
else:
self.norm_attn = nn.LayerNorm(d_model)
# Position-wise Feed-Forward
self.ff = nn.Sequential(nn.Linear(d_model, d_ff, bias=bias),
get_activation_fn(activation),
nn.Dropout(dropout),
nn.Linear(d_ff, d_model, bias=bias))
# Add & Norm
self.dropout_ffn = nn.Dropout(dropout)
if "batch" in norm.lower():
self.norm_ffn = nn.Sequential(Transpose(1,2), nn.BatchNorm1d(d_model), Transpose(1,2))
else:
self.norm_ffn = nn.LayerNorm(d_model)
self.pre_norm = pre_norm
self.store_attn = store_attn
def forward(self, src:Tensor, prev:Optional[Tensor]=None):
"""
src: tensor [bs x q_len x d_model]
"""
# Multi-Head attention sublayer
if self.pre_norm:
src = self.norm_attn(src)
## Multi-Head attention
if self.res_attention:
src2, attn, scores = self.self_attn(src, src, src, prev)
else:
src2, attn = self.self_attn(src, src, src)
if self.store_attn:
self.attn = attn
## Add & Norm
src = src + self.dropout_attn(src2) # Add: residual connection with residual dropout
if not self.pre_norm:
src = self.norm_attn(src)
# Feed-forward sublayer
if self.pre_norm:
src = self.norm_ffn(src)
## Position-wise Feed-Forward
src2 = self.ff(src)
## Add & Norm
src = src + self.dropout_ffn(src2) # Add: residual connection with residual dropout
if not self.pre_norm:
src = self.norm_ffn(src)
if self.res_attention:
return src, scores
else:
return src
@@ -0,0 +1,58 @@
from torch import nn
import collections
from collections import OrderedDict
import torch
import os
from datetime import timedelta
def init_ddp():
local_rank = int(os.environ.get('LOCAL_RANK'))
world_size = int(os.environ.get('WORLD_SIZE'))
rank = int(os.environ.get('RANK'))
torch.cuda.set_device(local_rank)
torch.distributed.init_process_group(
'nccl',
init_method='env://',
world_size=world_size,
rank=rank,
timeout=timedelta(seconds=600)
)
def nested_children(m: nn.Module):
children = dict(m.named_children())
output = {}
if children == {}:
# if module has no children; m is last child! :O
return m
else:
# look for children from children... to the last child!
for name, child in children.items():
try:
output[name] = nested_children(child)
except TypeError:
output[name] = nested_children(child)
return output
def flatten_dict(d, parent_key='', sep='_'):
items = []
for k, v in d.items():
new_key = parent_key + sep + k if parent_key else k
if isinstance(v, collections.MutableMapping):
items.extend(flatten_dict(v, new_key, sep=sep).items())
else:
items.append((new_key, v))
return dict(items)
def unwrap_model(model):
unwrapped_model = nested_children(model)
unwrapped_model = flatten_dict(unwrapped_model)
unwrapped_model = nn.Sequential(OrderedDict(unwrapped_model))
return unwrapped_model
@@ -0,0 +1,209 @@
# ALL scripts in this file come from Autoformer
if [ ! -d "./logs" ]; then
mkdir ./logs
fi
if [ ! -d "./logs/LongForecasting" ]; then
mkdir ./logs/LongForecasting
fi
random_seed=2021
model_name=Autoformer
for pred_len in 96 192 336 720
do
python -u run_longExp.py \
--random_seed $random_seed \
--is_training 1 \
--root_path ./dataset/ \
--data_path exchange_rate.csv \
--model_id exchange_96_$pred_len \
--model $model_name \
--data custom \
--features M \
--seq_len 96 \
--label_len 48 \
--pred_len $pred_len \
--e_layers 2 \
--d_layers 1 \
--factor 3 \
--enc_in 8 \
--dec_in 8 \
--c_out 8 \
--des 'Exp' \
--itr 1 \
--train_epochs 1 >logs/LongForecasting/$model_name'_exchange_rate_'$pred_len.log
python -u run_longExp.py \
--random_seed $random_seed \
--is_training 1 \
--root_path ./dataset/ \
--data_path electricity.csv \
--model_id electricity_96_$pred_len \
--model $model_name \
--data custom \
--features M \
--seq_len 96 \
--label_len 48 \
--pred_len $pred_len \
--e_layers 2 \
--d_layers 1 \
--factor 3 \
--enc_in 321 \
--dec_in 321 \
--c_out 321 \
--des 'Exp' \
--itr 1 >logs/LongForecasting/$model_name'_electricity_'$pred_len.log
python -u run_longExp.py \
--random_seed $random_seed \
--is_training 1 \
--root_path ./dataset/ \
--data_path traffic.csv \
--model_id traffic_96_$pred_len \
--model $model_name \
--data custom \
--features M \
--seq_len 96 \
--label_len 48 \
--pred_len $pred_len \
--e_layers 2 \
--d_layers 1 \
--factor 3 \
--enc_in 862 \
--dec_in 862 \
--c_out 862 \
--des 'Exp' \
--itr 1 \
--train_epochs 3 >logs/LongForecasting/$model_name'_traffic_'$pred_len.log
python -u run_longExp.py \
--random_seed $random_seed \
--is_training 1 \
--root_path ./dataset/ \
--data_path weather.csv \
--model_id weather_96_$pred_len \
--model $model_name \
--data custom \
--features M \
--seq_len 96 \
--label_len 48 \
--pred_len $pred_len \
--e_layers 2 \
--d_layers 1 \
--factor 3 \
--enc_in 21 \
--dec_in 21 \
--c_out 21 \
--des 'Exp' \
--itr 1 \
--train_epochs 2 >logs/LongForecasting/$model_name'_weather_'$pred_len.log
python -u run_longExp.py \
--random_seed $random_seed \
--is_training 1 \
--root_path ./dataset/ \
--data_path ETTh1.csv \
--model_id ETTh1_96_$pred_len \
--model $model_name \
--data ETTh1 \
--features M \
--seq_len 96 \
--label_len 48 \
--pred_len $pred_len \
--e_layers 2 \
--d_layers 1 \
--factor 3 \
--enc_in 7 \
--dec_in 7 \
--c_out 7 \
--des 'Exp' \
--itr 1 >logs/LongForecasting/$model_name'_Etth1_'$pred_len.log
python -u run_longExp.py \
--random_seed $random_seed \
--is_training 1 \
--root_path ./dataset/ \
--data_path ETTh2.csv \
--model_id ETTh2_96_$pred_len \
--model $model_name \
--data ETTh2 \
--features M \
--seq_len 96 \
--label_len 48 \
--pred_len $pred_len \
--e_layers 2 \
--d_layers 1 \
--factor 3 \
--enc_in 7 \
--dec_in 7 \
--c_out 7 \
--des 'Exp' \
--itr 1 >logs/LongForecasting/$model_name'_Etth2_'$pred_len.log
python -u run_longExp.py \
--random_seed $random_seed \
--is_training 1 \
--root_path ./dataset/ \
--data_path ETTm1.csv \
--model_id ETTm1_96_$pred_len \
--model $model_name \
--data ETTm1 \
--features M \
--seq_len 96 \
--label_len 48 \
--pred_len $pred_len \
--e_layers 2 \
--d_layers 1 \
--factor 3 \
--enc_in 7 \
--dec_in 7 \
--c_out 7 \
--des 'Exp' \
--itr 1 >logs/LongForecasting/$model_name'_Ettm1_'$pred_len.log
python -u run_longExp.py \
--random_seed $random_seed \
--is_training 1 \
--root_path ./dataset/ \
--data_path ETTm2.csv \
--model_id ETTm2_96_$pred_len \
--model $model_name \
--data ETTm2 \
--features M \
--seq_len 96 \
--label_len 48 \
--pred_len $pred_len \
--e_layers 2 \
--d_layers 1 \
--factor 3 \
--enc_in 7 \
--dec_in 7 \
--c_out 7 \
--des 'Exp' \
--itr 1 >logs/LongForecasting/$model_name'_Ettm2_'$pred_len.log
done
for pred_len in 24 36 48 60
do
python -u run_longExp.py \
--random_seed $random_seed \
--is_training 1 \
--root_path ./dataset/ \
--data_path national_illness.csv \
--model_id ili_36_$pred_len \
--model $model_name \
--data custom \
--features M \
--seq_len 36 \
--label_len 18 \
--pred_len $pred_len \
--e_layers 2 \
--d_layers 1 \
--factor 3 \
--enc_in 7 \
--dec_in 7 \
--c_out 7 \
--des 'Exp' \
--itr 1 >logs/LongForecasting/$model_name'_ili_'$pred_len.log
done
@@ -0,0 +1,21 @@
MIT License
Copyright (c) 2021 xxxx
Permission is hereby granted, free of charge, to any person obtaining a copy
of this software and associated documentation files (the "Software"), to deal
in the Software without restriction, including without limitation the rights
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
copies of the Software, and to permit persons to whom the Software is
furnished to do so, subject to the following conditions:
The above copyright notice and this permission notice shall be included in all
copies or substantial portions of the Software.
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
SOFTWARE.
@@ -0,0 +1,53 @@
# FEDformer
From https://github.com/MAZiqing/FEDformer
Frequency Enhanced Decomposed
Transformer (FEDformer) is more efficient than
standard Transformer with a linear complexity
to the sequence length.
Our empirical studies
with six benchmark datasets show that compared
with state-of-the-art methods, FEDformer can
reduce prediction error by 14.8% and 22.6%
for multivariate and univariate time series,
respectively.
## Get Started
1. Install Python 3.6, PyTorch 1.9.0.
2. Download data. You can obtain all the six benchmarks from xxxx.
3. Train the model. We provide the experiment scripts of all benchmarks under the folder `./scripts`. You can reproduce the experiment results by:
```bash
bash ./scripts/run_M.sh
bash ./scripts/run_S.sh
```
## Citation
If you find this repo useful, please cite our paper.
```
xxxxx
```
## Contact
If you have any question or want to use the code, please contact xxx@xxxx .
## Acknowledgement
We appreciate the following github repos a lot for their valuable code base or datasets:
https://github.com/thuml/Autoformer
https://github.com/zhouhaoyi/Informer2020
https://github.com/zhouhaoyi/ETDataset
https://github.com/laiguokun/multivariate-time-series-data
@@ -0,0 +1,52 @@
from data_provider.data_loader import Dataset_ETT_hour, Dataset_ETT_minute, Dataset_Custom,Dataset_sin
from torch.utils.data import DataLoader
data_dict = {
'ETTh1': Dataset_ETT_hour,
'ETTh2': Dataset_ETT_hour,
'ETTm1': Dataset_ETT_minute,
'ETTm2': Dataset_ETT_minute,
'custom': Dataset_Custom,
'sin':Dataset_sin,
}
def data_provider(args, flag):
Data = data_dict[args.data]
timeenc = 0 if args.embed != 'timeF' else 1
if flag == 'test':
shuffle_flag = False
drop_last = True
batch_size = args.batch_size
freq = args.freq
elif flag == 'pred':
shuffle_flag = False
drop_last = False
batch_size = 1
freq = args.detail_freq
Data = Dataset_Pred
else:
shuffle_flag = True
drop_last = True
batch_size = args.batch_size
freq = args.freq
data_set = Data(
root_path=args.root_path,
data_path=args.data_path,
flag=flag,
size=[args.seq_len, args.label_len, args.pred_len],
features=args.features,
target=args.target,
timeenc=timeenc,
freq=freq
)
print(flag, len(data_set))
data_loader = DataLoader(
data_set,
batch_size=batch_size,
shuffle=shuffle_flag,
num_workers=args.num_workers,
drop_last=drop_last)
return data_set, data_loader
@@ -0,0 +1,372 @@
import os
import numpy as np
import pandas as pd
import os
import torch
from torch.utils.data import Dataset, DataLoader
from sklearn.preprocessing import StandardScaler
from utils.timefeatures import time_features
import warnings
warnings.filterwarnings('ignore')
class Dataset_ETT_hour(Dataset):
def __init__(self, root_path, flag='train', size=None,
features='S', data_path='ETTh1.csv',
target='OT', scale=True, timeenc=0, freq='h'):
# size [seq_len, label_len, pred_len]
# info
if size == None:
self.seq_len = 24 * 4 * 4
self.label_len = 24 * 4
self.pred_len = 24 * 4
else:
self.seq_len = size[0]
self.label_len = size[1]
self.pred_len = size[2]
# init
assert flag in ['train', 'test', 'val']
type_map = {'train': 0, 'val': 1, 'test': 2}
self.set_type = type_map[flag]
self.features = features
self.target = target
self.scale = scale
self.timeenc = timeenc
self.freq = freq
self.root_path = root_path
self.data_path = data_path
self.__read_data__()
def __read_data__(self):
self.scaler = StandardScaler()
df_raw = pd.read_csv(os.path.join(self.root_path,
self.data_path))
border1s = [0, 12 * 30 * 24 - self.seq_len, 12 * 30 * 24 + 4 * 30 * 24 - self.seq_len]
border2s = [12 * 30 * 24, 12 * 30 * 24 + 4 * 30 * 24, 12 * 30 * 24 + 8 * 30 * 24]
border1 = border1s[self.set_type]
border2 = border2s[self.set_type]
if self.features == 'M' or self.features == 'MS':
cols_data = df_raw.columns[1:]
df_data = df_raw[cols_data]
elif self.features == 'S':
df_data = df_raw[[self.target]]
if self.scale:
train_data = df_data[border1s[0]:border2s[0]]
self.scaler.fit(train_data.values)
data = self.scaler.transform(df_data.values)
else:
data = df_data.values
df_stamp = df_raw[['date']][border1:border2]
df_stamp['date'] = pd.to_datetime(df_stamp.date)
if self.timeenc == 0:
df_stamp['month'] = df_stamp.date.apply(lambda row: row.month, 1)
df_stamp['day'] = df_stamp.date.apply(lambda row: row.day, 1)
df_stamp['weekday'] = df_stamp.date.apply(lambda row: row.weekday(), 1)
df_stamp['hour'] = df_stamp.date.apply(lambda row: row.hour, 1)
data_stamp = df_stamp.drop(['date'], axis=1).values
elif self.timeenc == 1:
data_stamp = time_features(pd.to_datetime(df_stamp['date'].values), freq=self.freq)
data_stamp = data_stamp.transpose(1, 0)
self.data_x = data[border1:border2]
self.data_y = data[border1:border2]
self.data_stamp = data_stamp
def __getitem__(self, index):
s_begin = index
s_end = s_begin + self.seq_len
r_begin = s_end - self.label_len
r_end = r_begin + self.label_len + self.pred_len
seq_x = self.data_x[s_begin:s_end]
seq_y = self.data_y[r_begin:r_end]
seq_x_mark = self.data_stamp[s_begin:s_end]
seq_y_mark = self.data_stamp[r_begin:r_end]
return seq_x, seq_y, seq_x_mark, seq_y_mark
def __len__(self):
return len(self.data_x) - self.seq_len - self.pred_len + 1
def inverse_transform(self, data):
return self.scaler.inverse_transform(data)
class Dataset_ETT_minute(Dataset):
def __init__(self, root_path, flag='train', size=None,
features='S', data_path='ETTm1.csv',
target='OT', scale=True, timeenc=0, freq='t'):
# size [seq_len, label_len, pred_len]
# info
if size == None:
self.seq_len = 24 * 4 * 4
self.label_len = 24 * 4
self.pred_len = 24 * 4
else:
self.seq_len = size[0]
self.label_len = size[1]
self.pred_len = size[2]
# init
assert flag in ['train', 'test', 'val']
type_map = {'train': 0, 'val': 1, 'test': 2}
self.set_type = type_map[flag]
self.features = features
self.target = target
self.scale = scale
self.timeenc = timeenc
self.freq = freq
self.root_path = root_path
self.data_path = data_path
self.__read_data__()
def __read_data__(self):
self.scaler = StandardScaler()
df_raw = pd.read_csv(os.path.join(self.root_path,
self.data_path))
border1s = [0, 12 * 30 * 24 * 4 - self.seq_len, 12 * 30 * 24 * 4 + 4 * 30 * 24 * 4 - self.seq_len]
border2s = [12 * 30 * 24 * 4, 12 * 30 * 24 * 4 + 4 * 30 * 24 * 4, 12 * 30 * 24 * 4 + 8 * 30 * 24 * 4]
border1 = border1s[self.set_type]
border2 = border2s[self.set_type]
if self.features == 'M' or self.features == 'MS':
cols_data = df_raw.columns[1:]
df_data = df_raw[cols_data]
elif self.features == 'S':
df_data = df_raw[[self.target]]
if self.scale:
train_data = df_data[border1s[0]:border2s[0]]
self.scaler.fit(train_data.values)
data = self.scaler.transform(df_data.values)
else:
data = df_data.values
df_stamp = df_raw[['date']][border1:border2]
df_stamp['date'] = pd.to_datetime(df_stamp.date)
if self.timeenc == 0:
df_stamp['month'] = df_stamp.date.apply(lambda row: row.month, 1)
df_stamp['day'] = df_stamp.date.apply(lambda row: row.day, 1)
df_stamp['weekday'] = df_stamp.date.apply(lambda row: row.weekday(), 1)
df_stamp['hour'] = df_stamp.date.apply(lambda row: row.hour, 1)
df_stamp['minute'] = df_stamp.date.apply(lambda row: row.minute, 1)
df_stamp['minute'] = df_stamp.minute.map(lambda x: x // 15)
data_stamp = df_stamp.drop(['date'], axis=1).values
elif self.timeenc == 1:
data_stamp = time_features(pd.to_datetime(df_stamp['date'].values), freq=self.freq)
data_stamp = data_stamp.transpose(1, 0)
self.data_x = data[border1:border2]
self.data_y = data[border1:border2]
self.data_stamp = data_stamp
def __getitem__(self, index):
s_begin = index
s_end = s_begin + self.seq_len
r_begin = s_end - self.label_len
r_end = r_begin + self.label_len + self.pred_len
seq_x = self.data_x[s_begin:s_end]
seq_y = self.data_y[r_begin:r_end]
seq_x_mark = self.data_stamp[s_begin:s_end]
seq_y_mark = self.data_stamp[r_begin:r_end]
return seq_x, seq_y, seq_x_mark, seq_y_mark
def __len__(self):
return len(self.data_x) - self.seq_len - self.pred_len + 1
def inverse_transform(self, data):
return self.scaler.inverse_transform(data)
class Dataset_Custom(Dataset):
def __init__(self, root_path, flag='train', size=None,
features='S', data_path='ETTh1.csv',
target='OT', scale=True, timeenc=0, freq='h'):
# size [seq_len, label_len, pred_len]
# info
if size == None:
self.seq_len = 24 * 4 * 4
self.label_len = 24 * 4
self.pred_len = 24 * 4
else:
self.seq_len = size[0]
self.label_len = size[1]
self.pred_len = size[2]
# init
assert flag in ['train', 'test', 'val']
type_map = {'train': 0, 'val': 1, 'test': 2}
self.set_type = type_map[flag]
self.features = features
self.target = target
self.scale = scale
self.timeenc = timeenc
self.freq = freq
self.root_path = root_path
self.data_path = data_path
self.__read_data__()
def __read_data__(self):
self.scaler = StandardScaler()
df_raw = pd.read_csv(os.path.join(self.root_path,
self.data_path))
'''
df_raw.columns: ['date', ...(other features), target feature]
'''
cols = list(df_raw.columns)
cols.remove(self.target)
cols.remove('date')
df_raw = df_raw[['date'] + cols + [self.target]]
# print(cols)
num_train = int(len(df_raw) * 0.7)
num_test = int(len(df_raw) * 0.2)
num_vali = len(df_raw) - num_train - num_test
border1s = [0, num_train - self.seq_len, len(df_raw) - num_test - self.seq_len]
border2s = [num_train, num_train + num_vali, len(df_raw)]
border1 = border1s[self.set_type]
border2 = border2s[self.set_type]
if self.features == 'M' or self.features == 'MS':
cols_data = df_raw.columns[1:]
df_data = df_raw[cols_data]
elif self.features == 'S':
df_data = df_raw[[self.target]]
if self.scale:
train_data = df_data[border1s[0]:border2s[0]]
self.scaler.fit(train_data.values)
data = self.scaler.transform(df_data.values)
else:
data = df_data.values
df_stamp = df_raw[['date']][border1:border2]
df_stamp['date'] = pd.to_datetime(df_stamp.date)
if self.timeenc == 0:
df_stamp['month'] = df_stamp.date.apply(lambda row: row.month, 1)
df_stamp['day'] = df_stamp.date.apply(lambda row: row.day, 1)
df_stamp['weekday'] = df_stamp.date.apply(lambda row: row.weekday(), 1)
df_stamp['hour'] = df_stamp.date.apply(lambda row: row.hour, 1)
data_stamp = df_stamp.drop(['date'], axis=1).values
elif self.timeenc == 1:
data_stamp = time_features(pd.to_datetime(df_stamp['date'].values), freq=self.freq)
data_stamp = data_stamp.transpose(1, 0)
self.data_x = data[border1:border2]
self.data_y = data[border1:border2]
self.data_stamp = data_stamp
def __getitem__(self, index):
s_begin = index
s_end = s_begin + self.seq_len
r_begin = s_end - self.label_len
r_end = r_begin + self.label_len + self.pred_len
seq_x = self.data_x[s_begin:s_end]
seq_y = self.data_y[r_begin:r_end]
seq_x_mark = self.data_stamp[s_begin:s_end]
seq_y_mark = self.data_stamp[r_begin:r_end]
return seq_x, seq_y, seq_x_mark, seq_y_mark
def __len__(self):
return len(self.data_x) - self.seq_len - self.pred_len + 1
def inverse_transform(self, data):
return self.scaler.inverse_transform(data)
class Dataset_sin(Dataset):
def __init__(self, root_path, flag='train', size=None,
features='S', data_path='sin.csv',
target='y', scale=True, timeenc=0, freq='h'):
# size [seq_len, label_len, pred_len]
# info
if size == None:
self.seq_len = 24 * 4 * 4
self.label_len = 24 * 4
self.pred_len = 24 * 4
else:
self.seq_len = size[0]
self.label_len = size[1]
self.pred_len = size[2]
# init
assert flag in ['train', 'test', 'val']
type_map = {'train': 0, 'val': 1, 'test': 2}
self.set_type = type_map[flag]
self.features = features
self.target = target
self.scale = scale
self.timeenc = timeenc
self.freq = freq
self.root_path = root_path
self.data_path = data_path
self.__read_data__()
def __read_data__(self):
self.scaler = StandardScaler()
df_raw = pd.read_csv(os.path.join(self.root_path,
self.data_path))
'''
df_raw.columns: ['date', ...(other features), target feature]
'''
cols = list(df_raw.columns)
print(cols)
cols.remove(self.target)
cols.remove('x')
df_raw = df_raw[['x'] + cols + [self.target]]
# print(cols)
num_train = int(len(df_raw) * 0.7)
num_test = int(len(df_raw) * 0.2)
num_vali = len(df_raw) - num_train - num_test
border1s = [0, num_train - self.seq_len, len(df_raw) - num_test - self.seq_len]
border2s = [num_train, num_train + num_vali, len(df_raw)]
border1 = border1s[self.set_type]
border2 = border2s[self.set_type]
df_data = df_raw[[self.target]]
if self.scale:
train_data = df_data[border1s[0]:border2s[0]]
self.scaler.fit(train_data.values)
data = self.scaler.transform(df_data.values)
else:
data = df_data.values
self.data_x = data[border1:border2]
self.data_y = data[border1:border2]
def __getitem__(self, index):
s_begin = index
s_end = s_begin + self.seq_len
r_begin = s_end - self.label_len
r_end = r_begin + self.label_len + self.pred_len
seq_x = self.data_x[s_begin:s_end]
seq_y = self.data_y[r_begin:r_end]
seq_x_mark = np.zeros_like(seq_x)
seq_y_mark = np.zeros_like(seq_y)
return seq_x, seq_y, seq_x_mark, seq_y_mark
def __len__(self):
return len(self.data_x) - self.seq_len - self.pred_len + 1
def inverse_transform(self, data):
return self.scaler.inverse_transform(data)
@@ -0,0 +1,37 @@
import os
import torch
import numpy as np
class Exp_Basic(object):
def __init__(self, args):
self.args = args
self.device = self._acquire_device()
self.model = self._build_model().to(self.device)
def _build_model(self):
raise NotImplementedError
return None
def _acquire_device(self):
if self.args.use_gpu:
os.environ["CUDA_VISIBLE_DEVICES"] = str(
self.args.gpu) if not self.args.use_multi_gpu else self.args.devices
device = torch.device('cuda:{}'.format(self.args.gpu))
print('Use GPU: cuda:{}'.format(self.args.gpu))
else:
device = torch.device('cpu')
print('Use CPU')
return device
def _get_data(self):
pass
def vali(self):
pass
def train(self):
pass
def test(self):
pass
@@ -0,0 +1,347 @@
from data_provider.data_factory import data_provider
from exp.exp_basic import Exp_Basic
from models import FEDformer, Autoformer, Informer, Transformer
from utils.tools import EarlyStopping, adjust_learning_rate, visual
from utils.metrics import metric
import numpy as np
import torch
import torch.nn as nn
from torch import optim
import os
import time
import warnings
import matplotlib.pyplot as plt
import numpy as np
import io
from scipy import stats
warnings.filterwarnings('ignore')
class Exp_Main(Exp_Basic):
def __init__(self, args):
super(Exp_Main, self).__init__(args)
def _build_model(self):
model_dict = {
'FEDformer': FEDformer,
'Autoformer': Autoformer,
'Transformer': Transformer,
'Informer': Informer,
}
model = model_dict[self.args.model].Model(self.args).float()
if self.args.use_multi_gpu and self.args.use_gpu:
model = nn.DataParallel(model, device_ids=self.args.device_ids)
return model
def _get_data(self, flag):
data_set, data_loader = data_provider(self.args, flag)
return data_set, data_loader
def _select_optimizer(self):
model_optim = optim.Adam(self.model.parameters(), lr=self.args.learning_rate)
return model_optim
def _select_criterion(self):
criterion = nn.MSELoss()
return criterion
def vali(self, vali_data, vali_loader, criterion):
total_loss = []
ks_test_96,ks_test_192,ks_test_336,ks_test_720,ks_test_96_back=[],[],[],[],[]
ks_result=[]
ks_test_96_raw,ks_test_192_raw,ks_test_336_raw,ks_test_720_raw,ks_test_96_back_raw=[],[],[],[],[]
self.model.eval()
input_len=720
with torch.no_grad():
for i, (batch_x, batch_y, batch_x_mark, batch_y_mark) in enumerate(vali_loader):
batch_x = batch_x.float().to(self.device)
batch_y = batch_y.float()
batch_x_mark = batch_x_mark.float().to(self.device)
batch_y_mark = batch_y_mark.float().to(self.device)
# decoder input
dec_inp = torch.zeros_like(batch_y[:, -self.args.pred_len:, :]).float()
dec_inp = torch.cat([batch_y[:, :self.args.label_len, :], dec_inp], dim=1).float().to(self.device)
# encoder - decoder
if self.args.use_amp:
with torch.cuda.amp.autocast():
if self.args.output_attention:
outputs = self.model(batch_x, batch_x_mark, dec_inp, batch_y_mark)[0]
else:
outputs = self.model(batch_x, batch_x_mark, dec_inp, batch_y_mark)
else:
if self.args.output_attention:
outputs = self.model(batch_x, batch_x_mark, dec_inp, batch_y_mark)[0]
else:
outputs = self.model(batch_x, batch_x_mark, dec_inp, batch_y_mark)
f_dim = -1 if self.args.features == 'MS' else 0
batch_y = batch_y[:, -self.args.pred_len:, f_dim:].to(self.device)
pred = outputs.detach().cpu()
true = batch_y.detach().cpu()
loss = criterion(pred, true)
total_loss.append(loss)
total_loss = np.average(total_loss)
self.model.train()
return total_loss
def train(self, setting):
train_data, train_loader = self._get_data(flag='train')
vali_data, vali_loader = self._get_data(flag='val')
test_data, test_loader = self._get_data(flag='test')
path = os.path.join(self.args.checkpoints, setting)
if not os.path.exists(path):
os.makedirs(path)
time_now = time.time()
train_steps = len(train_loader)
early_stopping = EarlyStopping(patience=self.args.patience, verbose=True)
model_optim = self._select_optimizer()
criterion = self._select_criterion()
if self.args.use_amp:
scaler = torch.cuda.amp.GradScaler()
for epoch in range(self.args.train_epochs):
iter_count = 0
train_loss = []
self.model.train()
epoch_time = time.time()
for i, (batch_x, batch_y, batch_x_mark, batch_y_mark) in enumerate(train_loader):
iter_count += 1
model_optim.zero_grad()
batch_x = batch_x.float().to(self.device)
batch_y = batch_y.float().to(self.device)
batch_x_mark = batch_x_mark.float().to(self.device)
batch_y_mark = batch_y_mark.float().to(self.device)
# decoder input
dec_inp = torch.zeros_like(batch_y[:, -self.args.pred_len:, :]).float()
dec_inp = torch.cat([batch_y[:, :self.args.label_len, :], dec_inp], dim=1).float().to(self.device)
# encoder - decoder
if self.args.use_amp:
with torch.cuda.amp.autocast():
if self.args.output_attention:
outputs = self.model(batch_x, batch_x_mark, dec_inp, batch_y_mark)[0]
else:
outputs = self.model(batch_x, batch_x_mark, dec_inp, batch_y_mark)
f_dim = -1 if self.args.features == 'MS' else 0
batch_y = batch_y[:, -self.args.pred_len:, f_dim:].to(self.device)
loss = criterion(outputs, batch_y)
train_loss.append(loss.item())
else:
if self.args.output_attention:
outputs = self.model(batch_x, batch_x_mark, dec_inp, batch_y_mark)[0]
else:
outputs = self.model(batch_x, batch_x_mark, dec_inp, batch_y_mark)
f_dim = -1 if self.args.features == 'MS' else 0
batch_y = batch_y[:, -self.args.pred_len:, f_dim:].to(self.device)
# if i==0:
# from scipy import stats
# pred = outputs.detach().cpu().numpy()
# true = batch_y.detach().cpu().numpy()
# input_data = batch_x.detach().cpu().numpy()
# plot_index1=np.arange(input_data.shape[1])
# plot_index2=np.arange(input_data.shape[1],input_data.shape[1]+pred.shape[1])
# plt.cla()
# plt.plot(plot_index1,input_data[0,:,-1:],label='input')
# plt.plot(plot_index2,pred[0,:,-1:],label="predict")
# plt.plot(plot_index2,true[0,:,-1:],label="true")
# print('KS test1',stats.kstest(input_data[0,:,-1:].reshape(-1),pred[0,-input_data.shape[1]:,-1:].reshape(-1)))
# plt.legend()
# #f = io.BytesIO()
# plt.savefig("sample0.png",format="png")
# #plt.clf()
# plt.cla()
# plt.plot(plot_index1,input_data[8,:,-1:],label='input')
# plt.plot(plot_index2,pred[8,:,-1:],label="predict")
# plt.plot(plot_index2,true[8,:,-1:],label="true")
# plt.legend()
# plt.savefig("sample1.png",format="png")
# print('KS test2',stats.kstest(input_data[8,:,-1:].reshape(-1),pred[8,-input_data.shape[1]:,-1:].reshape(-1)))
# raise Exception('aaa')
loss = criterion(outputs, batch_y)
train_loss.append(loss.item())
if (i + 1) % 100 == 0:
# print("\titers: {0}, epoch: {1} | loss: {2:.7f}".format(i + 1, epoch + 1, loss.item()))
speed = (time.time() - time_now) / iter_count
left_time = speed * ((self.args.train_epochs - epoch) * train_steps - i)
# print('\tspeed: {:.4f}s/iter; left time: {:.4f}s'.format(speed, left_time))
iter_count = 0
time_now = time.time()
if self.args.use_amp:
scaler.scale(loss).backward()
scaler.step(model_optim)
scaler.update()
else:
loss.backward()
model_optim.step()
print("Epoch: {} cost time: {}".format(epoch + 1, time.time() - epoch_time))
train_loss = np.average(train_loss)
vali_loss = self.vali(vali_data, vali_loader, criterion)
test_loss = self.vali(test_data, test_loader, criterion)
print("Epoch: {0}, Steps: {1} | Train Loss: {2:.7f} Vali Loss: {3:.7f} Test Loss: {4:.7f}".format(
epoch + 1, train_steps, train_loss, vali_loss, test_loss))
early_stopping(vali_loss, self.model, path)
if early_stopping.early_stop:
print("Early stopping")
break
adjust_learning_rate(model_optim, epoch + 1, self.args)
best_model_path = path + '/' + 'checkpoint.pth'
self.model.load_state_dict(torch.load(best_model_path))
return self.model
def test(self, setting, test=0):
test_data, test_loader = self._get_data(flag='test')
if test:
print('loading model')
self.model.load_state_dict(torch.load(os.path.join('./checkpoints/' + setting, 'checkpoint.pth')))
preds = []
trues = []
folder_path = './test_results/' + setting + '/'
if not os.path.exists(folder_path):
os.makedirs(folder_path)
self.model.eval()
with torch.no_grad():
for i, (batch_x, batch_y, batch_x_mark, batch_y_mark) in enumerate(test_loader):
batch_x = batch_x.float().to(self.device)
batch_y = batch_y.float().to(self.device)
batch_x_mark = batch_x_mark.float().to(self.device)
batch_y_mark = batch_y_mark.float().to(self.device)
# decoder input
dec_inp = torch.zeros_like(batch_y[:, -self.args.pred_len:, :]).float()
dec_inp = torch.cat([batch_y[:, :self.args.label_len, :], dec_inp], dim=1).float().to(self.device)
# encoder - decoder
if self.args.use_amp:
with torch.cuda.amp.autocast():
if self.args.output_attention:
outputs = self.model(batch_x, batch_x_mark, dec_inp, batch_y_mark)[0]
else:
outputs = self.model(batch_x, batch_x_mark, dec_inp, batch_y_mark)
else:
if self.args.output_attention:
outputs = self.model(batch_x, batch_x_mark, dec_inp, batch_y_mark)[0]
else:
outputs = self.model(batch_x, batch_x_mark, dec_inp, batch_y_mark)
f_dim = -1 if self.args.features == 'MS' else 0
batch_y = batch_y[:, -self.args.pred_len:, f_dim:].to(self.device)
outputs = outputs.detach().cpu().numpy()
batch_y = batch_y.detach().cpu().numpy()
pred = outputs # outputs.detach().cpu().numpy() # .squeeze()
true = batch_y # batch_y.detach().cpu().numpy() # .squeeze()
preds.append(pred)
trues.append(true)
if i % 20 == 0:
input = batch_x.detach().cpu().numpy()
gt = np.concatenate((input[0, :, -1], true[0, :, -1]), axis=0)
pd = np.concatenate((input[0, :, -1], pred[0, :, -1]), axis=0)
visual(gt, pd, os.path.join(folder_path, str(i) + '.pdf'))
preds = np.array(preds)
trues = np.array(trues)
print('test shape:', preds.shape, trues.shape)
preds = preds.reshape(-1, preds.shape[-2], preds.shape[-1])
trues = trues.reshape(-1, trues.shape[-2], trues.shape[-1])
print('test shape:', preds.shape, trues.shape)
# result save
folder_path = './results/' + setting + '/'
if not os.path.exists(folder_path):
os.makedirs(folder_path)
mae, mse, rmse, mape, mspe = metric(preds, trues)
print('mse:{}, mae:{}'.format(mse, mae))
f = open("result.txt", 'a')
f.write(setting + " \n")
f.write('mse:{}, mae:{}'.format(mse, mae))
f.write('\n')
f.write('\n')
f.close()
np.save(folder_path + 'metrics.npy', np.array([mae, mse, rmse, mape, mspe]))
np.save(folder_path + 'pred.npy', preds)
np.save(folder_path + 'true.npy', trues)
return
def predict(self, setting, load=False):
pred_data, pred_loader = self._get_data(flag='pred')
if load:
path = os.path.join(self.args.checkpoints, setting)
best_model_path = path + '/' + 'checkpoint.pth'
self.model.load_state_dict(torch.load(best_model_path))
preds = []
self.model.eval()
with torch.no_grad():
for i, (batch_x, batch_y, batch_x_mark, batch_y_mark) in enumerate(pred_loader):
batch_x = batch_x.float().to(self.device)
batch_y = batch_y.float()
batch_x_mark = batch_x_mark.float().to(self.device)
batch_y_mark = batch_y_mark.float().to(self.device)
# decoder input
dec_inp = torch.zeros_like(batch_y[:, -self.args.pred_len:, :]).float()
dec_inp = torch.cat([batch_y[:, :self.args.label_len, :], dec_inp], dim=1).float().to(self.device)
# encoder - decoder
if self.args.use_amp:
with torch.cuda.amp.autocast():
if self.args.output_attention:
outputs = self.model(batch_x, batch_x_mark, dec_inp, batch_y_mark)[0]
else:
outputs = self.model(batch_x, batch_x_mark, dec_inp, batch_y_mark)
else:
if self.args.output_attention:
outputs = self.model(batch_x, batch_x_mark, dec_inp, batch_y_mark)[0]
else:
outputs = self.model(batch_x, batch_x_mark, dec_inp, batch_y_mark)
pred = outputs.detach().cpu().numpy() # .squeeze()
preds.append(pred)
preds = np.array(preds)
preds = preds.reshape(-1, preds.shape[-2], preds.shape[-1])
# result save
folder_path = './results/' + setting + '/'
if not os.path.exists(folder_path):
os.makedirs(folder_path)
np.save(folder_path + 'real_prediction.npy', preds)
return
@@ -0,0 +1,222 @@
import time
import torch
import torch.nn as nn
import numpy as np
import math
from torch.nn.functional import interpolate
def decor_time(func):
def func2(*args, **kw):
now = time.time()
y = func(*args, **kw)
t = time.time() - now
print('call <{}>, time={}'.format(func.__name__, t))
return y
return func2
class AutoCorrelation(nn.Module):
"""
AutoCorrelation Mechanism with the following two phases:
(1) period-based dependencies discovery
(2) time delay aggregation
This block can replace the self-attention family mechanism seamlessly.
"""
def __init__(self, mask_flag=True, factor=1, scale=None, attention_dropout=0.1, output_attention=False, configs=None):
super(AutoCorrelation, self).__init__()
print('Autocorrelation used !')
self.factor = factor
self.scale = scale
self.mask_flag = mask_flag
self.output_attention = output_attention
self.dropout = nn.Dropout(attention_dropout)
self.agg = None
self.use_wavelet = configs.wavelet
# @decor_time
def time_delay_agg_training(self, values, corr):
"""
SpeedUp version of Autocorrelation (a batch-normalization style design)
This is for the training phase.
"""
head = values.shape[1]
channel = values.shape[2]
length = values.shape[3]
# find top k
top_k = int(self.factor * math.log(length))
mean_value = torch.mean(torch.mean(corr, dim=1), dim=1)
index = torch.topk(torch.mean(mean_value, dim=0), top_k, dim=-1)[1]
weights = torch.stack([mean_value[:, index[i]] for i in range(top_k)], dim=-1)
# update corr
tmp_corr = torch.softmax(weights, dim=-1)
# aggregation
tmp_values = values
delays_agg = torch.zeros_like(values).float()
for i in range(top_k):
pattern = torch.roll(tmp_values, -int(index[i]), -1)
delays_agg = delays_agg + pattern * \
(tmp_corr[:, i].unsqueeze(1).unsqueeze(1).unsqueeze(1).repeat(1, head, channel, length))
return delays_agg # size=[B, H, d, S]
def time_delay_agg_inference(self, values, corr):
"""
SpeedUp version of Autocorrelation (a batch-normalization style design)
This is for the inference phase.
"""
batch = values.shape[0]
head = values.shape[1]
channel = values.shape[2]
length = values.shape[3]
# index init
init_index = torch.arange(length).unsqueeze(0).unsqueeze(0).unsqueeze(0).repeat(batch, head, channel, 1).cuda()
# find top k
top_k = int(self.factor * math.log(length))
mean_value = torch.mean(torch.mean(corr, dim=1), dim=1)
weights = torch.topk(mean_value, top_k, dim=-1)[0]
delay = torch.topk(mean_value, top_k, dim=-1)[1]
# update corr
tmp_corr = torch.softmax(weights, dim=-1)
# aggregation
tmp_values = values.repeat(1, 1, 1, 2)
delays_agg = torch.zeros_like(values).float()
for i in range(top_k):
tmp_delay = init_index + delay[:, i].unsqueeze(1).unsqueeze(1).unsqueeze(1).repeat(1, head, channel, length)
pattern = torch.gather(tmp_values, dim=-1, index=tmp_delay)
delays_agg = delays_agg + pattern * \
(tmp_corr[:, i].unsqueeze(1).unsqueeze(1).unsqueeze(1).repeat(1, head, channel, length))
return delays_agg
def time_delay_agg_full(self, values, corr):
"""
Standard version of Autocorrelation
"""
batch = values.shape[0]
head = values.shape[1]
channel = values.shape[2]
length = values.shape[3]
# index init
init_index = torch.arange(length).unsqueeze(0).unsqueeze(0).unsqueeze(0).repeat(batch, head, channel, 1).cuda()
# find top k
top_k = int(self.factor * math.log(length))
weights = torch.topk(corr, top_k, dim=-1)[0]
delay = torch.topk(corr, top_k, dim=-1)[1]
# update corr
tmp_corr = torch.softmax(weights, dim=-1)
# aggregation
tmp_values = values.repeat(1, 1, 1, 2)
delays_agg = torch.zeros_like(values).float()
for i in range(top_k):
tmp_delay = init_index + delay[..., i].unsqueeze(-1)
pattern = torch.gather(tmp_values, dim=-1, index=tmp_delay)
delays_agg = delays_agg + pattern * (tmp_corr[..., i].unsqueeze(-1))
return delays_agg
def forward(self, queries, keys, values, attn_mask):
B, L, H, E = queries.shape
_, S, _, D = values.shape
if L > S:
zeros = torch.zeros_like(queries[:, :(L - S), :]).float()
values = torch.cat([values, zeros], dim=1)
keys = torch.cat([keys, zeros], dim=1)
else:
values = values[:, :L, :, :]
keys = keys[:, :L, :, :]
# period-based dependencies
if self.use_wavelet != 2:
if self.use_wavelet == 1:
j_list = self.j_list
queries = queries.reshape([B, L, -1])
keys = keys.reshape([B, L, -1])
Ql, Qh_list = self.dwt1d(queries.transpose(1, 2)) # [B, H*D, L]
Kl, Kh_list = self.dwt1d(keys.transpose(1, 2))
qs = [queries.transpose(1, 2)] + Qh_list + [Ql] # [B, H*D, L]
ks = [keys.transpose(1, 2)] + Kh_list + [Kl]
q_list = []
k_list = []
for q, k, j in zip(qs, ks, j_list):
q_list += [interpolate(q, scale_factor=j, mode='linear')[:, :, -L:]]
k_list += [interpolate(k, scale_factor=j, mode='linear')[:, :, -L:]]
queries = torch.stack([i.reshape([B, H, E, L]) for i in q_list], dim=3).reshape([B, H, -1, L]).permute(0, 3, 1, 2)
keys = torch.stack([i.reshape([B, H, E, L]) for i in k_list], dim=3).reshape([B, H, -1, L]).permute(0, 3, 1, 2)
else:
pass
q_fft = torch.fft.rfft(queries.permute(0, 2, 3, 1).contiguous(), dim=-1) # size=[B, H, E, L]
k_fft = torch.fft.rfft(keys.permute(0, 2, 3, 1).contiguous(), dim=-1)
res = q_fft * torch.conj(k_fft)
corr = torch.fft.irfft(res, dim=-1) # size=[B, H, E, L]
# time delay agg
if self.training:
V = self.time_delay_agg_training(values.permute(0, 2, 3, 1).contiguous(), corr).permute(0, 3, 1, 2) # [B, L, H, E], [B, H, E, L] -> [B, L, H, E]
else:
V = self.time_delay_agg_inference(values.permute(0, 2, 3, 1).contiguous(), corr).permute(0, 3, 1, 2)
else:
V_list = []
queries = queries.reshape([B, L, -1])
keys = keys.reshape([B, L, -1])
values = values.reshape([B, L, -1])
Ql, Qh_list = self.dwt1d(queries.transpose(1, 2)) # [B, H*D, L]
Kl, Kh_list = self.dwt1d(keys.transpose(1, 2))
Vl, Vh_list = self.dwt1d(values.transpose(1, 2))
qs = Qh_list + [Ql] # [B, H*D, L]
ks = Kh_list + [Kl]
vs = Vh_list + [Vl]
for q, k, v in zip(qs, ks, vs):
q = q.reshape([B, H, E, -1])
k = k.reshape([B, H, E, -1])
v = v.reshape([B, H, E, -1]).permute(0, 3, 1, 2)
q_fft = torch.fft.rfft(q.contiguous(), dim=-1)
k_fft = torch.fft.rfft(k.contiguous(), dim=-1)
res = q_fft * torch.conj(k_fft)
corr = torch.fft.irfft(res, dim=-1) # [B, H, E, L]
if self.training:
V = self.time_delay_agg_training(v.permute(0, 2, 3, 1).contiguous(), corr).permute(0, 3, 1, 2)
else:
V = self.time_delay_agg_inference(v.permute(0, 2, 3, 1).contiguous(), corr).permute(0, 3, 1, 2)
V_list += [V]
Vl = V_list[-1].reshape([B, -1, H*E]).transpose(1, 2)
Vh_list = [i.reshape([B, -1, H*E]).transpose(1, 2) for i in V_list[:-1]]
V = self.dwt1div((Vl, Vh_list)).reshape([B, H, E, -1]).permute(0, 3, 1, 2)
# corr = self.dwt1div((V_list[-1], V_list[:-1]))
if self.output_attention:
return (V.contiguous(), corr.permute(0, 3, 1, 2)) # size = [B, L, H, E]
else:
return (V.contiguous(), None)
class AutoCorrelationLayer(nn.Module):
def __init__(self, correlation, d_model, n_heads, d_keys=None,
d_values=None):
super(AutoCorrelationLayer, self).__init__()
d_keys = d_keys or (d_model // n_heads)
d_values = d_values or (d_model // n_heads)
self.inner_correlation = correlation
self.query_projection = nn.Linear(d_model, d_keys * n_heads)
self.key_projection = nn.Linear(d_model, d_keys * n_heads)
self.value_projection = nn.Linear(d_model, d_values * n_heads)
self.out_projection = nn.Linear(d_values * n_heads, d_model)
self.n_heads = n_heads
def forward(self, queries, keys, values, attn_mask):
B, L, _ = queries.shape
_, S, _ = keys.shape
H = self.n_heads
queries = self.query_projection(queries).view(B, L, H, -1)
keys = self.key_projection(keys).view(B, S, H, -1)
values = self.value_projection(values).view(B, S, H, -1)
out, attn = self.inner_correlation(
queries,
keys,
values,
attn_mask
)
out = out.view(B, L, -1)
return self.out_projection(out), attn
@@ -0,0 +1,219 @@
import torch
import torch.nn as nn
import torch.nn.functional as F
import math
from layers.SelfAttention_Family import FullAttention
class my_Layernorm(nn.Module):
"""
Special designed layernorm for the seasonal part
"""
def __init__(self, channels):
super(my_Layernorm, self).__init__()
self.layernorm = nn.LayerNorm(channels)
def forward(self, x):
x_hat = self.layernorm(x)
bias = torch.mean(x_hat, dim=1).unsqueeze(1).repeat(1, x.shape[1], 1)
return x_hat - bias
class moving_avg(nn.Module):
"""
Moving average block to highlight the trend of time series
"""
def __init__(self, kernel_size, stride):
super(moving_avg, self).__init__()
self.kernel_size = kernel_size
self.avg = nn.AvgPool1d(kernel_size=kernel_size, stride=stride, padding=0)
def forward(self, x):
# padding on the both ends of time series
front = x[:, 0:1, :].repeat(1, self.kernel_size - 1-math.floor((self.kernel_size - 1) // 2), 1)
end = x[:, -1:, :].repeat(1, math.floor((self.kernel_size - 1) // 2), 1)
x = torch.cat([front, x, end], dim=1)
x = self.avg(x.permute(0, 2, 1))
x = x.permute(0, 2, 1)
return x
class series_decomp(nn.Module):
"""
Series decomposition block
"""
def __init__(self, kernel_size):
super(series_decomp, self).__init__()
self.moving_avg = moving_avg(kernel_size, stride=1)
def forward(self, x):
moving_mean = self.moving_avg(x)
res = x - moving_mean
return res, moving_mean
class series_decomp_multi(nn.Module):
"""
Series decomposition block
"""
def __init__(self, kernel_size):
super(series_decomp_multi, self).__init__()
self.moving_avg = [moving_avg(kernel, stride=1) for kernel in kernel_size]
self.layer = torch.nn.Linear(1, len(kernel_size))
def forward(self, x):
moving_mean=[]
for func in self.moving_avg:
moving_avg = func(x)
moving_mean.append(moving_avg.unsqueeze(-1))
moving_mean=torch.cat(moving_mean,dim=-1)
moving_mean = torch.sum(moving_mean*nn.Softmax(-1)(self.layer(x.unsqueeze(-1))),dim=-1)
res = x - moving_mean
return res, moving_mean
class FourierDecomp(nn.Module):
def __init__(self):
super(FourierDecomp, self).__init__()
pass
def forward(self, x):
x_ft = torch.fft.rfft(x, dim=-1)
class EncoderLayer(nn.Module):
"""
Autoformer encoder layer with the progressive decomposition architecture
"""
def __init__(self, attention, d_model, d_ff=None, moving_avg=25, dropout=0.1, activation="relu"):
super(EncoderLayer, self).__init__()
d_ff = d_ff or 4 * d_model
self.attention = attention
self.conv1 = nn.Conv1d(in_channels=d_model, out_channels=d_ff, kernel_size=1, bias=False)
self.conv2 = nn.Conv1d(in_channels=d_ff, out_channels=d_model, kernel_size=1, bias=False)
if isinstance(moving_avg, list):
self.decomp1 = series_decomp_multi(moving_avg)
self.decomp2 = series_decomp_multi(moving_avg)
else:
self.decomp1 = series_decomp(moving_avg)
self.decomp2 = series_decomp(moving_avg)
self.dropout = nn.Dropout(dropout)
self.activation = F.relu if activation == "relu" else F.gelu
def forward(self, x, attn_mask=None):
new_x, attn = self.attention(
x, x, x,
attn_mask=attn_mask
)
x = x + self.dropout(new_x)
x, _ = self.decomp1(x)
y = x
y = self.dropout(self.activation(self.conv1(y.transpose(-1, 1))))
y = self.dropout(self.conv2(y).transpose(-1, 1))
res, _ = self.decomp2(x + y)
return res, attn
class Encoder(nn.Module):
"""
Autoformer encoder
"""
def __init__(self, attn_layers, conv_layers=None, norm_layer=None):
super(Encoder, self).__init__()
self.attn_layers = nn.ModuleList(attn_layers)
self.conv_layers = nn.ModuleList(conv_layers) if conv_layers is not None else None
self.norm = norm_layer
def forward(self, x, attn_mask=None):
attns = []
if self.conv_layers is not None:
for attn_layer, conv_layer in zip(self.attn_layers, self.conv_layers):
x, attn = attn_layer(x, attn_mask=attn_mask)
x = conv_layer(x)
attns.append(attn)
x, attn = self.attn_layers[-1](x)
attns.append(attn)
else:
for attn_layer in self.attn_layers:
x, attn = attn_layer(x, attn_mask=attn_mask)
attns.append(attn)
if self.norm is not None:
x = self.norm(x)
return x, attns
class DecoderLayer(nn.Module):
"""
Autoformer decoder layer with the progressive decomposition architecture
"""
def __init__(self, self_attention, cross_attention, d_model, c_out, d_ff=None,
moving_avg=25, dropout=0.1, activation="relu"):
super(DecoderLayer, self).__init__()
d_ff = d_ff or 4 * d_model
self.self_attention = self_attention
self.cross_attention = cross_attention
self.conv1 = nn.Conv1d(in_channels=d_model, out_channels=d_ff, kernel_size=1, bias=False)
self.conv2 = nn.Conv1d(in_channels=d_ff, out_channels=d_model, kernel_size=1, bias=False)
if isinstance(moving_avg, list):
self.decomp1 = series_decomp_multi(moving_avg)
self.decomp2 = series_decomp_multi(moving_avg)
self.decomp3 = series_decomp_multi(moving_avg)
else:
self.decomp1 = series_decomp(moving_avg)
self.decomp2 = series_decomp(moving_avg)
self.decomp3 = series_decomp(moving_avg)
self.dropout = nn.Dropout(dropout)
self.projection = nn.Conv1d(in_channels=d_model, out_channels=c_out, kernel_size=3, stride=1, padding=1,
padding_mode='circular', bias=False)
self.activation = F.relu if activation == "relu" else F.gelu
def forward(self, x, cross, x_mask=None, cross_mask=None):
x = x + self.dropout(self.self_attention(
x, x, x,
attn_mask=x_mask
)[0])
x, trend1 = self.decomp1(x)
x = x + self.dropout(self.cross_attention(
x, cross, cross,
attn_mask=cross_mask
)[0])
x, trend2 = self.decomp2(x)
y = x
y = self.dropout(self.activation(self.conv1(y.transpose(-1, 1))))
y = self.dropout(self.conv2(y).transpose(-1, 1))
x, trend3 = self.decomp3(x + y)
residual_trend = trend1 + trend2 + trend3
residual_trend = self.projection(residual_trend.permute(0, 2, 1)).transpose(1, 2)
return x, residual_trend
class Decoder(nn.Module):
"""
Autoformer encoder
"""
def __init__(self, layers, norm_layer=None, projection=None):
super(Decoder, self).__init__()
self.layers = nn.ModuleList(layers)
self.norm = norm_layer
self.projection = projection
def forward(self, x, cross, x_mask=None, cross_mask=None, trend=None):
for layer in self.layers:
x, residual_trend = layer(x, cross, x_mask=x_mask, cross_mask=cross_mask)
trend = trend + residual_trend
if self.norm is not None:
x = self.norm(x)
if self.projection is not None:
x = self.projection(x)
return x, trend
@@ -0,0 +1,178 @@
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.nn.utils import weight_norm
import math
class PositionalEmbedding(nn.Module):
def __init__(self, d_model, max_len=5000):
super(PositionalEmbedding, self).__init__()
# Compute the positional encodings once in log space.
pe = torch.zeros(max_len, d_model).float()
pe.require_grad = False
position = torch.arange(0, max_len).float().unsqueeze(1)
div_term = (torch.arange(0, d_model, 2).float() * -(math.log(10000.0) / d_model)).exp()
pe[:, 0::2] = torch.sin(position * div_term)
pe[:, 1::2] = torch.cos(position * div_term)
pe = pe.unsqueeze(0)
self.register_buffer('pe', pe)
def forward(self, x):
return self.pe[:, :x.size(1)]
class TokenEmbedding(nn.Module):
def __init__(self, c_in, d_model):
super(TokenEmbedding, self).__init__()
padding = 1 if torch.__version__ >= '1.5.0' else 2
self.tokenConv = nn.Conv1d(in_channels=c_in, out_channels=d_model,
kernel_size=3, padding=padding, padding_mode='circular', bias=False)
for m in self.modules():
if isinstance(m, nn.Conv1d):
nn.init.kaiming_normal_(m.weight, mode='fan_in', nonlinearity='leaky_relu')
def forward(self, x):
x = self.tokenConv(x.permute(0, 2, 1)).transpose(1, 2)
return x
class FixedEmbedding(nn.Module):
def __init__(self, c_in, d_model):
super(FixedEmbedding, self).__init__()
w = torch.zeros(c_in, d_model).float()
w.require_grad = False
position = torch.arange(0, c_in).float().unsqueeze(1)
div_term = (torch.arange(0, d_model, 2).float() * -(math.log(10000.0) / d_model)).exp()
w[:, 0::2] = torch.sin(position * div_term)
w[:, 1::2] = torch.cos(position * div_term)
self.emb = nn.Embedding(c_in, d_model)
self.emb.weight = nn.Parameter(w, requires_grad=False)
def forward(self, x):
return self.emb(x).detach()
class TemporalEmbedding(nn.Module):
def __init__(self, d_model, embed_type='fixed', freq='h'):
super(TemporalEmbedding, self).__init__()
minute_size = 4
hour_size = 24
weekday_size = 7
day_size = 32
month_size = 13
Embed = FixedEmbedding if embed_type == 'fixed' else nn.Embedding
if freq == 't':
self.minute_embed = Embed(minute_size, d_model)
self.hour_embed = Embed(hour_size, d_model)
self.weekday_embed = Embed(weekday_size, d_model)
self.day_embed = Embed(day_size, d_model)
self.month_embed = Embed(month_size, d_model)
def forward(self, x):
x = x.long()
minute_x = self.minute_embed(x[:, :, 4]) if hasattr(self, 'minute_embed') else 0.
hour_x = self.hour_embed(x[:, :, 3])
weekday_x = self.weekday_embed(x[:, :, 2])
day_x = self.day_embed(x[:, :, 1])
month_x = self.month_embed(x[:, :, 0])
return hour_x + weekday_x + day_x + month_x + minute_x
class TimeFeatureEmbedding(nn.Module):
def __init__(self, d_model, embed_type='timeF', freq='h'):
super(TimeFeatureEmbedding, self).__init__()
freq_map = {'h': 4, 't': 5, 's': 6, 'm': 1, 'a': 1, 'w': 2, 'd': 3, 'b': 3}
d_inp = freq_map[freq]
self.embed = nn.Linear(d_inp, d_model, bias=False)
def forward(self, x):
return self.embed(x)
class DataEmbedding(nn.Module):
def __init__(self, c_in, d_model, embed_type='fixed', freq='h', dropout=0.1):
super(DataEmbedding, self).__init__()
self.value_embedding = TokenEmbedding(c_in=c_in, d_model=d_model)
self.position_embedding = PositionalEmbedding(d_model=d_model)
self.temporal_embedding = TemporalEmbedding(d_model=d_model, embed_type=embed_type,
freq=freq) if embed_type != 'timeF' else TimeFeatureEmbedding(
d_model=d_model, embed_type=embed_type, freq=freq)
self.dropout = nn.Dropout(p=dropout)
def forward(self, x, x_mark):
x = self.value_embedding(x) + self.temporal_embedding(x_mark) + self.position_embedding(x)
return self.dropout(x)
class DataEmbedding_onlypos(nn.Module):
def __init__(self, c_in, d_model, embed_type='fixed', freq='h', dropout=0.1):
super(DataEmbedding_onlypos, self).__init__()
self.value_embedding = TokenEmbedding(c_in=c_in, d_model=d_model)
self.position_embedding = PositionalEmbedding(d_model=d_model)
self.dropout = nn.Dropout(p=dropout)
def forward(self, x, x_mark):
x = self.value_embedding(x) + self.position_embedding(x)
return self.dropout(x)
class DataEmbedding_wo_pos(nn.Module):
def __init__(self, c_in, d_model, embed_type='fixed', freq='h', dropout=0.1):
super(DataEmbedding_wo_pos, self).__init__()
self.value_embedding = TokenEmbedding(c_in=c_in, d_model=d_model)
self.position_embedding = PositionalEmbedding(d_model=d_model)
self.temporal_embedding = TemporalEmbedding(d_model=d_model, embed_type=embed_type,
freq=freq) if embed_type != 'timeF' else TimeFeatureEmbedding(
d_model=d_model, embed_type=embed_type, freq=freq)
self.dropout = nn.Dropout(p=dropout)
def forward(self, x, x_mark):
# try:
x = self.value_embedding(x) + self.temporal_embedding(x_mark)
# except:
# a = 1
return self.dropout(x)
class DataEmbedding_wo_pos_temp(nn.Module):
def __init__(self, c_in, d_model, embed_type='fixed', freq='h', dropout=0.1):
super(DataEmbedding_wo_pos_temp, self).__init__()
self.value_embedding = TokenEmbedding(c_in=c_in, d_model=d_model)
self.position_embedding = PositionalEmbedding(d_model=d_model)
self.temporal_embedding = TemporalEmbedding(d_model=d_model, embed_type=embed_type,
freq=freq) if embed_type != 'timeF' else TimeFeatureEmbedding(
d_model=d_model, embed_type=embed_type, freq=freq)
self.dropout = nn.Dropout(p=dropout)
def forward(self, x, x_mark):
x = self.value_embedding(x)
return self.dropout(x)
class DataEmbedding_wo_temp(nn.Module):
def __init__(self, c_in, d_model, embed_type='fixed', freq='h', dropout=0.1):
super(DataEmbedding_wo_temp, self).__init__()
self.value_embedding = TokenEmbedding(c_in=c_in, d_model=d_model)
self.position_embedding = PositionalEmbedding(d_model=d_model)
self.temporal_embedding = TemporalEmbedding(d_model=d_model, embed_type=embed_type,
freq=freq) if embed_type != 'timeF' else TimeFeatureEmbedding(
d_model=d_model, embed_type=embed_type, freq=freq)
self.dropout = nn.Dropout(p=dropout)
def forward(self, x, x_mark):
x = self.value_embedding(x) + self.position_embedding(x)
return self.dropout(x)
@@ -0,0 +1,137 @@
# coding=utf-8
# author=maziqing
# email=maziqing.mzq@alibaba-inc.com
import numpy as np
import torch
import torch.nn as nn
def get_frequency_modes(seq_len, modes=64, mode_select_method='random'):
"""
get modes on frequency domain:
'random' means sampling randomly;
'else' means sampling the lowest modes;
"""
modes = min(modes, seq_len//2)
if mode_select_method == 'random':
index = list(range(0, seq_len // 2))
np.random.shuffle(index)
index = index[:modes]
else:
index = list(range(0, modes))
index.sort()
return index
# ########## fourier layer #############
class FourierBlock(nn.Module):
def __init__(self, in_channels, out_channels, seq_len, modes=0, mode_select_method='random'):
super(FourierBlock, self).__init__()
print('fourier enhanced block used!')
"""
1D Fourier block. It performs representation learning on frequency domain,
it does FFT, linear transform, and Inverse FFT.
"""
# get modes on frequency domain
self.index = get_frequency_modes(seq_len, modes=modes, mode_select_method=mode_select_method)
print('modes={}, index={}'.format(modes, self.index))
self.scale = (1 / (in_channels * out_channels))
self.weights1 = nn.Parameter(
self.scale * torch.rand(8, in_channels // 8, out_channels // 8, len(self.index), dtype=torch.cfloat))
# Complex multiplication
def compl_mul1d(self, input, weights):
# (batch, in_channel, x ), (in_channel, out_channel, x) -> (batch, out_channel, x)
return torch.einsum("bhi,hio->bho", input, weights)
def forward(self, q, k, v, mask):
# size = [B, L, H, E]
B, L, H, E = q.shape
x = q.permute(0, 2, 3, 1)
# Compute Fourier coefficients
x_ft = torch.fft.rfft(x, dim=-1)
# Perform Fourier neural operations
out_ft = torch.zeros(B, H, E, L // 2 + 1, device=x.device, dtype=torch.cfloat)
for wi, i in enumerate(self.index):
if i >= x_ft.shape[3] or wi >= out_ft.shape[3]:
continue
out_ft[:, :, :, wi] = self.compl_mul1d(x_ft[:, :, :, i], self.weights1[:, :, :, wi])
# Return to time domain
x = torch.fft.irfft(out_ft, n=x.size(-1))
return (x, None)
# ########## Fourier Cross Former ####################
class FourierCrossAttention(nn.Module):
def __init__(self, in_channels, out_channels, seq_len_q, seq_len_kv, modes=64, mode_select_method='random',
activation='tanh', policy=0):
super(FourierCrossAttention, self).__init__()
print(' fourier enhanced cross attention used!')
"""
1D Fourier Cross Attention layer. It does FFT, linear transform, attention mechanism and Inverse FFT.
"""
self.activation = activation
self.in_channels = in_channels
self.out_channels = out_channels
# get modes for queries and keys (& values) on frequency domain
self.index_q = get_frequency_modes(seq_len_q, modes=modes, mode_select_method=mode_select_method)
self.index_kv = get_frequency_modes(seq_len_kv, modes=modes, mode_select_method=mode_select_method)
print('modes_q={}, index_q={}'.format(len(self.index_q), self.index_q))
print('modes_kv={}, index_kv={}'.format(len(self.index_kv), self.index_kv))
self.scale = (1 / (in_channels * out_channels))
self.weights1 = nn.Parameter(
self.scale * torch.rand(8, in_channels // 8, out_channels // 8, len(self.index_q), dtype=torch.cfloat))
# Complex multiplication
def compl_mul1d(self, input, weights):
# (batch, in_channel, x ), (in_channel, out_channel, x) -> (batch, out_channel, x)
return torch.einsum("bhi,hio->bho", input, weights)
def forward(self, q, k, v, mask):
# size = [B, L, H, E]
B, L, H, E = q.shape
xq = q.permute(0, 2, 3, 1) # size = [B, H, E, L]
xk = k.permute(0, 2, 3, 1)
xv = v.permute(0, 2, 3, 1)
# Compute Fourier coefficients
xq_ft_ = torch.zeros(B, H, E, len(self.index_q), device=xq.device, dtype=torch.cfloat)
xq_ft = torch.fft.rfft(xq, dim=-1)
for i, j in enumerate(self.index_q):
if j >= xq_ft.shape[3]:
continue
xq_ft_[:, :, :, i] = xq_ft[:, :, :, j]
xk_ft_ = torch.zeros(B, H, E, len(self.index_kv), device=xq.device, dtype=torch.cfloat)
xk_ft = torch.fft.rfft(xk, dim=-1)
for i, j in enumerate(self.index_kv):
if j >= xk_ft.shape[3]:
continue
xk_ft_[:, :, :, i] = xk_ft[:, :, :, j]
# perform attention mechanism on frequency domain
xqk_ft = (torch.einsum("bhex,bhey->bhxy", xq_ft_, xk_ft_))
if self.activation == 'tanh':
xqk_ft = xqk_ft.tanh()
elif self.activation == 'softmax':
xqk_ft = torch.softmax(abs(xqk_ft), dim=-1)
xqk_ft = torch.complex(xqk_ft, torch.zeros_like(xqk_ft))
else:
raise Exception('{} actiation function is not implemented'.format(self.activation))
xqkv_ft = torch.einsum("bhxy,bhey->bhex", xqk_ft, xk_ft_)
xqkvw = torch.einsum("bhex,heox->bhox", xqkv_ft, self.weights1)
out_ft = torch.zeros(B, H, E, L // 2 + 1, device=xq.device, dtype=torch.cfloat)
for i, j in enumerate(self.index_q):
if i >= xqkvw.shape[3] or j >= out_ft.shape[3]:
continue
out_ft[:, :, :, j] = xqkvw[:, :, :, i]
# Return to time domain
out = torch.fft.irfft(out_ft / self.in_channels / self.out_channels, n=xq.size(-1))
return (out, None)
@@ -0,0 +1,379 @@
import torch
import numpy as np
import torch.nn as nn
import torch.nn.functional as F
from torch import Tensor
from typing import List, Tuple
import math
from functools import partial
from einops import rearrange, reduce, repeat
from torch import nn, einsum, diagonal
from math import log2, ceil
import pdb
from utils.masking import LocalMask
from layers.utils import get_filter
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
class MultiWaveletTransform(nn.Module):
"""
1D multiwavelet block.
"""
def __init__(self, ich=1, k=8, alpha=16, c=128,
nCZ=1, L=0, base='legendre', attention_dropout=0.1):
super(MultiWaveletTransform, self).__init__()
print('base', base)
self.k = k
self.c = c
self.L = L
self.nCZ = nCZ
self.Lk0 = nn.Linear(ich, c * k)
self.Lk1 = nn.Linear(c * k, ich)
self.ich = ich
self.MWT_CZ = nn.ModuleList(MWT_CZ1d(k, alpha, L, c, base) for i in range(nCZ))
def forward(self, queries, keys, values, attn_mask):
B, L, H, E = queries.shape
_, S, _, D = values.shape
if L > S:
zeros = torch.zeros_like(queries[:, :(L - S), :]).float()
values = torch.cat([values, zeros], dim=1)
keys = torch.cat([keys, zeros], dim=1)
else:
values = values[:, :L, :, :]
keys = keys[:, :L, :, :]
values = values.view(B, L, -1)
V = self.Lk0(values).view(B, L, self.c, -1)
for i in range(self.nCZ):
V = self.MWT_CZ[i](V)
if i < self.nCZ - 1:
V = F.relu(V)
V = self.Lk1(V.view(B, L, -1))
V = V.view(B, L, -1, D)
return (V.contiguous(), None)
class MultiWaveletCross(nn.Module):
"""
1D Multiwavelet Cross Attention layer.
"""
def __init__(self, in_channels, out_channels, seq_len_q, seq_len_kv, modes, c=64,
k=8, ich=512,
L=0,
base='legendre',
mode_select_method='random',
initializer=None, activation='tanh',
**kwargs):
super(MultiWaveletCross, self).__init__()
print('base', base)
self.c = c
self.k = k
self.L = L
H0, H1, G0, G1, PHI0, PHI1 = get_filter(base, k)
H0r = H0 @ PHI0
G0r = G0 @ PHI0
H1r = H1 @ PHI1
G1r = G1 @ PHI1
H0r[np.abs(H0r) < 1e-8] = 0
H1r[np.abs(H1r) < 1e-8] = 0
G0r[np.abs(G0r) < 1e-8] = 0
G1r[np.abs(G1r) < 1e-8] = 0
self.max_item = 3
self.attn1 = FourierCrossAttentionW(in_channels=in_channels, out_channels=out_channels, seq_len_q=seq_len_q,
seq_len_kv=seq_len_kv, modes=modes, activation=activation,
mode_select_method=mode_select_method)
self.attn2 = FourierCrossAttentionW(in_channels=in_channels, out_channels=out_channels, seq_len_q=seq_len_q,
seq_len_kv=seq_len_kv, modes=modes, activation=activation,
mode_select_method=mode_select_method)
self.attn3 = FourierCrossAttentionW(in_channels=in_channels, out_channels=out_channels, seq_len_q=seq_len_q,
seq_len_kv=seq_len_kv, modes=modes, activation=activation,
mode_select_method=mode_select_method)
self.attn4 = FourierCrossAttentionW(in_channels=in_channels, out_channels=out_channels, seq_len_q=seq_len_q,
seq_len_kv=seq_len_kv, modes=modes, activation=activation,
mode_select_method=mode_select_method)
self.T0 = nn.Linear(k, k)
self.register_buffer('ec_s', torch.Tensor(
np.concatenate((H0.T, H1.T), axis=0)))
self.register_buffer('ec_d', torch.Tensor(
np.concatenate((G0.T, G1.T), axis=0)))
self.register_buffer('rc_e', torch.Tensor(
np.concatenate((H0r, G0r), axis=0)))
self.register_buffer('rc_o', torch.Tensor(
np.concatenate((H1r, G1r), axis=0)))
self.Lk = nn.Linear(ich, c * k)
self.Lq = nn.Linear(ich, c * k)
self.Lv = nn.Linear(ich, c * k)
self.out = nn.Linear(c * k, ich)
self.modes1 = modes
def forward(self, q, k, v, mask=None):
B, N, H, E = q.shape # (B, N, H, E) torch.Size([3, 768, 8, 2])
_, S, _, _ = k.shape # (B, S, H, E) torch.Size([3, 96, 8, 2])
q = q.view(q.shape[0], q.shape[1], -1)
k = k.view(k.shape[0], k.shape[1], -1)
v = v.view(v.shape[0], v.shape[1], -1)
q = self.Lq(q)
q = q.view(q.shape[0], q.shape[1], self.c, self.k)
k = self.Lk(k)
k = k.view(k.shape[0], k.shape[1], self.c, self.k)
v = self.Lv(v)
v = v.view(v.shape[0], v.shape[1], self.c, self.k)
if N > S:
zeros = torch.zeros_like(q[:, :(N - S), :]).float()
v = torch.cat([v, zeros], dim=1)
k = torch.cat([k, zeros], dim=1)
else:
v = v[:, :N, :, :]
k = k[:, :N, :, :]
ns = math.floor(np.log2(N))
nl = pow(2, math.ceil(np.log2(N)))
extra_q = q[:, 0:nl - N, :, :]
extra_k = k[:, 0:nl - N, :, :]
extra_v = v[:, 0:nl - N, :, :]
q = torch.cat([q, extra_q], 1)
k = torch.cat([k, extra_k], 1)
v = torch.cat([v, extra_v], 1)
Ud_q = torch.jit.annotate(List[Tuple[Tensor]], [])
Ud_k = torch.jit.annotate(List[Tuple[Tensor]], [])
Ud_v = torch.jit.annotate(List[Tuple[Tensor]], [])
Us_q = torch.jit.annotate(List[Tensor], [])
Us_k = torch.jit.annotate(List[Tensor], [])
Us_v = torch.jit.annotate(List[Tensor], [])
Ud = torch.jit.annotate(List[Tensor], [])
Us = torch.jit.annotate(List[Tensor], [])
# decompose
for i in range(ns - self.L):
# print('q shape',q.shape)
d, q = self.wavelet_transform(q)
Ud_q += [tuple([d, q])]
Us_q += [d]
for i in range(ns - self.L):
d, k = self.wavelet_transform(k)
Ud_k += [tuple([d, k])]
Us_k += [d]
for i in range(ns - self.L):
d, v = self.wavelet_transform(v)
Ud_v += [tuple([d, v])]
Us_v += [d]
for i in range(ns - self.L):
dk, sk = Ud_k[i], Us_k[i]
dq, sq = Ud_q[i], Us_q[i]
dv, sv = Ud_v[i], Us_v[i]
Ud += [self.attn1(dq[0], dk[0], dv[0], mask)[0] + self.attn2(dq[1], dk[1], dv[1], mask)[0]]
Us += [self.attn3(sq, sk, sv, mask)[0]]
v = self.attn4(q, k, v, mask)[0]
# reconstruct
for i in range(ns - 1 - self.L, -1, -1):
v = v + Us[i]
v = torch.cat((v, Ud[i]), -1)
v = self.evenOdd(v)
v = self.out(v[:, :N, :, :].contiguous().view(B, N, -1))
return (v.contiguous(), None)
def wavelet_transform(self, x):
xa = torch.cat([x[:, ::2, :, :],
x[:, 1::2, :, :],
], -1)
d = torch.matmul(xa, self.ec_d)
s = torch.matmul(xa, self.ec_s)
return d, s
def evenOdd(self, x):
B, N, c, ich = x.shape # (B, N, c, k)
assert ich == 2 * self.k
x_e = torch.matmul(x, self.rc_e)
x_o = torch.matmul(x, self.rc_o)
x = torch.zeros(B, N * 2, c, self.k,
device=x.device)
x[..., ::2, :, :] = x_e
x[..., 1::2, :, :] = x_o
return x
class FourierCrossAttentionW(nn.Module):
def __init__(self, in_channels, out_channels, seq_len_q, seq_len_kv, modes=16, activation='tanh',
mode_select_method='random'):
super(FourierCrossAttentionW, self).__init__()
print('corss fourier correlation used!')
self.in_channels = in_channels
self.out_channels = out_channels
self.modes1 = modes
self.activation = activation
def forward(self, q, k, v, mask):
B, L, E, H = q.shape
xq = q.permute(0, 3, 2, 1) # size = [B, H, E, L] torch.Size([3, 8, 64, 512])
xk = k.permute(0, 3, 2, 1)
xv = v.permute(0, 3, 2, 1)
self.index_q = list(range(0, min(int(L // 2), self.modes1)))
self.index_k_v = list(range(0, min(int(xv.shape[3] // 2), self.modes1)))
# Compute Fourier coefficients
xq_ft_ = torch.zeros(B, H, E, len(self.index_q), device=xq.device, dtype=torch.cfloat)
xq_ft = torch.fft.rfft(xq, dim=-1)
for i, j in enumerate(self.index_q):
xq_ft_[:, :, :, i] = xq_ft[:, :, :, j]
xk_ft_ = torch.zeros(B, H, E, len(self.index_k_v), device=xq.device, dtype=torch.cfloat)
xk_ft = torch.fft.rfft(xk, dim=-1)
for i, j in enumerate(self.index_k_v):
xk_ft_[:, :, :, i] = xk_ft[:, :, :, j]
xqk_ft = (torch.einsum("bhex,bhey->bhxy", xq_ft_, xk_ft_))
if self.activation == 'tanh':
xqk_ft = xqk_ft.tanh()
elif self.activation == 'softmax':
xqk_ft = torch.softmax(abs(xqk_ft), dim=-1)
xqk_ft = torch.complex(xqk_ft, torch.zeros_like(xqk_ft))
else:
raise Exception('{} actiation function is not implemented'.format(self.activation))
xqkv_ft = torch.einsum("bhxy,bhey->bhex", xqk_ft, xk_ft_)
xqkvw = xqkv_ft
out_ft = torch.zeros(B, H, E, L // 2 + 1, device=xq.device, dtype=torch.cfloat)
for i, j in enumerate(self.index_q):
out_ft[:, :, :, j] = xqkvw[:, :, :, i]
out = torch.fft.irfft(out_ft / self.in_channels / self.out_channels, n=xq.size(-1)).permute(0, 3, 2, 1)
# size = [B, L, H, E]
return (out, None)
class sparseKernelFT1d(nn.Module):
def __init__(self,
k, alpha, c=1,
nl=1,
initializer=None,
**kwargs):
super(sparseKernelFT1d, self).__init__()
self.modes1 = alpha
self.scale = (1 / (c * k * c * k))
self.weights1 = nn.Parameter(self.scale * torch.rand(c * k, c * k, self.modes1, dtype=torch.cfloat))
self.weights1.requires_grad = True
self.k = k
def compl_mul1d(self, x, weights):
# (batch, in_channel, x ), (in_channel, out_channel, x) -> (batch, out_channel, x)
return torch.einsum("bix,iox->box", x, weights)
def forward(self, x):
B, N, c, k = x.shape # (B, N, c, k)
x = x.view(B, N, -1)
x = x.permute(0, 2, 1)
x_fft = torch.fft.rfft(x)
# Multiply relevant Fourier modes
l = min(self.modes1, N // 2 + 1)
# l = N//2+1
out_ft = torch.zeros(B, c * k, N // 2 + 1, device=x.device, dtype=torch.cfloat)
out_ft[:, :, :l] = self.compl_mul1d(x_fft[:, :, :l], self.weights1[:, :, :l])
x = torch.fft.irfft(out_ft, n=N)
x = x.permute(0, 2, 1).view(B, N, c, k)
return x
# ##
class MWT_CZ1d(nn.Module):
def __init__(self,
k=3, alpha=64,
L=0, c=1,
base='legendre',
initializer=None,
**kwargs):
super(MWT_CZ1d, self).__init__()
self.k = k
self.L = L
H0, H1, G0, G1, PHI0, PHI1 = get_filter(base, k)
H0r = H0 @ PHI0
G0r = G0 @ PHI0
H1r = H1 @ PHI1
G1r = G1 @ PHI1
H0r[np.abs(H0r) < 1e-8] = 0
H1r[np.abs(H1r) < 1e-8] = 0
G0r[np.abs(G0r) < 1e-8] = 0
G1r[np.abs(G1r) < 1e-8] = 0
self.max_item = 3
self.A = sparseKernelFT1d(k, alpha, c)
self.B = sparseKernelFT1d(k, alpha, c)
self.C = sparseKernelFT1d(k, alpha, c)
self.T0 = nn.Linear(k, k)
self.register_buffer('ec_s', torch.Tensor(
np.concatenate((H0.T, H1.T), axis=0)))
self.register_buffer('ec_d', torch.Tensor(
np.concatenate((G0.T, G1.T), axis=0)))
self.register_buffer('rc_e', torch.Tensor(
np.concatenate((H0r, G0r), axis=0)))
self.register_buffer('rc_o', torch.Tensor(
np.concatenate((H1r, G1r), axis=0)))
def forward(self, x):
B, N, c, k = x.shape # (B, N, k)
ns = math.floor(np.log2(N))
nl = pow(2, math.ceil(np.log2(N)))
extra_x = x[:, 0:nl - N, :, :]
x = torch.cat([x, extra_x], 1)
Ud = torch.jit.annotate(List[Tensor], [])
Us = torch.jit.annotate(List[Tensor], [])
# decompose
for i in range(ns - self.L):
# print('x shape',x.shape)
d, x = self.wavelet_transform(x)
Ud += [self.A(d) + self.B(x)]
Us += [self.C(d)]
x = self.T0(x) # coarsest scale transform
# reconstruct
for i in range(ns - 1 - self.L, -1, -1):
x = x + Us[i]
x = torch.cat((x, Ud[i]), -1)
x = self.evenOdd(x)
x = x[:, :N, :, :]
return x
def wavelet_transform(self, x):
xa = torch.cat([x[:, ::2, :, :],
x[:, 1::2, :, :],
], -1)
d = torch.matmul(xa, self.ec_d)
s = torch.matmul(xa, self.ec_s)
return d, s
def evenOdd(self, x):
B, N, c, ich = x.shape # (B, N, c, k)
assert ich == 2 * self.k
x_e = torch.matmul(x, self.rc_e)
x_o = torch.matmul(x, self.rc_o)
x = torch.zeros(B, N * 2, c, self.k,
device=x.device)
x[..., ::2, :, :] = x_e
x[..., 1::2, :, :] = x_o
return x
@@ -0,0 +1,198 @@
import torch
import torch.nn as nn
import torch.nn.functional as F
import matplotlib.pyplot as plt
import numpy as np
import math
from math import sqrt
from utils.masking import TriangularCausalMask, ProbMask
import os
class FullAttention(nn.Module):
def __init__(self, mask_flag=True, factor=5, scale=None, attention_dropout=0.1, output_attention=False):
super(FullAttention, self).__init__()
self.scale = scale
self.mask_flag = mask_flag
self.output_attention = output_attention
self.dropout = nn.Dropout(attention_dropout)
def forward(self, queries, keys, values, attn_mask):
B, L, H, E = queries.shape
_, S, _, D = values.shape
scale = self.scale or 1. / sqrt(E)
# print('queries shape',queries.shape)
# print('keys shape',keys.shape)
# print('values shape',values.shape)
scores = torch.einsum("blhe,bshe->bhls", queries, keys)
if self.mask_flag:
if attn_mask is None:
attn_mask = TriangularCausalMask(B, L, device=queries.device)
scores.masked_fill_(attn_mask.mask, -np.inf)
A = self.dropout(torch.softmax(scale * scores, dim=-1))
V = torch.einsum("bhls,bshd->blhd", A, values)
#print('output shape',V.shape)
if self.output_attention:
return (V.contiguous(), A)
else:
return (V.contiguous(), None)
class SparseAttention(nn.Module):
def __init__(self, mask_flag=True, factor=5, scale=None, attention_dropout=0.1, output_attention=False):
super(FullAttention, self).__init__()
self.scale = scale
self.mask_flag = mask_flag
self.output_attention = output_attention
self.dropout = nn.Dropout(attention_dropout)
def forward(self, queries, keys, values, attn_mask=None):
B, L, H, E = queries.shape
_, S, _, D = values.shape
scale = self.scale or 1. / sqrt(E)
scores = torch.einsum("blhe,bshe->bhls", queries, keys)
if attn_mask is None:
attn_mask = LocalMask(B, L, device=queries.device)
scores.masked_fill_(attn_mask.mask, -np.inf)
A = self.dropout(torch.softmax(scale * scores, dim=-1))
V = torch.einsum("bhls,bshd->blhd", A, values)
#print('output shape',V.shape)
if self.output_attention:
return (V.contiguous(), A)
else:
return (V.contiguous(), None)
class ProbAttention(nn.Module):
def __init__(self, mask_flag=True, factor=5, scale=None, attention_dropout=0.1, output_attention=False):
super(ProbAttention, self).__init__()
self.factor = factor
self.scale = scale
self.mask_flag = mask_flag
self.output_attention = output_attention
self.dropout = nn.Dropout(attention_dropout)
def _prob_QK(self, Q, K, sample_k, n_top): # n_top: c*ln(L_q)
# Q [B, H, L, D]
B, H, L_K, E = K.shape
_, _, L_Q, _ = Q.shape
# calculate the sampled Q_K
K_expand = K.unsqueeze(-3).expand(B, H, L_Q, L_K, E)
index_sample = torch.randint(L_K, (L_Q, sample_k)) # real U = U_part(factor*ln(L_k))*L_q
K_sample = K_expand[:, :, torch.arange(L_Q).unsqueeze(1), index_sample, :]
Q_K_sample = torch.matmul(Q.unsqueeze(-2), K_sample.transpose(-2, -1)).squeeze()
# find the Top_k query with sparisty measurement
M = Q_K_sample.max(-1)[0] - torch.div(Q_K_sample.sum(-1), L_K)
M_top = M.topk(n_top, sorted=False)[1]
# use the reduced Q to calculate Q_K
Q_reduce = Q[torch.arange(B)[:, None, None],
torch.arange(H)[None, :, None],
M_top, :] # factor*ln(L_q)
Q_K = torch.matmul(Q_reduce, K.transpose(-2, -1)) # factor*ln(L_q)*L_k
return Q_K, M_top
def _get_initial_context(self, V, L_Q):
B, H, L_V, D = V.shape
if not self.mask_flag:
# V_sum = V.sum(dim=-2)
V_sum = V.mean(dim=-2)
contex = V_sum.unsqueeze(-2).expand(B, H, L_Q, V_sum.shape[-1]).clone()
else: # use mask
assert (L_Q == L_V) # requires that L_Q == L_V, i.e. for self-attention only
contex = V.cumsum(dim=-2)
return contex
def _update_context(self, context_in, V, scores, index, L_Q, attn_mask):
B, H, L_V, D = V.shape
if self.mask_flag:
attn_mask = ProbMask(B, H, L_Q, index, scores, device=V.device)
scores.masked_fill_(attn_mask.mask, -np.inf)
attn = torch.softmax(scores, dim=-1) # nn.Softmax(dim=-1)(scores)
context_in[torch.arange(B)[:, None, None],
torch.arange(H)[None, :, None],
index, :] = torch.matmul(attn, V).type_as(context_in)
if self.output_attention:
attns = (torch.ones([B, H, L_V, L_V]) / L_V).type_as(attn).to(attn.device)
attns[torch.arange(B)[:, None, None], torch.arange(H)[None, :, None], index, :] = attn
return (context_in, attns)
else:
return (context_in, None)
def forward(self, queries, keys, values, attn_mask):
B, L_Q, H, D = queries.shape
_, L_K, _, _ = keys.shape
queries = queries.transpose(2, 1)
keys = keys.transpose(2, 1)
values = values.transpose(2, 1)
U_part = self.factor * np.ceil(np.log(L_K)).astype('int').item() # c*ln(L_k)
u = self.factor * np.ceil(np.log(L_Q)).astype('int').item() # c*ln(L_q)
U_part = U_part if U_part < L_K else L_K
u = u if u < L_Q else L_Q
scores_top, index = self._prob_QK(queries, keys, sample_k=U_part, n_top=u)
# add scale factor
scale = self.scale or 1. / sqrt(D)
if scale is not None:
scores_top = scores_top * scale
# get the context
context = self._get_initial_context(values, L_Q)
# update the context with selected top_k queries
context, attn = self._update_context(context, values, scores_top, index, L_Q, attn_mask)
return context.contiguous(), attn
class AttentionLayer(nn.Module):
def __init__(self, attention, d_model, n_heads, d_keys=None,
d_values=None):
super(AttentionLayer, self).__init__()
d_keys = d_keys or (d_model // n_heads)
d_values = d_values or (d_model // n_heads)
self.inner_attention = attention
self.query_projection = nn.Linear(d_model, d_keys * n_heads)
self.key_projection = nn.Linear(d_model, d_keys * n_heads)
self.value_projection = nn.Linear(d_model, d_values * n_heads)
self.out_projection = nn.Linear(d_values * n_heads, d_model)
self.n_heads = n_heads
def forward(self, queries, keys, values, attn_mask):
B, L, _ = queries.shape
_, S, _ = keys.shape
H = self.n_heads
queries = self.query_projection(queries).view(B, L, H, -1)
keys = self.key_projection(keys).view(B, S, H, -1)
values = self.value_projection(values).view(B, S, H, -1)
out, attn = self.inner_attention(
queries,
keys,
values,
attn_mask
)
out = out.view(B, L, -1)
return self.out_projection(out), attn
@@ -0,0 +1,131 @@
import torch
import torch.nn as nn
import torch.nn.functional as F
class ConvLayer(nn.Module):
def __init__(self, c_in):
super(ConvLayer, self).__init__()
self.downConv = nn.Conv1d(in_channels=c_in,
out_channels=c_in,
kernel_size=3,
padding=2,
padding_mode='circular')
self.norm = nn.BatchNorm1d(c_in)
self.activation = nn.ELU()
self.maxPool = nn.MaxPool1d(kernel_size=3, stride=2, padding=1)
def forward(self, x):
x = self.downConv(x.permute(0, 2, 1))
x = self.norm(x)
x = self.activation(x)
x = self.maxPool(x)
x = x.transpose(1, 2)
return x
class EncoderLayer(nn.Module):
def __init__(self, attention, d_model, d_ff=None, dropout=0.1, activation="relu"):
super(EncoderLayer, self).__init__()
d_ff = d_ff or 4 * d_model
self.attention = attention
self.conv1 = nn.Conv1d(in_channels=d_model, out_channels=d_ff, kernel_size=1)
self.conv2 = nn.Conv1d(in_channels=d_ff, out_channels=d_model, kernel_size=1)
self.norm1 = nn.LayerNorm(d_model)
self.norm2 = nn.LayerNorm(d_model)
self.dropout = nn.Dropout(dropout)
self.activation = F.relu if activation == "relu" else F.gelu
def forward(self, x, attn_mask=None):
new_x, attn = self.attention(
x, x, x,
attn_mask=attn_mask
)
x = x + self.dropout(new_x)
y = x = self.norm1(x)
y = self.dropout(self.activation(self.conv1(y.transpose(-1, 1))))
y = self.dropout(self.conv2(y).transpose(-1, 1))
return self.norm2(x + y), attn
class Encoder(nn.Module):
def __init__(self, attn_layers, conv_layers=None, norm_layer=None):
super(Encoder, self).__init__()
self.attn_layers = nn.ModuleList(attn_layers)
self.conv_layers = nn.ModuleList(conv_layers) if conv_layers is not None else None
self.norm = norm_layer
def forward(self, x, attn_mask=None):
# x [B, L, D]
attns = []
if self.conv_layers is not None:
for attn_layer, conv_layer in zip(self.attn_layers, self.conv_layers):
x, attn = attn_layer(x, attn_mask=attn_mask)
x = conv_layer(x)
attns.append(attn)
x, attn = self.attn_layers[-1](x)
attns.append(attn)
else:
for attn_layer in self.attn_layers:
x, attn = attn_layer(x, attn_mask=attn_mask)
attns.append(attn)
if self.norm is not None:
x = self.norm(x)
return x, attns
class DecoderLayer(nn.Module):
def __init__(self, self_attention, cross_attention, d_model, d_ff=None,
dropout=0.1, activation="relu"):
super(DecoderLayer, self).__init__()
d_ff = d_ff or 4 * d_model
self.self_attention = self_attention
self.cross_attention = cross_attention
self.conv1 = nn.Conv1d(in_channels=d_model, out_channels=d_ff, kernel_size=1)
self.conv2 = nn.Conv1d(in_channels=d_ff, out_channels=d_model, kernel_size=1)
self.norm1 = nn.LayerNorm(d_model)
self.norm2 = nn.LayerNorm(d_model)
self.norm3 = nn.LayerNorm(d_model)
self.dropout = nn.Dropout(dropout)
self.activation = F.relu if activation == "relu" else F.gelu
def forward(self, x, cross, x_mask=None, cross_mask=None):
x = x + self.dropout(self.self_attention(
x, x, x,
attn_mask=x_mask
)[0])
x = self.norm1(x)
x = x + self.dropout(self.cross_attention(
x, cross, cross,
attn_mask=cross_mask
)[0])
y = x = self.norm2(x)
y = self.dropout(self.activation(self.conv1(y.transpose(-1, 1))))
y = self.dropout(self.conv2(y).transpose(-1, 1))
return self.norm3(x + y)
class Decoder(nn.Module):
def __init__(self, layers, norm_layer=None, projection=None):
super(Decoder, self).__init__()
self.layers = nn.ModuleList(layers)
self.norm = norm_layer
self.projection = projection
def forward(self, x, cross, x_mask=None, cross_mask=None):
for layer in self.layers:
x = layer(x, cross, x_mask=x_mask, cross_mask=cross_mask)
if self.norm is not None:
x = self.norm(x)
if self.projection is not None:
x = self.projection(x)
return x
@@ -0,0 +1,389 @@
import torch
import torch.nn as nn
import numpy as np
from functools import partial
from scipy.special import eval_legendre
from sympy import Poly, legendre, Symbol, chebyshevt
def legendreDer(k, x):
def _legendre(k, x):
return (2*k+1) * eval_legendre(k, x)
out = 0
for i in np.arange(k-1,-1,-2):
out += _legendre(i, x)
return out
def phi_(phi_c, x, lb = 0, ub = 1):
mask = np.logical_or(x<lb, x>ub) * 1.0
return np.polynomial.polynomial.Polynomial(phi_c)(x) * (1-mask)
def get_phi_psi(k, base):
x = Symbol('x')
phi_coeff = np.zeros((k,k))
phi_2x_coeff = np.zeros((k,k))
if base == 'legendre':
for ki in range(k):
coeff_ = Poly(legendre(ki, 2*x-1), x).all_coeffs()
phi_coeff[ki,:ki+1] = np.flip(np.sqrt(2*ki+1) * np.array(coeff_).astype(np.float64))
coeff_ = Poly(legendre(ki, 4*x-1), x).all_coeffs()
phi_2x_coeff[ki,:ki+1] = np.flip(np.sqrt(2) * np.sqrt(2*ki+1) * np.array(coeff_).astype(np.float64))
psi1_coeff = np.zeros((k, k))
psi2_coeff = np.zeros((k, k))
for ki in range(k):
psi1_coeff[ki,:] = phi_2x_coeff[ki,:]
for i in range(k):
a = phi_2x_coeff[ki,:ki+1]
b = phi_coeff[i, :i+1]
prod_ = np.convolve(a, b)
prod_[np.abs(prod_)<1e-8] = 0
proj_ = (prod_ * 1/(np.arange(len(prod_))+1) * np.power(0.5, 1+np.arange(len(prod_)))).sum()
psi1_coeff[ki,:] -= proj_ * phi_coeff[i,:]
psi2_coeff[ki,:] -= proj_ * phi_coeff[i,:]
for j in range(ki):
a = phi_2x_coeff[ki,:ki+1]
b = psi1_coeff[j, :]
prod_ = np.convolve(a, b)
prod_[np.abs(prod_)<1e-8] = 0
proj_ = (prod_ * 1/(np.arange(len(prod_))+1) * np.power(0.5, 1+np.arange(len(prod_)))).sum()
psi1_coeff[ki,:] -= proj_ * psi1_coeff[j,:]
psi2_coeff[ki,:] -= proj_ * psi2_coeff[j,:]
a = psi1_coeff[ki,:]
prod_ = np.convolve(a, a)
prod_[np.abs(prod_)<1e-8] = 0
norm1 = (prod_ * 1/(np.arange(len(prod_))+1) * np.power(0.5, 1+np.arange(len(prod_)))).sum()
a = psi2_coeff[ki,:]
prod_ = np.convolve(a, a)
prod_[np.abs(prod_)<1e-8] = 0
norm2 = (prod_ * 1/(np.arange(len(prod_))+1) * (1-np.power(0.5, 1+np.arange(len(prod_))))).sum()
norm_ = np.sqrt(norm1 + norm2)
psi1_coeff[ki,:] /= norm_
psi2_coeff[ki,:] /= norm_
psi1_coeff[np.abs(psi1_coeff)<1e-8] = 0
psi2_coeff[np.abs(psi2_coeff)<1e-8] = 0
phi = [np.poly1d(np.flip(phi_coeff[i,:])) for i in range(k)]
psi1 = [np.poly1d(np.flip(psi1_coeff[i,:])) for i in range(k)]
psi2 = [np.poly1d(np.flip(psi2_coeff[i,:])) for i in range(k)]
elif base == 'chebyshev':
for ki in range(k):
if ki == 0:
phi_coeff[ki,:ki+1] = np.sqrt(2/np.pi)
phi_2x_coeff[ki,:ki+1] = np.sqrt(2/np.pi) * np.sqrt(2)
else:
coeff_ = Poly(chebyshevt(ki, 2*x-1), x).all_coeffs()
phi_coeff[ki,:ki+1] = np.flip(2/np.sqrt(np.pi) * np.array(coeff_).astype(np.float64))
coeff_ = Poly(chebyshevt(ki, 4*x-1), x).all_coeffs()
phi_2x_coeff[ki,:ki+1] = np.flip(np.sqrt(2) * 2 / np.sqrt(np.pi) * np.array(coeff_).astype(np.float64))
phi = [partial(phi_, phi_coeff[i,:]) for i in range(k)]
x = Symbol('x')
kUse = 2*k
roots = Poly(chebyshevt(kUse, 2*x-1)).all_roots()
x_m = np.array([rt.evalf(20) for rt in roots]).astype(np.float64)
# x_m[x_m==0.5] = 0.5 + 1e-8 # add small noise to avoid the case of 0.5 belonging to both phi(2x) and phi(2x-1)
# not needed for our purpose here, we use even k always to avoid
wm = np.pi / kUse / 2
psi1_coeff = np.zeros((k, k))
psi2_coeff = np.zeros((k, k))
psi1 = [[] for _ in range(k)]
psi2 = [[] for _ in range(k)]
for ki in range(k):
psi1_coeff[ki,:] = phi_2x_coeff[ki,:]
for i in range(k):
proj_ = (wm * phi[i](x_m) * np.sqrt(2)* phi[ki](2*x_m)).sum()
psi1_coeff[ki,:] -= proj_ * phi_coeff[i,:]
psi2_coeff[ki,:] -= proj_ * phi_coeff[i,:]
for j in range(ki):
proj_ = (wm * psi1[j](x_m) * np.sqrt(2) * phi[ki](2*x_m)).sum()
psi1_coeff[ki,:] -= proj_ * psi1_coeff[j,:]
psi2_coeff[ki,:] -= proj_ * psi2_coeff[j,:]
psi1[ki] = partial(phi_, psi1_coeff[ki,:], lb = 0, ub = 0.5)
psi2[ki] = partial(phi_, psi2_coeff[ki,:], lb = 0.5, ub = 1)
norm1 = (wm * psi1[ki](x_m) * psi1[ki](x_m)).sum()
norm2 = (wm * psi2[ki](x_m) * psi2[ki](x_m)).sum()
norm_ = np.sqrt(norm1 + norm2)
psi1_coeff[ki,:] /= norm_
psi2_coeff[ki,:] /= norm_
psi1_coeff[np.abs(psi1_coeff)<1e-8] = 0
psi2_coeff[np.abs(psi2_coeff)<1e-8] = 0
psi1[ki] = partial(phi_, psi1_coeff[ki,:], lb = 0, ub = 0.5+1e-16)
psi2[ki] = partial(phi_, psi2_coeff[ki,:], lb = 0.5+1e-16, ub = 1)
return phi, psi1, psi2
def get_filter(base, k):
def psi(psi1, psi2, i, inp):
mask = (inp<=0.5) * 1.0
return psi1[i](inp) * mask + psi2[i](inp) * (1-mask)
if base not in ['legendre', 'chebyshev']:
raise Exception('Base not supported')
x = Symbol('x')
H0 = np.zeros((k,k))
H1 = np.zeros((k,k))
G0 = np.zeros((k,k))
G1 = np.zeros((k,k))
PHI0 = np.zeros((k,k))
PHI1 = np.zeros((k,k))
phi, psi1, psi2 = get_phi_psi(k, base)
if base == 'legendre':
roots = Poly(legendre(k, 2*x-1)).all_roots()
x_m = np.array([rt.evalf(20) for rt in roots]).astype(np.float64)
wm = 1/k/legendreDer(k,2*x_m-1)/eval_legendre(k-1,2*x_m-1)
for ki in range(k):
for kpi in range(k):
H0[ki, kpi] = 1/np.sqrt(2) * (wm * phi[ki](x_m/2) * phi[kpi](x_m)).sum()
G0[ki, kpi] = 1/np.sqrt(2) * (wm * psi(psi1, psi2, ki, x_m/2) * phi[kpi](x_m)).sum()
H1[ki, kpi] = 1/np.sqrt(2) * (wm * phi[ki]((x_m+1)/2) * phi[kpi](x_m)).sum()
G1[ki, kpi] = 1/np.sqrt(2) * (wm * psi(psi1, psi2, ki, (x_m+1)/2) * phi[kpi](x_m)).sum()
PHI0 = np.eye(k)
PHI1 = np.eye(k)
elif base == 'chebyshev':
x = Symbol('x')
kUse = 2*k
roots = Poly(chebyshevt(kUse, 2*x-1)).all_roots()
x_m = np.array([rt.evalf(20) for rt in roots]).astype(np.float64)
# x_m[x_m==0.5] = 0.5 + 1e-8 # add small noise to avoid the case of 0.5 belonging to both phi(2x) and phi(2x-1)
# not needed for our purpose here, we use even k always to avoid
wm = np.pi / kUse / 2
for ki in range(k):
for kpi in range(k):
H0[ki, kpi] = 1/np.sqrt(2) * (wm * phi[ki](x_m/2) * phi[kpi](x_m)).sum()
G0[ki, kpi] = 1/np.sqrt(2) * (wm * psi(psi1, psi2, ki, x_m/2) * phi[kpi](x_m)).sum()
H1[ki, kpi] = 1/np.sqrt(2) * (wm * phi[ki]((x_m+1)/2) * phi[kpi](x_m)).sum()
G1[ki, kpi] = 1/np.sqrt(2) * (wm * psi(psi1, psi2, ki, (x_m+1)/2) * phi[kpi](x_m)).sum()
PHI0[ki, kpi] = (wm * phi[ki](2*x_m) * phi[kpi](2*x_m)).sum() * 2
PHI1[ki, kpi] = (wm * phi[ki](2*x_m-1) * phi[kpi](2*x_m-1)).sum() * 2
PHI0[np.abs(PHI0)<1e-8] = 0
PHI1[np.abs(PHI1)<1e-8] = 0
H0[np.abs(H0)<1e-8] = 0
H1[np.abs(H1)<1e-8] = 0
G0[np.abs(G0)<1e-8] = 0
G1[np.abs(G1)<1e-8] = 0
return H0, H1, G0, G1, PHI0, PHI1
def train(model, train_loader, optimizer, epoch, device, verbose = 0,
lossFn = None, lr_schedule=None,
post_proc = lambda args: args):
if lossFn is None:
lossFn = nn.MSELoss()
model.train()
total_loss = 0.
for batch_idx, (data, target) in enumerate(train_loader):
bs = len(data)
data, target = data.to(device), target.to(device)
optimizer.zero_grad()
output = model(data)
target = post_proc(target)
output = post_proc(output)
loss = lossFn(output.view(bs, -1), target.view(bs, -1))
loss.backward()
optimizer.step()
total_loss += loss.sum().item()
if lr_schedule is not None: lr_schedule.step()
if verbose>0:
print('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format(
epoch, batch_idx * len(data), len(train_loader.dataset),
100. * batch_idx / len(train_loader), loss.item()))
return total_loss/len(train_loader.dataset)
def test(model, test_loader, device, verbose=0, lossFn=None,
post_proc = lambda args: args):
model.eval()
if lossFn is None:
lossFn = nn.MSELoss()
total_loss = 0.
predictions = []
with torch.no_grad():
for data, target in test_loader:
bs = len(data)
data, target = data.to(device), target.to(device)
output = model(data)
output = post_proc(output)
loss = lossFn(output.view(bs, -1), target.view(bs, -1))
total_loss += loss.sum().item()
return total_loss/len(test_loader.dataset)
# Till EoF
# taken from FNO paper:
# https://github.com/zongyi-li/fourier_neural_operator
# normalization, pointwise gaussian
class UnitGaussianNormalizer(object):
def __init__(self, x, eps=0.00001):
super(UnitGaussianNormalizer, self).__init__()
# x could be in shape of ntrain*n or ntrain*T*n or ntrain*n*T
self.mean = torch.mean(x, 0)
self.std = torch.std(x, 0)
self.eps = eps
def encode(self, x):
x = (x - self.mean) / (self.std + self.eps)
return x
def decode(self, x, sample_idx=None):
if sample_idx is None:
std = self.std + self.eps # n
mean = self.mean
else:
if len(self.mean.shape) == len(sample_idx[0].shape):
std = self.std[sample_idx] + self.eps # batch*n
mean = self.mean[sample_idx]
if len(self.mean.shape) > len(sample_idx[0].shape):
std = self.std[:,sample_idx]+ self.eps # T*batch*n
mean = self.mean[:,sample_idx]
# x is in shape of batch*n or T*batch*n
x = (x * std) + mean
return x
def cuda(self):
self.mean = self.mean.cuda()
self.std = self.std.cuda()
def cpu(self):
self.mean = self.mean.cpu()
self.std = self.std.cpu()
# normalization, Gaussian
class GaussianNormalizer(object):
def __init__(self, x, eps=0.00001):
super(GaussianNormalizer, self).__init__()
self.mean = torch.mean(x)
self.std = torch.std(x)
self.eps = eps
def encode(self, x):
x = (x - self.mean) / (self.std + self.eps)
return x
def decode(self, x, sample_idx=None):
x = (x * (self.std + self.eps)) + self.mean
return x
def cuda(self):
self.mean = self.mean.cuda()
self.std = self.std.cuda()
def cpu(self):
self.mean = self.mean.cpu()
self.std = self.std.cpu()
# normalization, scaling by range
class RangeNormalizer(object):
def __init__(self, x, low=0.0, high=1.0):
super(RangeNormalizer, self).__init__()
mymin = torch.min(x, 0)[0].view(-1)
mymax = torch.max(x, 0)[0].view(-1)
self.a = (high - low)/(mymax - mymin)
self.b = -self.a*mymax + high
def encode(self, x):
s = x.size()
x = x.view(s[0], -1)
x = self.a*x + self.b
x = x.view(s)
return x
def decode(self, x):
s = x.size()
x = x.view(s[0], -1)
x = (x - self.b)/self.a
x = x.view(s)
return x
class LpLoss(object):
def __init__(self, d=2, p=2, size_average=True, reduction=True):
super(LpLoss, self).__init__()
#Dimension and Lp-norm type are postive
assert d > 0 and p > 0
self.d = d
self.p = p
self.reduction = reduction
self.size_average = size_average
def abs(self, x, y):
num_examples = x.size()[0]
#Assume uniform mesh
h = 1.0 / (x.size()[1] - 1.0)
all_norms = (h**(self.d/self.p))*torch.norm(x.view(num_examples,-1) - y.view(num_examples,-1), self.p, 1)
if self.reduction:
if self.size_average:
return torch.mean(all_norms)
else:
return torch.sum(all_norms)
return all_norms
def rel(self, x, y):
num_examples = x.size()[0]
diff_norms = torch.norm(x.reshape(num_examples,-1) - y.reshape(num_examples,-1), self.p, 1)
y_norms = torch.norm(y.reshape(num_examples,-1), self.p, 1)
if self.reduction:
if self.size_average:
return torch.mean(diff_norms/y_norms)
else:
return torch.sum(diff_norms/y_norms)
return diff_norms/y_norms
def __call__(self, x, y):
return self.rel(x, y)
@@ -0,0 +1,102 @@
# coding=utf-8
# author=maziqing
# email=maziqing.mzq@alibaba-inc.com
import torch
import torch.nn as nn
import torch.nn.functional as F
from layers.Embed import DataEmbedding, DataEmbedding_wo_pos
from layers.AutoCorrelation import AutoCorrelation, AutoCorrelationLayer
from layers.Autoformer_EncDec import Encoder, Decoder, EncoderLayer, DecoderLayer, my_Layernorm, series_decomp
class Model(nn.Module):
"""
Autoformer is the first method to achieve the series-wise connection,
with inherent O(LlogL) complexity
"""
def __init__(self, configs):
super(Model, self).__init__()
self.seq_len = configs.seq_len
self.label_len = configs.label_len
self.pred_len = configs.pred_len
self.output_attention = configs.output_attention
# Decomp
kernel_size = configs.moving_avg
self.decomp = series_decomp(kernel_size)
# Embedding
# The series-wise connection inherently contains the sequential information.
# Thus, we can discard the position embedding of transformers.
self.enc_embedding = DataEmbedding_wo_pos(configs.enc_in, configs.d_model, configs.embed, configs.freq,
configs.dropout)
self.dec_embedding = DataEmbedding_wo_pos(configs.dec_in, configs.d_model, configs.embed, configs.freq,
configs.dropout)
# Encoder
self.encoder = Encoder(
[
EncoderLayer(
AutoCorrelationLayer(
AutoCorrelation(False, configs.factor, attention_dropout=configs.dropout,
output_attention=configs.output_attention),
configs.d_model, configs.n_heads),
configs.d_model,
configs.d_ff,
moving_avg=configs.moving_avg,
dropout=configs.dropout,
activation=configs.activation
) for l in range(configs.e_layers)
],
norm_layer=my_Layernorm(configs.d_model)
)
# Decoder
self.decoder = Decoder(
[
DecoderLayer(
AutoCorrelationLayer(
AutoCorrelation(True, configs.factor, attention_dropout=configs.dropout,
output_attention=False),
configs.d_model, configs.n_heads),
AutoCorrelationLayer(
AutoCorrelation(False, configs.factor, attention_dropout=configs.dropout,
output_attention=False),
configs.d_model, configs.n_heads),
configs.d_model,
configs.c_out,
configs.d_ff,
moving_avg=configs.moving_avg,
dropout=configs.dropout,
activation=configs.activation,
)
for l in range(configs.d_layers)
],
norm_layer=my_Layernorm(configs.d_model),
projection=nn.Linear(configs.d_model, configs.c_out, bias=True)
)
def forward(self, x_enc, x_mark_enc, x_dec, x_mark_dec,
enc_self_mask=None, dec_self_mask=None, dec_enc_mask=None):
# decomp init
mean = torch.mean(x_enc, dim=1).unsqueeze(1).repeat(1, self.pred_len, 1)
zeros = torch.zeros([x_dec.shape[0], self.pred_len, x_dec.shape[2]], device=x_enc.device)
seasonal_init, trend_init = self.decomp(x_enc)
# decoder input
trend_init = torch.cat([trend_init[:, -self.label_len:, :], mean], dim=1)
seasonal_init = torch.cat([seasonal_init[:, -self.label_len:, :], zeros], dim=1)
# enc
enc_out = self.enc_embedding(x_enc, x_mark_enc)
enc_out, attns = self.encoder(enc_out, attn_mask=enc_self_mask)
# dec
dec_out = self.dec_embedding(seasonal_init, x_mark_dec)
seasonal_part, trend_part = self.decoder(dec_out, enc_out, x_mask=dec_self_mask, cross_mask=dec_enc_mask,
trend=trend_init)
# final
dec_out = trend_part + seasonal_part
if self.output_attention:
return dec_out[:, -self.pred_len:, :], attns
else:
return dec_out[:, -self.pred_len:, :] # [B, L, D]
@@ -0,0 +1,227 @@
import torch
import torch.nn as nn
import torch.nn.functional as F
from layers.Embed import DataEmbedding, DataEmbedding_wo_pos,DataEmbedding_wo_pos_temp,DataEmbedding_wo_temp
from layers.AutoCorrelation import AutoCorrelation, AutoCorrelationLayer
from layers.FourierCorrelation import FourierBlock, FourierCrossAttention
from layers.MultiWaveletCorrelation import MultiWaveletCross, MultiWaveletTransform
from layers.SelfAttention_Family import FullAttention, ProbAttention
# from layers.FED_wo_decomp import Encoder, Decoder, EncoderLayer, DecoderLayer, my_Layernorm, series_decomp, series_decomp_multi
from layers.Autoformer_EncDec import Encoder, Decoder, EncoderLayer, DecoderLayer, my_Layernorm, series_decomp, series_decomp_multi
import math
import numpy as np
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
class Model(nn.Module):
"""
FEDformer performs the attention mechanism on frequency domain and achieved O(N) complexity
"""
def __init__(self, configs):
super(Model, self).__init__()
self.version = configs.version
self.mode_select = configs.mode_select
self.modes = configs.modes
self.seq_len = configs.seq_len
self.label_len = configs.label_len
self.pred_len = configs.pred_len
self.output_attention = configs.output_attention
# Decomp
kernel_size = configs.moving_avg
if isinstance(kernel_size, list):
self.decomp = series_decomp_multi(kernel_size)
else:
self.decomp = series_decomp(kernel_size)
# Embedding
# The series-wise connection inherently contains the sequential information.
# Thus, we can discard the position embedding of transformers.
# self.enc_embedding = DataEmbedding_wo_pos(configs.enc_in, configs.d_model, configs.embed, configs.freq,
# configs.dropout)
# self.dec_embedding = DataEmbedding_wo_pos(configs.dec_in, configs.d_model, configs.embed, configs.freq,
# configs.dropout)
if configs.embed_type == 0:
self.enc_embedding = DataEmbedding_wo_pos(configs.enc_in, configs.d_model, configs.embed, configs.freq,
configs.dropout)
self.dec_embedding = DataEmbedding_wo_pos(configs.dec_in, configs.d_model, configs.embed, configs.freq,
configs.dropout)
elif configs.embed_type == 1:
self.enc_embedding = DataEmbedding(configs.enc_in, configs.d_model, configs.embed, configs.freq,
configs.dropout)
self.dec_embedding = DataEmbedding(configs.dec_in, configs.d_model, configs.embed, configs.freq,
configs.dropout)
elif configs.embed_type == 2:
self.enc_embedding = DataEmbedding_wo_pos_temp(configs.enc_in, configs.d_model, configs.embed, configs.freq,
configs.dropout)
self.dec_embedding = DataEmbedding_wo_pos_temp(configs.dec_in, configs.d_model, configs.embed, configs.freq,
configs.dropout)
elif configs.embed_type == 3:
self.enc_embedding = DataEmbedding_wo_temp(configs.enc_in, configs.d_model, configs.embed, configs.freq,
configs.dropout)
self.dec_embedding = DataEmbedding_wo_temp(configs.dec_in, configs.d_model, configs.embed, configs.freq,
configs.dropout)
if configs.version == 'Wavelets':
encoder_self_att = MultiWaveletTransform(ich=configs.d_model, L=configs.L, base=configs.base)
decoder_self_att = MultiWaveletTransform(ich=configs.d_model, L=configs.L, base=configs.base)
decoder_cross_att = MultiWaveletCross(in_channels=configs.d_model,
out_channels=configs.d_model,
seq_len_q=self.seq_len // 2 + self.pred_len,
seq_len_kv=self.seq_len,
modes=configs.modes,
ich=configs.d_model,
base=configs.base,
activation=configs.cross_activation)
else:
encoder_self_att = FourierBlock(in_channels=configs.d_model,
out_channels=configs.d_model,
seq_len=self.seq_len,
modes=configs.modes,
mode_select_method=configs.mode_select)
decoder_self_att = FourierBlock(in_channels=configs.d_model,
out_channels=configs.d_model,
seq_len=self.seq_len//2+self.pred_len,
modes=configs.modes,
mode_select_method=configs.mode_select)
decoder_cross_att = FourierCrossAttention(in_channels=configs.d_model,
out_channels=configs.d_model,
seq_len_q=self.seq_len//2+self.pred_len,
seq_len_kv=self.seq_len,
modes=configs.modes,
mode_select_method=configs.mode_select)
# Encoder
enc_modes = int(min(configs.modes, configs.seq_len//2))
dec_modes = int(min(configs.modes, (configs.seq_len//2+configs.pred_len)//2))
print('enc_modes: {}, dec_modes: {}'.format(enc_modes, dec_modes))
self.encoder = Encoder(
[
EncoderLayer(
AutoCorrelationLayer(
encoder_self_att,
configs.d_model, configs.n_heads),
configs.d_model,
configs.d_ff,
moving_avg=configs.moving_avg,
dropout=configs.dropout,
activation=configs.activation
) for l in range(configs.e_layers)
],
norm_layer=my_Layernorm(configs.d_model)
)
# Decoder
self.decoder = Decoder(
[
DecoderLayer(
AutoCorrelationLayer(
decoder_self_att,
configs.d_model, configs.n_heads),
AutoCorrelationLayer(
decoder_cross_att,
configs.d_model, configs.n_heads),
configs.d_model,
configs.c_out,
configs.d_ff,
moving_avg=configs.moving_avg,
dropout=configs.dropout,
activation=configs.activation,
)
for l in range(configs.d_layers)
],
norm_layer=my_Layernorm(configs.d_model),
projection=nn.Linear(configs.d_model, configs.c_out, bias=True)
)
# def forward(self, x_enc, x_mark_enc, x_dec, x_mark_dec,
# enc_self_mask=None, dec_self_mask=None, dec_enc_mask=None):
# # decomp init
# mean = torch.mean(x_enc, dim=1).unsqueeze(1).repeat(1, self.pred_len, 1)
# # zeros = torch.zeros([x_dec.shape[0], self.pred_len, x_dec.shape[2]]).to(device) # cuda()
# # seasonal_init, trend_init = self.decomp(x_enc)
# # decoder input
# dec_in = torch.cat([x_enc[:, -self.label_len:, :], mean], dim=1)
# dec_out = F.pad(x_enc[:, -self.label_len:, :], (0, 0, 0, self.pred_len))
# # enc
# enc_out = self.enc_embedding(x_enc, x_mark_enc)
# enc_out, attns = self.encoder(enc_out, attn_mask=enc_self_mask)
# # dec
# dec_out = self.dec_embedding(dec_out, x_mark_dec)
# dec_out, _ = self.decoder(dec_out, enc_out, x_mask=dec_self_mask, cross_mask=dec_enc_mask,
# trend=dec_in)
# # final
# # dec_out = trend_part + seasonal_part
# if self.output_attention:
# return dec_out[:, -self.pred_len:, :], attns
# else:
# return dec_out[:, -self.pred_len:, :] # [B, L, D]
def forward(self, x_enc, x_mark_enc, x_dec, x_mark_dec,
enc_self_mask=None, dec_self_mask=None, dec_enc_mask=None):
# decomp init
mean = torch.mean(x_enc, dim=1).unsqueeze(1).repeat(1, self.pred_len, 1)
zeros = torch.zeros([x_dec.shape[0], self.pred_len, x_dec.shape[2]]).to(device) # cuda()
seasonal_init, trend_init = self.decomp(x_enc)
# decoder input
trend_init = torch.cat([trend_init[:, -self.label_len:, :], mean], dim=1)
seasonal_init = F.pad(seasonal_init[:, -self.label_len:, :], (0, 0, 0, self.pred_len))
# enc
enc_out = self.enc_embedding(x_enc, x_mark_enc)
enc_out, attns = self.encoder(enc_out, attn_mask=enc_self_mask)
# dec
dec_out = self.dec_embedding(seasonal_init, x_mark_dec)
seasonal_part, trend_part = self.decoder(dec_out, enc_out, x_mask=dec_self_mask, cross_mask=dec_enc_mask,
trend=trend_init)
# final
dec_out = trend_part + seasonal_part
if self.output_attention:
return dec_out[:, -self.pred_len:, :], attns
else:
return dec_out[:, -self.pred_len:, :] # [B, L, D]
if __name__ == '__main__':
class Configs(object):
ab = 0
modes = 32
mode_select = 'random'
# version = 'Fourier'
version = 'Wavelets'
moving_avg = [12, 24]
L = 1
base = 'legendre'
cross_activation = 'tanh'
seq_len = 96
label_len = 48
pred_len = 96
output_attention = True
enc_in = 7
dec_in = 7
d_model = 16
embed = 'timeF'
dropout = 0.05
freq = 'h'
factor = 1
n_heads = 8
d_ff = 16
e_layers = 2
d_layers = 1
c_out = 7
activation = 'gelu'
wavelet = 0
configs = Configs()
model = Model(configs)
print('parameter number is {}'.format(sum(p.numel() for p in model.parameters())))
enc = torch.randn([3, configs.seq_len, 7])
enc_mark = torch.randn([3, configs.seq_len, 4])
dec = torch.randn([3, configs.seq_len//2+configs.pred_len, 7])
dec_mark = torch.randn([3, configs.seq_len//2+configs.pred_len, 4])
out = model.forward(enc, enc_mark, dec, dec_mark)
print(out)
@@ -0,0 +1,80 @@
import torch
import torch.nn as nn
import torch.nn.functional as F
from utils.masking import TriangularCausalMask, ProbMask
from layers.Transformer_EncDec import Decoder, DecoderLayer, Encoder, EncoderLayer, ConvLayer
from layers.SelfAttention_Family import FullAttention, ProbAttention, AttentionLayer
from layers.Embed import DataEmbedding
import numpy as np
class Model(nn.Module):
"""
Informer with Propspare attention in O(LlogL) complexity
"""
def __init__(self, configs):
super(Model, self).__init__()
self.pred_len = configs.pred_len
self.output_attention = configs.output_attention
# Embedding
self.enc_embedding = DataEmbedding(configs.enc_in, configs.d_model, configs.embed, configs.freq,
configs.dropout)
self.dec_embedding = DataEmbedding(configs.dec_in, configs.d_model, configs.embed, configs.freq,
configs.dropout)
# Encoder
self.encoder = Encoder(
[
EncoderLayer(
AttentionLayer(
ProbAttention(False, configs.factor, attention_dropout=configs.dropout,
output_attention=configs.output_attention),
configs.d_model, configs.n_heads),
configs.d_model,
configs.d_ff,
dropout=configs.dropout,
activation=configs.activation
) for l in range(configs.e_layers)
],
[
ConvLayer(
configs.d_model
) for l in range(configs.e_layers - 1)
] if configs.distil else None,
norm_layer=torch.nn.LayerNorm(configs.d_model)
)
# Decoder
self.decoder = Decoder(
[
DecoderLayer(
AttentionLayer(
ProbAttention(True, configs.factor, attention_dropout=configs.dropout, output_attention=False),
configs.d_model, configs.n_heads),
AttentionLayer(
ProbAttention(False, configs.factor, attention_dropout=configs.dropout, output_attention=False),
configs.d_model, configs.n_heads),
configs.d_model,
configs.d_ff,
dropout=configs.dropout,
activation=configs.activation,
)
for l in range(configs.d_layers)
],
norm_layer=torch.nn.LayerNorm(configs.d_model),
projection=nn.Linear(configs.d_model, configs.c_out, bias=True)
)
def forward(self, x_enc, x_mark_enc, x_dec, x_mark_dec,
enc_self_mask=None, dec_self_mask=None, dec_enc_mask=None):
enc_out = self.enc_embedding(x_enc, x_mark_enc)
enc_out, attns = self.encoder(enc_out, attn_mask=enc_self_mask)
dec_out = self.dec_embedding(x_dec, x_mark_dec)
dec_out = self.decoder(dec_out, enc_out, x_mask=dec_self_mask, cross_mask=dec_enc_mask)
if self.output_attention:
return dec_out[:, -self.pred_len:, :], attns
else:
return dec_out[:, -self.pred_len:, :] # [B, L, D]
@@ -0,0 +1,71 @@
import torch
import torch.nn as nn
import torch.nn.functional as F
from layers.Transformer_EncDec import Decoder, DecoderLayer, Encoder, EncoderLayer, ConvLayer
from layers.SelfAttention_Family import FullAttention, AttentionLayer
from layers.Embed import DataEmbedding
class Model(nn.Module):
"""
Vanilla Transformer with O(L^2) complexity
"""
def __init__(self, configs):
super(Model, self).__init__()
self.pred_len = configs.pred_len
self.output_attention = configs.output_attention
# Embedding
self.enc_embedding = DataEmbedding(configs.enc_in, configs.d_model, configs.embed, configs.freq,
configs.dropout)
self.dec_embedding = DataEmbedding(configs.dec_in, configs.d_model, configs.embed, configs.freq,
configs.dropout)
# Encoder
self.encoder = Encoder(
[
EncoderLayer(
AttentionLayer(
FullAttention(False, configs.factor, attention_dropout=configs.dropout,
output_attention=configs.output_attention), configs.d_model, configs.n_heads),
configs.d_model,
configs.d_ff,
dropout=configs.dropout,
activation=configs.activation
) for l in range(configs.e_layers)
],
norm_layer=torch.nn.LayerNorm(configs.d_model)
)
# Decoder
self.decoder = Decoder(
[
DecoderLayer(
AttentionLayer(
FullAttention(True, configs.factor, attention_dropout=configs.dropout, output_attention=False),
configs.d_model, configs.n_heads),
AttentionLayer(
FullAttention(False, configs.factor, attention_dropout=configs.dropout, output_attention=False),
configs.d_model, configs.n_heads),
configs.d_model,
configs.d_ff,
dropout=configs.dropout,
activation=configs.activation,
)
for l in range(configs.d_layers)
],
norm_layer=torch.nn.LayerNorm(configs.d_model),
projection=nn.Linear(configs.d_model, configs.c_out, bias=True)
)
def forward(self, x_enc, x_mark_enc, x_dec, x_mark_dec,
enc_self_mask=None, dec_self_mask=None, dec_enc_mask=None):
enc_out = self.enc_embedding(x_enc, x_mark_enc)
enc_out, attns = self.encoder(enc_out, attn_mask=enc_self_mask)
dec_out = self.dec_embedding(x_dec, x_mark_dec)
dec_out = self.decoder(dec_out, enc_out, x_mask=dec_self_mask, cross_mask=dec_enc_mask)
if self.output_attention:
return dec_out[:, -self.pred_len:, :], attns
else:
return dec_out[:, -self.pred_len:, :] # [B, L, D]
@@ -0,0 +1,165 @@
import argparse
import os
import torch
from exp.exp_main import Exp_Main
import random
import numpy as np
fix_seed = 2021
random.seed(fix_seed)
torch.manual_seed(fix_seed)
np.random.seed(fix_seed)
parser = argparse.ArgumentParser(description='Autoformer & Transformer family for Time Series Forecasting')
# basic config
parser.add_argument('--is_training', type=int, default=1, help='status')
parser.add_argument('--task_id', type=str, default='test', help='task id')
parser.add_argument('--model', type=str, default='Reformer',
help='model name, options: [FEDformer, Autoformer, Informer, Transformer]')
# supplementary config for FEDformer model
parser.add_argument('--version', type=str, default='Fourier',
help='for FEDformer, there are two versions to choose, options: [Fourier, Wavelets]')
parser.add_argument('--mode_select', type=str, default='random',
help='for FEDformer, there are two mode selection method, options: [random, low]')
parser.add_argument('--modes', type=int, default=64, help='modes to be selected random 64')
parser.add_argument('--L', type=int, default=3, help='ignore level')
parser.add_argument('--base', type=str, default='legendre', help='mwt base')
parser.add_argument('--cross_activation', type=str, default='tanh',
help='mwt cross atention activation function tanh or softmax')
# data loader
parser.add_argument('--data', type=str, default='ETTh1', help='dataset type')
parser.add_argument('--root_path', type=str, default='../dataset', help='root path of the data file')
parser.add_argument('--data_path', type=str, default='ETTh1.csv', help='data file')
parser.add_argument('--features', type=str, default='M',
help='forecasting task, options:[M, S, MS]; M:multivariate predict multivariate, '
'S:univariate predict univariate, MS:multivariate predict univariate')
parser.add_argument('--target', type=str, default='OT', help='target feature in S or MS task')
parser.add_argument('--freq', type=str, default='h',
help='freq for time features encoding, options:[s:secondly, t:minutely, h:hourly, d:daily, '
'b:business days, w:weekly, m:monthly], you can also use more detailed freq like 15min or 3h')
parser.add_argument('--checkpoints', type=str, default='./checkpoints/', help='location of model checkpoints')
# forecasting task
parser.add_argument('--seq_len', type=int, default=96, help='input sequence length')
parser.add_argument('--label_len', type=int, default=48, help='start token length')
parser.add_argument('--pred_len', type=int, default=96, help='prediction sequence length')
parser.add_argument('--embed_type', type=int, default=0, help='prediction sequence length')
# parser.add_argument('--cross_activation', type=str, default='tanh'
# model define
parser.add_argument('--enc_in', type=int, default=7, help='encoder input size')
parser.add_argument('--dec_in', type=int, default=7, help='decoder input size')
parser.add_argument('--c_out', type=int, default=7, help='output size')
parser.add_argument('--d_model', type=int, default=512, help='dimension of model')
parser.add_argument('--n_heads', type=int, default=8, help='num of heads')
parser.add_argument('--e_layers', type=int, default=2, help='num of encoder layers')
parser.add_argument('--d_layers', type=int, default=1, help='num of decoder layers')
parser.add_argument('--d_ff', type=int, default=2048, help='dimension of fcn')
parser.add_argument('--moving_avg', default=[24], help='window size of moving average')
parser.add_argument('--factor', type=int, default=1, help='attn factor')
parser.add_argument('--distil', action='store_false',
help='whether to use distilling in encoder, using this argument means not using distilling',
default=True)
parser.add_argument('--dropout', type=float, default=0.05, help='dropout')
parser.add_argument('--embed', type=str, default='timeF',
help='time features encoding, options:[timeF, fixed, learned]')
parser.add_argument('--activation', type=str, default='gelu', help='activation')
parser.add_argument('--output_attention', action='store_true', help='whether to output attention in ecoder')
parser.add_argument('--do_predict', action='store_true', help='whether to predict unseen future data')
# optimization
parser.add_argument('--num_workers', type=int, default=10, help='data loader num workers')
parser.add_argument('--itr', type=int, default=1, help='experiments times')
parser.add_argument('--train_epochs', type=int, default=10, help='train epochs')
parser.add_argument('--batch_size', type=int, default=32, help='batch size of train input data')
parser.add_argument('--patience', type=int, default=3, help='early stopping patience')
parser.add_argument('--learning_rate', type=float, default=0.0001, help='optimizer learning rate')
parser.add_argument('--des', type=str, default='test', help='exp description')
parser.add_argument('--loss', type=str, default='mse', help='loss function')
parser.add_argument('--lradj', type=str, default='type1', help='adjust learning rate')
parser.add_argument('--use_amp', action='store_true', help='use automatic mixed precision training', default=False)
# GPU
parser.add_argument('--use_gpu', type=bool, default=True, help='use gpu')
parser.add_argument('--gpu', type=int, default=0, help='gpu')
parser.add_argument('--use_multi_gpu', action='store_true', help='use multiple gpus', default=False)
parser.add_argument('--devices', type=str, default='0,1', help='device ids of multi gpus')
args = parser.parse_args()
args.use_gpu = True if torch.cuda.is_available() and args.use_gpu else False
if args.use_gpu and args.use_multi_gpu:
args.dvices = args.devices.replace(' ', '')
device_ids = args.devices.split(',')
args.device_ids = [int(id_) for id_ in device_ids]
args.gpu = args.device_ids[0]
print('Args in experiment:')
print(args)
Exp = Exp_Main
if args.is_training:
for ii in range(args.itr):
# setting record of experiments
setting = '{}_{}_{}_modes{}_{}_ft{}_sl{}_ll{}_pl{}_dm{}_nh{}_el{}_dl{}_df{}_fc{}_eb{}_dt{}_{}_{}'.format(
args.task_id,
args.model,
args.mode_select,
args.modes,
args.data,
args.features,
args.seq_len,
args.label_len,
args.pred_len,
args.d_model,
args.n_heads,
args.e_layers,
args.d_layers,
args.d_ff,
args.factor,
args.embed,
args.distil,
args.des,
ii)
exp = Exp(args) # set experiments
print('>>>>>>>start training : {}>>>>>>>>>>>>>>>>>>>>>>>>>>'.format(setting))
exp.train(setting)
print('>>>>>>>testing : {}<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<'.format(setting))
exp.test(setting)
if args.do_predict:
print('>>>>>>>predicting : {}<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<'.format(setting))
exp.predict(setting, True)
torch.cuda.empty_cache()
else:
ii = 0
setting = '{}_{}_{}_ft{}_sl{}_ll{}_pl{}_dm{}_nh{}_el{}_dl{}_df{}_fc{}_eb{}_dt{}_{}_{}'.format(args.model_id,
args.model,
args.data,
args.features,
args.seq_len,
args.label_len,
args.pred_len,
args.d_model,
args.n_heads,
args.e_layers,
args.d_layers,
args.d_ff,
args.factor,
args.embed,
args.distil,
args.des, ii)
exp = Exp(args) # set experiments
print('>>>>>>>testing : {}<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<'.format(setting))
exp.test(setting, test=1)
torch.cuda.empty_cache()
@@ -0,0 +1,202 @@
# cd FEDformer
if [ ! -d "../logs" ]; then
mkdir ../logs
fi
if [ ! -d "../logs/LongForecasting" ]; then
mkdir ../logs/LongForecasting
fi
for preLen in 96 192 336 720
do
# ETTm1
python -u run.py \
--is_training 1 \
--data_path ETTm1.csv \
--task_id ETTm1 \
--model FEDformer \
--data ETTm1 \
--features M \
--seq_len 96 \
--label_len 48 \
--pred_len $preLen \
--e_layers 2 \
--d_layers 1 \
--factor 3 \
--enc_in 7 \
--dec_in 7 \
--c_out 7 \
--des 'Exp' \
--d_model 512 \
--itr 1 >../logs/LongForecasting/FEDformer_ETTm1_$pred_len.log
# ETTh1
python -u run.py \
--is_training 1 \
--data_path ETTh1.csv \
--task_id ETTh1 \
--model FEDformer \
--data ETTh1 \
--features S \
--seq_len 96 \
--label_len 48 \
--pred_len $preLen \
--e_layers 2 \
--d_layers 1 \
--factor 3 \
--enc_in 7 \
--dec_in 7 \
--c_out 7 \
--des 'Exp' \
--d_model 512 \
--itr 1 >../logs/LongForecasting/FEDformer_ETTh1_$pred_len.log
# ETTm2
python -u run.py \
--is_training 1 \
--data_path ETTm2.csv \
--task_id ETTm2 \
--model FEDformer \
--data ETTm2 \
--features M \
--seq_len 96 \
--label_len 48 \
--pred_len $preLen \
--e_layers 2 \
--d_layers 1 \
--factor 3 \
--enc_in 7 \
--dec_in 7 \
--c_out 7 \
--des 'Exp' \
--d_model 512 \
--itr 1 >../logs/LongForecasting/FEDformer_ETTm2_$pred_len.log
# ETTh2
python -u run.py \
--is_training 1 \
--data_path ETTh2.csv \
--task_id ETTh2 \
--model FEDformer \
--data ETTh2 \
--features M \
--seq_len 96 \
--label_len 48 \
--pred_len $preLen \
--e_layers 2 \
--d_layers 1 \
--factor 3 \
--enc_in 7 \
--dec_in 7 \
--c_out 7 \
--des 'Exp' \
--d_model 512 \
--itr 1 >../logs/LongForecasting/FEDformer_ETTh2_$pred_len.log
# electricity
python -u run.py \
--is_training 1 \
--data_path electricity.csv \
--task_id ECL \
--model FEDformer \
--data custom \
--features M \
--seq_len 96 \
--label_len 48 \
--pred_len $preLen \
--e_layers 2 \
--d_layers 1 \
--factor 3 \
--enc_in 321 \
--dec_in 321 \
--c_out 321 \
--des 'Exp' \
--itr 1 >../logs/LongForecasting/FEDformer_electricity_$pred_len.log
# exchange
python -u run.py \
--is_training 1 \
--data_path exchange_rate.csv \
--task_id Exchange \
--model FEDformer \
--data custom \
--features S \
--seq_len 96 \
--label_len 48 \
--pred_len $preLen \
--e_layers 2 \
--d_layers 1 \
--factor 3 \
--enc_in 8 \
--dec_in 8 \
--c_out 8 \
--des 'Exp' \
--itr 1 >../logs/LongForecasting/FEDformer_exchange_rate_$pred_len.log
# traffic
python -u run.py \
--is_training 1 \
--data_path traffic.csv \
--task_id traffic \
--model FEDformer \
--data custom \
--features S \
--seq_len 96 \
--label_len 48 \
--pred_len $preLen \
--e_layers 2 \
--d_layers 1 \
--factor 3 \
--enc_in 862 \
--dec_in 862 \
--c_out 862 \
--des 'Exp' \
--itr 1 \
--train_epochs 3 >../logs/LongForecasting/FEDformer_traffic_$pred_len.log
# weather
python -u run.py \
--is_training 1 \
--data_path weather.csv \
--task_id weather \
--model FEDformer \
--data custom \
--features S \
--seq_len 96 \
--label_len 48 \
--pred_len $preLen \
--e_layers 2 \
--d_layers 1 \
--factor 3 \
--enc_in 21 \
--dec_in 21 \
--c_out 21 \
--des 'Exp' \
--itr 1 >../logs/LongForecasting/FEDformer_weather_$pred_len.log
done
for preLen in 24 36 48 60
do
# illness
python -u run.py \
--is_training 1 \
--data_path national_illness.csv \
--task_id ili \
--model FEDformer \
--data custom \
--features S \
--seq_len 36 \
--label_len 18 \
--pred_len $preLen \
--e_layers 2 \
--d_layers 1 \
--factor 3 \
--enc_in 7 \
--dec_in 7 \
--c_out 7 \
--des 'Exp' \
--itr 1 >../logs/LongForecasting/FEDformer_ili_$pred_len.log
done
# cd ..
@@ -0,0 +1,237 @@
# cd FEDformer
if [ ! -d "../logs" ]; then
mkdir ../logs
fi
if [ ! -d "../logs/LookBackWindow" ]; then
mkdir ../logs/LookBackWindow
fi
for seqLen in 36 48 60 72 144 288
do
for pred_len in 24 576
do
python -u run.py \
--is_training 1 \
--root_path .../dataset/ \
--data_path ETTm1.csv \
--task_id ETTm1 \
--model FEDformer \
--data ETTm1 \
--features M \
--seq_len $seqLen \
--label_len 48 \
--pred_len $pred_len \
--e_layers 2 \
--d_layers 1 \
--factor 3 \
--enc_in 7 \
--dec_in 7 \
--c_out 7 \
--des 'Exp' \
--d_model 512 \
--itr 1 >../logs/LookBackWindow/FEDformer_ETTm2_$seqLen'_'$pred_len.log
python -u run.py \
--is_training 1 \
--root_path .../dataset/ \
--data_path ETTm2.csv \
--task_id ETTm2 \
--model FEDformer \
--data ETTm2 \
--features M \
--seq_len $seqLen \
--label_len 48 \
--pred_len $pred_len \
--e_layers 2 \
--d_layers 1 \
--factor 3 \
--enc_in 7 \
--dec_in 7 \
--c_out 7 \
--des 'Exp' \
--d_model 512 \
--itr 1 >../logs/LookBackWindow/FEDformer_ETTm2_$seqLen'_'$pred_len.log
done
done
for seqLen in 48 72 120 144 168 192 336 720
do
for pred_len in 24 720
do
# ETTh1
python -u run.py \
--is_training 1 \
--root_path .../dataset/ \
--data_path ETTh1.csv \
--task_id ETTh1 \
--model FEDformer \
--data ETTh1 \
--features M \
--seq_len $seqLen \
--label_len 48 \
--pred_len $pred_len \
--e_layers 2 \
--d_layers 1 \
--factor 3 \
--enc_in 7 \
--dec_in 7 \
--c_out 7 \
--des 'Exp' \
--d_model 512 \
--itr 1 >../logs/LookBackWindow/FEDformer_ETTh1_$seqLen'_'$pred_len.log
# ETTh2
python -u run.py \
--is_training 1 \
--root_path .../dataset/ \
--data_path ETTh2.csv \
--task_id ETTh2 \
--model FEDformer \
--data ETTh2 \
--features M \
--seq_len $seqLen \
--label_len 48 \
--pred_len $pred_len \
--e_layers 2 \
--d_layers 1 \
--factor 3 \
--enc_in 7 \
--dec_in 7 \
--c_out 7 \
--des 'Exp' \
--d_model 512 \
--itr 1 >../logs/LookBackWindow/FEDformer_ETTh2_$seqLen'_'$pred_len.log
## electricity
python -u run.py \
--is_training 1 \
--root_path .../dataset/ \
--data_path electricity.csv \
--task_id ECL \
--model FEDformer \
--data custom \
--features M \
--seq_len $seqLen \
--label_len 48 \
--pred_len $pred_len \
--e_layers 2 \
--d_layers 1 \
--factor 3 \
--enc_in 321 \
--dec_in 321 \
--c_out 321 \
--des 'Exp' \
--itr 1 >../logs/LookBackWindow/FEDformer_electricity_$seqLen'_'$pred_len.log
# exchange
python -u run.py \
--is_training 1 \
--root_path .../dataset/ \
--data_path exchange_rate.csv \
--task_id Exchange \
--model FEDformer \
--data custom \
--features M \
--seq_len $seqLen \
--label_len 48 \
--pred_len $pred_len \
--e_layers 2 \
--d_layers 1 \
--factor 3 \
--enc_in 8 \
--dec_in 8 \
--c_out 8 \
--des 'Exp' \
--itr 1 >../logs/LookBackWindow/FEDformer_exchange_rate_$seqLen'_'$pred_len.log
# traffic
python -u run.py \
--is_training 1 \
--root_path .../dataset/ \
--data_path traffic.csv \
--task_id traffic \
--model FEDformer \
--data custom \
--features M \
--seq_len $seqLen \
--label_len 48 \
--pred_len $pred_len \
--e_layers 2 \
--d_layers 1 \
--factor 3 \
--enc_in 862 \
--dec_in 862 \
--c_out 862 \
--des 'Exp' \
--itr 1 \
--train_epochs 3 >../logs/LookBackWindow/FEDformer_traffic_$seqLen'_'$pred_len.log
# weather
python -u run.py \
--is_training 1 \
--root_path .../dataset/ \
--data_path weather.csv \
--task_id weather \
--model FEDformer \
--data custom \
--features M \
--seq_len $seqLen \
--label_len 48 \
--pred_len $pred_len \
--e_layers 2 \
--d_layers 1 \
--factor 3 \
--enc_in 21 \
--dec_in 21 \
--c_out 21 \
--des 'Exp' \
--itr 1 >../logs/LookBackWindow/FEDformer_weather_$seqLen'_'$pred_len.log
done
done
for seqLen in 26 52 78 104 130 156 208
do
# illness
python -u run.py \
--is_training 1 \
--root_path .../dataset/ \
--data_path national_illness.csv \
--task_id ili \
--model FEDformer \
--data custom \
--features M \
--seq_len $seqLen \
--label_len 18 \
--pred_len 24 \
--e_layers 2 \
--d_layers 1 \
--factor 3 \
--enc_in 7 \
--dec_in 7 \
--c_out 7 \
--des 'Exp' \
--itr 1 >../logs/LookBackWindow/FEDformer_ili_$seqLen'_'24.log
python -u run.py \
--is_training 1 \
--root_path .../dataset/ \
--data_path national_illness.csv \
--task_id ili \
--model FEDformer \
--data custom \
--features M \
--seq_len $seqLen \
--label_len 18 \
--pred_len 60 \
--e_layers 2 \
--d_layers 1 \
--factor 3 \
--enc_in 7 \
--dec_in 7 \
--c_out 7 \
--des 'Exp' \
--itr 1 >../logs/LookBackWindow/FEDformer_ili_$seqLen'_'60.log
done
# cd ..
@@ -0,0 +1,39 @@
import torch
import numpy as np
import math
class TriangularCausalMask():
def __init__(self, B, L, device="cpu"):
mask_shape = [B, 1, L, L]
with torch.no_grad():
self._mask = torch.triu(torch.ones(mask_shape, dtype=torch.bool), diagonal=1).to(device)
@property
def mask(self):
return self._mask
class ProbMask():
def __init__(self, B, H, L, index, scores, device="cpu"):
_mask = torch.ones(L, scores.shape[-1], dtype=torch.bool).to(device).triu(1)
_mask_ex = _mask[None, None, :].expand(B, H, L, scores.shape[-1])
indicator = _mask_ex[torch.arange(B)[:, None, None],
torch.arange(H)[None, :, None],
index, :].to(device)
self._mask = indicator.view(scores.shape).to(device)
@property
def mask(self):
return self._mask
class LocalMask():
def __init__(self, B, L,S,device="cpu"):
mask_shape = [B, 1, L, S]
with torch.no_grad():
self.len = math.ceil(np.log2(L))
self._mask1 = torch.triu(torch.ones(mask_shape, dtype=torch.bool), diagonal=1).to(device)
self._mask2 = ~torch.triu(torch.ones(mask_shape,dtype=torch.bool),diagonal=-self.len).to(device)
self._mask = self._mask1+self._mask2
@property
def mask(self):
return self._mask
@@ -0,0 +1,51 @@
import numpy as np
def RSE(pred, true):
return np.sqrt(np.sum((true - pred) ** 2)) / np.sqrt(np.sum((true - true.mean()) ** 2))
def CORR(pred, true):
u = ((true - true.mean(0)) * (pred - pred.mean(0))).sum(0)
d = np.sqrt(((true - true.mean(0)) ** 2 * (pred - pred.mean(0)) ** 2).sum(0))
return (u / d).mean(-1)
def MAE(pred, true):
return np.mean(np.abs(pred - true))
def MSE(pred, true):
return np.mean((pred - true) ** 2)
def RMSE(pred, true):
return np.sqrt(MSE(pred, true))
def MAPE(pred, true):
return np.mean(np.abs((pred - true) / true))
def MSPE(pred, true):
return np.mean(np.square((pred - true) / true))
def metric(pred, true):
mae = MAE(pred, true)
mse = MSE(pred, true)
rmse = RMSE(pred, true)
mape = MAPE(pred, true)
mspe = MSPE(pred, true)
return mae, mse, rmse, mape, mspe
def metric2(pred, true):
mae = MAE(pred, true)
mse = MSE(pred, true)
rmse = RMSE(pred, true)
mape = MAPE(pred, true)
mspe = MSPE(pred, true)
rse = RSE(pred, true)
corr = CORR(pred, true)
return mae, mse, rmse, mape, mspe, rse, corr
@@ -0,0 +1,134 @@
from typing import List
import numpy as np
import pandas as pd
from pandas.tseries import offsets
from pandas.tseries.frequencies import to_offset
class TimeFeature:
def __init__(self):
pass
def __call__(self, index: pd.DatetimeIndex) -> np.ndarray:
pass
def __repr__(self):
return self.__class__.__name__ + "()"
class SecondOfMinute(TimeFeature):
"""Minute of hour encoded as value between [-0.5, 0.5]"""
def __call__(self, index: pd.DatetimeIndex) -> np.ndarray:
return index.second / 59.0 - 0.5
class MinuteOfHour(TimeFeature):
"""Minute of hour encoded as value between [-0.5, 0.5]"""
def __call__(self, index: pd.DatetimeIndex) -> np.ndarray:
return index.minute / 59.0 - 0.5
class HourOfDay(TimeFeature):
"""Hour of day encoded as value between [-0.5, 0.5]"""
def __call__(self, index: pd.DatetimeIndex) -> np.ndarray:
return index.hour / 23.0 - 0.5
class DayOfWeek(TimeFeature):
"""Hour of day encoded as value between [-0.5, 0.5]"""
def __call__(self, index: pd.DatetimeIndex) -> np.ndarray:
return index.dayofweek / 6.0 - 0.5
class DayOfMonth(TimeFeature):
"""Day of month encoded as value between [-0.5, 0.5]"""
def __call__(self, index: pd.DatetimeIndex) -> np.ndarray:
return (index.day - 1) / 30.0 - 0.5
class DayOfYear(TimeFeature):
"""Day of year encoded as value between [-0.5, 0.5]"""
def __call__(self, index: pd.DatetimeIndex) -> np.ndarray:
return (index.dayofyear - 1) / 365.0 - 0.5
class MonthOfYear(TimeFeature):
"""Month of year encoded as value between [-0.5, 0.5]"""
def __call__(self, index: pd.DatetimeIndex) -> np.ndarray:
return (index.month - 1) / 11.0 - 0.5
class WeekOfYear(TimeFeature):
"""Week of year encoded as value between [-0.5, 0.5]"""
def __call__(self, index: pd.DatetimeIndex) -> np.ndarray:
return (index.isocalendar().week - 1) / 52.0 - 0.5
def time_features_from_frequency_str(freq_str: str) -> List[TimeFeature]:
"""
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.YearEnd: [],
offsets.QuarterEnd: [MonthOfYear],
offsets.MonthEnd: [MonthOfYear],
offsets.Week: [DayOfMonth, WeekOfYear],
offsets.Day: [DayOfWeek, DayOfMonth, DayOfYear],
offsets.BusinessDay: [DayOfWeek, DayOfMonth, DayOfYear],
offsets.Hour: [HourOfDay, DayOfWeek, DayOfMonth, DayOfYear],
offsets.Minute: [
MinuteOfHour,
HourOfDay,
DayOfWeek,
DayOfMonth,
DayOfYear,
],
offsets.Second: [
SecondOfMinute,
MinuteOfHour,
HourOfDay,
DayOfWeek,
DayOfMonth,
DayOfYear,
],
}
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]
supported_freq_msg = f"""
Unsupported frequency {freq_str}
The following frequencies are supported:
Y - yearly
alias: A
M - monthly
W - weekly
D - daily
B - business days
H - hourly
T - minutely
alias: min
S - secondly
"""
raise RuntimeError(supported_freq_msg)
def time_features(dates, freq='h'):
return np.vstack([feat(dates) for feat in time_features_from_frequency_str(freq)])
@@ -0,0 +1,88 @@
import numpy as np
import torch
import matplotlib.pyplot as plt
plt.switch_backend('agg')
def adjust_learning_rate(optimizer, epoch, args):
# lr = args.learning_rate * (0.2 ** (epoch // 2))
if args.lradj == 'type1':
lr_adjust = {epoch: args.learning_rate * (0.5 ** ((epoch - 1) // 1))}
elif args.lradj == 'type2':
lr_adjust = {
2: 5e-5, 4: 1e-5, 6: 5e-6, 8: 1e-6,
10: 5e-7, 15: 1e-7, 20: 5e-8
}
elif args.lradj =='type3':
lr_adjust = {epoch: args.learning_rate}
elif args.lradj == 'type4':
lr_adjust = {epoch: args.learning_rate * (0.9 ** ((epoch - 1) // 1))}
if epoch in lr_adjust.keys():
lr = lr_adjust[epoch]
for param_group in optimizer.param_groups:
param_group['lr'] = lr
print('Updating learning rate to {}'.format(lr))
class EarlyStopping:
def __init__(self, patience=7, verbose=False, delta=0):
self.patience = patience
self.verbose = verbose
self.counter = 0
self.best_score = None
self.early_stop = False
self.val_loss_min = np.Inf
self.delta = delta
def __call__(self, val_loss, model, path):
score = -val_loss
if self.best_score is None:
self.best_score = score
self.save_checkpoint(val_loss, model, path)
elif score < self.best_score + self.delta:
self.counter += 1
print(f'EarlyStopping counter: {self.counter} out of {self.patience}')
if self.counter >= self.patience:
self.early_stop = True
else:
self.best_score = score
self.save_checkpoint(val_loss, model, path)
self.counter = 0
def save_checkpoint(self, val_loss, model, path):
if self.verbose:
print(f'Validation loss decreased ({self.val_loss_min:.6f} --> {val_loss:.6f}). Saving model ...')
torch.save(model.state_dict(), path + '/' + 'checkpoint.pth')
self.val_loss_min = val_loss
class dotdict(dict):
"""dot.notation access to dictionary attributes"""
__getattr__ = dict.get
__setattr__ = dict.__setitem__
__delattr__ = dict.__delitem__
class StandardScaler():
def __init__(self, mean, std):
self.mean = mean
self.std = std
def transform(self, data):
return (data - self.mean) / self.std
def inverse_transform(self, data):
return (data * self.std) + self.mean
def visual(true, preds=None, name='./pic/test.pdf'):
"""
Results visualization
"""
plt.figure()
plt.plot(true, label='GroundTruth', linewidth=2)
if preds is not None:
plt.plot(preds, label='Prediction', linewidth=2)
plt.legend()
plt.savefig(name, bbox_inches='tight')
@@ -0,0 +1,209 @@
# ALL scripts in this file come from Autoformer
if [ ! -d "./logs" ]; then
mkdir ./logs
fi
if [ ! -d "./logs/LongForecasting" ]; then
mkdir ./logs/LongForecasting
fi
random_seed=2021
model_name=Informer
for pred_len in 96 192 336 720
do
python -u run_longExp.py \
--random_seed $random_seed \
--is_training 1 \
--root_path ./dataset/ \
--data_path exchange_rate.csv \
--model_id exchange_96_$pred_len \
--model $model_name \
--data custom \
--features M \
--seq_len 96 \
--label_len 48 \
--pred_len $pred_len \
--e_layers 2 \
--d_layers 1 \
--factor 3 \
--enc_in 8 \
--dec_in 8 \
--c_out 8 \
--des 'Exp' \
--itr 1 \
--train_epochs 1 >logs/LongForecasting/$model_name'_exchange_rate_'$pred_len.log
python -u run_longExp.py \
--random_seed $random_seed \
--is_training 1 \
--root_path ./dataset/ \
--data_path electricity.csv \
--model_id electricity_96_$pred_len \
--model $model_name \
--data custom \
--features M \
--seq_len 96 \
--label_len 48 \
--pred_len $pred_len \
--e_layers 2 \
--d_layers 1 \
--factor 3 \
--enc_in 321 \
--dec_in 321 \
--c_out 321 \
--des 'Exp' \
--itr 1 >logs/LongForecasting/$model_name'_electricity_'$pred_len.log
python -u run_longExp.py \
--random_seed $random_seed \
--is_training 1 \
--root_path ./dataset/ \
--data_path traffic.csv \
--model_id traffic_96_$pred_len \
--model $model_name \
--data custom \
--features M \
--seq_len 96 \
--label_len 48 \
--pred_len $pred_len \
--e_layers 2 \
--d_layers 1 \
--factor 3 \
--enc_in 862 \
--dec_in 862 \
--c_out 862 \
--des 'Exp' \
--itr 1 \
--train_epochs 3 >logs/LongForecasting/$model_name'_traffic_'$pred_len.log
python -u run_longExp.py \
--random_seed $random_seed \
--is_training 1 \
--root_path ./dataset/ \
--data_path weather.csv \
--model_id weather_96_$pred_len \
--model $model_name \
--data custom \
--features M \
--seq_len 96 \
--label_len 48 \
--pred_len $pred_len \
--e_layers 2 \
--d_layers 1 \
--factor 3 \
--enc_in 21 \
--dec_in 21 \
--c_out 21 \
--des 'Exp' \
--itr 1 \
--train_epochs 2 >logs/LongForecasting/$model_name'_weather_'$pred_len.log
python -u run_longExp.py \
--random_seed $random_seed \
--is_training 1 \
--root_path ./dataset/ \
--data_path ETTh1.csv \
--model_id ETTh1_96_$pred_len \
--model $model_name \
--data ETTh1 \
--features M \
--seq_len 96 \
--label_len 48 \
--pred_len $pred_len \
--e_layers 2 \
--d_layers 1 \
--factor 3 \
--enc_in 7 \
--dec_in 7 \
--c_out 7 \
--des 'Exp' \
--itr 1 >logs/LongForecasting/$model_name'_Etth1_'$pred_len.log
python -u run_longExp.py \
--random_seed $random_seed \
--is_training 1 \
--root_path ./dataset/ \
--data_path ETTh2.csv \
--model_id ETTh2_96_$pred_len \
--model $model_name \
--data ETTh2 \
--features M \
--seq_len 96 \
--label_len 48 \
--pred_len $pred_len \
--e_layers 2 \
--d_layers 1 \
--factor 3 \
--enc_in 7 \
--dec_in 7 \
--c_out 7 \
--des 'Exp' \
--itr 1 >logs/LongForecasting/$model_name'_Etth2_'$pred_len.log
python -u run_longExp.py \
--random_seed $random_seed \
--is_training 1 \
--root_path ./dataset/ \
--data_path ETTm1.csv \
--model_id ETTm1_96_$pred_len \
--model $model_name \
--data ETTm1 \
--features M \
--seq_len 96 \
--label_len 48 \
--pred_len $pred_len \
--e_layers 2 \
--d_layers 1 \
--factor 3 \
--enc_in 7 \
--dec_in 7 \
--c_out 7 \
--des 'Exp' \
--itr 1 >logs/LongForecasting/$model_name'_Ettm1_'$pred_len.log
python -u run_longExp.py \
--random_seed $random_seed \
--is_training 1 \
--root_path ./dataset/ \
--data_path ETTm2.csv \
--model_id ETTm2_96_$pred_len \
--model $model_name \
--data ETTm2 \
--features M \
--seq_len 96 \
--label_len 48 \
--pred_len $pred_len \
--e_layers 2 \
--d_layers 1 \
--factor 3 \
--enc_in 7 \
--dec_in 7 \
--c_out 7 \
--des 'Exp' \
--itr 1 >logs/LongForecasting/$model_name'_Ettm2_'$pred_len.log
done
for pred_len in 24 36 48 60
do
python -u run_longExp.py \
--random_seed $random_seed \
--is_training 1 \
--root_path ./dataset/ \
--data_path national_illness.csv \
--model_id ili_36_$pred_len \
--model $model_name \
--data custom \
--features M \
--seq_len 36 \
--label_len 18 \
--pred_len $pred_len \
--e_layers 2 \
--d_layers 1 \
--factor 3 \
--enc_in 7 \
--dec_in 7 \
--c_out 7 \
--des 'Exp' \
--itr 1 >logs/LongForecasting/$model_name'_ili_'$pred_len.log
done
@@ -0,0 +1,7 @@
Legal Disclaimer
Within this source code, the comments in Chinese shall be the original, governing version. Any comment in other languages are for reference only. In the event of any conflict between the Chinese language version comments and other language version comments, the Chinese language version shall prevail.
法律免责声明
关于代码注释部分,中文注释为官方版本,其它语言注释仅做参考。中文注释可能与其它语言注释存在不一致,当中文注释与其它语言注释存在不一致时,请以中文注释为准。
@@ -0,0 +1,203 @@
Copyright 2022 - The AntGroup-RI team. All rights reserved.
Apache License
Version 2.0, January 2004
http://www.apache.org/licenses/
TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
1. Definitions.
"License" shall mean the terms and conditions for use, reproduction,
and distribution as defined by Sections 1 through 9 of this document.
"Licensor" shall mean the copyright owner or entity authorized by
the copyright owner that is granting the License.
"Legal Entity" shall mean the union of the acting entity and all
other entities that control, are controlled by, or are under common
control with that entity. For the purposes of this definition,
"control" means (i) the power, direct or indirect, to cause the
direction or management of such entity, whether by contract or
otherwise, or (ii) ownership of fifty percent (50%) or more of the
outstanding shares, or (iii) beneficial ownership of such entity.
"You" (or "Your") shall mean an individual or Legal Entity
exercising permissions granted by this License.
"Source" form shall mean the preferred form for making modifications,
including but not limited to software source code, documentation
source, and configuration files.
"Object" form shall mean any form resulting from mechanical
transformation or translation of a Source form, including but
not limited to compiled object code, generated documentation,
and conversions to other media types.
"Work" shall mean the work of authorship, whether in Source or
Object form, made available under the License, as indicated by a
copyright notice that is included in or attached to the work
(an example is provided in the Appendix below).
"Derivative Works" shall mean any work, whether in Source or Object
form, that is based on (or derived from) the Work and for which the
editorial revisions, annotations, elaborations, or other modifications
represent, as a whole, an original work of authorship. For the purposes
of this License, Derivative Works shall not include works that remain
separable from, or merely link (or bind by name) to the interfaces of,
the Work and Derivative Works thereof.
"Contribution" shall mean any work of authorship, including
the original version of the Work and any modifications or additions
to that Work or Derivative Works thereof, that is intentionally
submitted to Licensor for inclusion in the Work by the copyright owner
or by an individual or Legal Entity authorized to submit on behalf of
the copyright owner. For the purposes of this definition, "submitted"
means any form of electronic, verbal, or written communication sent
to the Licensor or its representatives, including but not limited to
communication on electronic mailing lists, source code control systems,
and issue tracking systems that are managed by, or on behalf of, the
Licensor for the purpose of discussing and improving the Work, but
excluding communication that is conspicuously marked or otherwise
designated in writing by the copyright owner as "Not a Contribution."
"Contributor" shall mean Licensor and any individual or Legal Entity
on behalf of whom a Contribution has been received by Licensor and
subsequently incorporated within the Work.
2. Grant of Copyright License. Subject to the terms and conditions of
this License, each Contributor hereby grants to You a perpetual,
worldwide, non-exclusive, no-charge, royalty-free, irrevocable
copyright license to reproduce, prepare Derivative Works of,
publicly display, publicly perform, sublicense, and distribute the
Work and such Derivative Works in Source or Object form.
3. Grant of Patent License. Subject to the terms and conditions of
this License, each Contributor hereby grants to You a perpetual,
worldwide, non-exclusive, no-charge, royalty-free, irrevocable
(except as stated in this section) patent license to make, have made,
use, offer to sell, sell, import, and otherwise transfer the Work,
where such license applies only to those patent claims licensable
by such Contributor that are necessarily infringed by their
Contribution(s) alone or by combination of their Contribution(s)
with the Work to which such Contribution(s) was submitted. If You
institute patent litigation against any entity (including a
cross-claim or counterclaim in a lawsuit) alleging that the Work
or a Contribution incorporated within the Work constitutes direct
or contributory patent infringement, then any patent licenses
granted to You under this License for that Work shall terminate
as of the date such litigation is filed.
4. Redistribution. You may reproduce and distribute copies of the
Work or Derivative Works thereof in any medium, with or without
modifications, and in Source or Object form, provided that You
meet the following conditions:
(a) You must give any other recipients of the Work or
Derivative Works a copy of this License; and
(b) You must cause any modified files to carry prominent notices
stating that You changed the files; and
(c) You must retain, in the Source form of any Derivative Works
that You distribute, all copyright, patent, trademark, and
attribution notices from the Source form of the Work,
excluding those notices that do not pertain to any part of
the Derivative Works; and
(d) If the Work includes a "NOTICE" text file as part of its
distribution, then any Derivative Works that You distribute must
include a readable copy of the attribution notices contained
within such NOTICE file, excluding those notices that do not
pertain to any part of the Derivative Works, in at least one
of the following places: within a NOTICE text file distributed
as part of the Derivative Works; within the Source form or
documentation, if provided along with the Derivative Works; or,
within a display generated by the Derivative Works, if and
wherever such third-party notices normally appear. The contents
of the NOTICE file are for informational purposes only and
do not modify the License. You may add Your own attribution
notices within Derivative Works that You distribute, alongside
or as an addendum to the NOTICE text from the Work, provided
that such additional attribution notices cannot be construed
as modifying the License.
You may add Your own copyright statement to Your modifications and
may provide additional or different license terms and conditions
for use, reproduction, or distribution of Your modifications, or
for any such Derivative Works as a whole, provided Your use,
reproduction, and distribution of the Work otherwise complies with
the conditions stated in this License.
5. Submission of Contributions. Unless You explicitly state otherwise,
any Contribution intentionally submitted for inclusion in the Work
by You to the Licensor shall be under the terms and conditions of
this License, without any additional terms or conditions.
Notwithstanding the above, nothing herein shall supersede or modify
the terms of any separate license agreement you may have executed
with Licensor regarding such Contributions.
6. Trademarks. This License does not grant permission to use the trade
names, trademarks, service marks, or product names of the Licensor,
except as required for reasonable and customary use in describing the
origin of the Work and reproducing the content of the NOTICE file.
7. Disclaimer of Warranty. Unless required by applicable law or
agreed to in writing, Licensor provides the Work (and each
Contributor provides its Contributions) on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
implied, including, without limitation, any warranties or conditions
of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
PARTICULAR PURPOSE. You are solely responsible for determining the
appropriateness of using or redistributing the Work and assume any
risks associated with Your exercise of permissions under this License.
8. Limitation of Liability. In no event and under no legal theory,
whether in tort (including negligence), contract, or otherwise,
unless required by applicable law (such as deliberate and grossly
negligent acts) or agreed to in writing, shall any Contributor be
liable to You for damages, including any direct, indirect, special,
incidental, or consequential damages of any character arising as a
result of this License or out of the use or inability to use the
Work (including but not limited to damages for loss of goodwill,
work stoppage, computer failure or malfunction, or any and all
other commercial damages or losses), even if such Contributor
has been advised of the possibility of such damages.
9. Accepting Warranty or Additional Liability. While redistributing
the Work or Derivative Works thereof, You may choose to offer,
and charge a fee for, acceptance of support, warranty, indemnity,
or other liability obligations and/or rights consistent with this
License. However, in accepting such obligations, You may act only
on Your own behalf and on Your sole responsibility, not on behalf
of any other Contributor, and only if You agree to indemnify,
defend, and hold each Contributor harmless for any liability
incurred by, or claims asserted against, such Contributor by reason
of your accepting any such warranty or additional liability.
END OF TERMS AND CONDITIONS
APPENDIX: How to apply the Apache License to your work.
To apply the Apache License to your work, attach the following
boilerplate notice, with the fields enclosed by brackets "[]"
replaced with your own identifying information. (Don't include
the brackets!) The text should be enclosed in the appropriate
comment syntax for the file format. We also recommend that a
file or class name and description of purpose be included on the
same "printed page" as the copyright notice for easier
identification within third-party archives.
Copyright [yyyy] [name of copyright owner]
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
@@ -0,0 +1,112 @@
# Pyraformer: Low-complexity Pyramidal Attention for Long-range Time Series Modeling and Forecasting
This is the Pytorch implementation of Pyraformer (Pyramidal Attention based Transformer) in the ICLR paper: [Pyraformer: Low-complexity Pyramidal Attention for Long-range Time Series Modeling and Forecasting](https://openreview.net/pdf?id=0EXmFzUn5I).
From https://github.com/alipay/Pyraformer
![The network architecture of Pyraformer.](./img/Figure_1.png)
<center><b>Figure 1.</b> The network architecture of Pyraformer.</center>
## Pyramidal Attention
As demonstrated in Figure 2, we leverage a pyramidal graph to describe the temporal dependencies of the observed time series in a multiresolution fashion. We can decompose the pyramidal graph into two parts: the inter-scale and the intra-scale connections. The inter-scale connections form a C-ary tree, in which each parent has C children. For example, if we associate the finest scale of the pyramidal graph with hourly observations of the original time series, the nodes at coarser scales can be regarded as the daily, weekly, and even monthly features of the time series. As a consequence, the pyramidal graph offers a multiresolution representation of the original time series. Furthermore, it is easier to capture long-range dependencies (e.g., monthly dependence) in the coarser scales by simply connecting the neighboring nodes via the intra-scale connections. In other words, the coarser scales are instrumental in describing long-range correlations in a manner that is graphically far more parsimonious than could be solely captured with a single, finest scale model.
![The Pyramidal Attention Mechanism.](./img/Figure_2.png#center)
<center><b>Figure 2.</b> The Pyramidal Attention Mechanism.</center>
## Requirements
* Python 3.7
* pytorch 1.8.0
* CUDA 11.1
* TVM 0.8.0 (optional)
Dependencies can be installed by:
pip install -r requirements.txt
If you are using CUDA 11.1, you can use the compiled TVM runtime version in the our code to run PAM-TVM. Due to the short history length in the experiments, PAM-TVM does not provide a speed increase. If you want to compile our PAM-TVM kernel yourself, see [here](https://tvm.apache.org/docs/install/index.html) to compile TVM 0.8.0 first.
## Data preparetion
The four datasets (Electricity, Wind, ETT and App Flow) used in this paper can be downloaded from the following links:
* [Electricity](https://archive.ics.uci.edu/ml/datasets/ElectricityLoadDiagrams20112014)
* [Wind](https://www.kaggle.com/sohier/30-years-of-european-wind-generation)
* [ETT](https://github.com/zhouhaoyi/ETDataset)
* [App Flow](https://github.com/alipay/Pyraformer/blob/master/data/app_zone_rpc_hour_encrypted.zip)
The downloaded datasets can be put in the 'data' directory. For single step forecasting, we preprocess Electricity, Wind and App Flow using scripts preprocess_elect.py, preprocess_wind.py and preprocess_flow.py respectively. You can also download preprocessed data [here](https://drive.google.com/drive/folders/1-b9tR6Tgmx48smPMetzAhVSV7-95im3X?usp=sharing). and put them in the 'data' directory. The directory structure looks like:
${CODE_ROOT}
......
|-- data
|-- elect
|-- test_data_elect.npy
|-- train_data_elect.npy
......
|-- flow
......
|-- wind
......
|-- ETT
|-- ETTh1.csv
|-- ETTh2.csv
|-- ETTm1.csv
|-- ETTm2.csv
|-- LD2011_2014.txt
|-- synthetic.npy
Where synthetic.npy is generated by running:
python simulate_sin.py
## Training
To perform long-range forecasting, run:
sh scripts/Pyraformer_LR_FC.sh
To perform single step forecasting, run:
sh scripts/Pyraformer_SS.sh
The meaning of each command line argument is explained in long_range_main.py and single_step_main.py, respectively.
## Evaluate
Evaluation can be done by adding the -eval option to the command line. We provide pretrained models [here](https://drive.google.com/drive/folders/15av5ZhHG8tbX8HuxZNNDGBybdnuxzA83?usp=sharing). The downloaded models should be put in the 'models' directory. The directory structure is as follows:
${CODE_ROOT}
......
|-- models
|-- LongRange
|-- elect
|-- 168
|-- best_iter0.pth
|-- best_iter1.pth
|-- best_iter2.pth
|-- best_iter3.pth
|-- best_iter4.pth
|-- 336
......
|-- 720
......
|-- ETTh1
......
|-- ETTm1
......
|-- SingleStep
|-- elect
|-- best_model.pth
|-- flow
|-- best_model.pth
|-- wind
|-- best_model.pth
Below are evaluation examples:
python long_range_main.py -data ETTh1 -input_size 168 -predict_step 168 -n_head 6 -eval
python single_step_main.py -data_path data/elect/ -dataset elect -eval
## Citation
@inproceedings{liu2022pyraformer,
title={Pyraformer: Low-Complexity Pyramidal Attention for Long-Range Time Series Modeling and Forecasting},
author={Liu, Shizhan and Yu, Hang and Liao, Cong and Li, Jianguo and Lin, Weiyao and Liu, Alex X and Dustdar, Schahram},
booktitle={International Conference on Learning Representations},
year={2022}
}
@@ -0,0 +1,719 @@
import os
import pandas as pd
from torch.utils.data import Dataset, DataLoader
from utils.tools import StandardScaler
from utils.timefeatures import time_features
import numpy as np
import torch
import warnings
warnings.filterwarnings('ignore')
"""Long range dataloader"""
class Dataset_ETT_hour(Dataset):
def __init__(self, root_path, flag='train', size=None, data_path='ETTh1.csv', dataset='ETTh1', inverse=False):
# size [seq_len, label_len, pred_len]
# info
if size == None:
self.seq_len = 24*4*4
self.pred_len = 24*4
else:
self.seq_len = size[0]
self.pred_len = size[1]
# init
assert flag in ['train', 'test', 'val']
type_map = {'train':0, 'val':1, 'test':2}
self.set_type = type_map[flag]
self.inverse = inverse
self.root_path = root_path
self.data_path = data_path
self.__read_data__()
def __read_data__(self):
self.scaler = StandardScaler()
df_raw = pd.read_csv(os.path.join(self.root_path,
self.data_path))
border1s = [0, 12*30*24 - self.seq_len, 12*30*24+4*30*24 - self.seq_len]
border2s = [12*30*24, 12*30*24+4*30*24, 12*30*24+8*30*24]
border1 = border1s[self.set_type]
border2 = border2s[self.set_type]
cols_data = df_raw.columns[1:]
df_data = df_raw[cols_data]
train_data = df_data[border1s[0]:border2s[0]]
self.scaler.fit(train_data.values)
data = self.scaler.transform(df_data.values)
df_stamp = df_raw[['date']][border1:border2]
df_stamp['date'] = pd.to_datetime(df_stamp.date)
data_stamp = time_features(df_stamp, timeenc=1, freq='h')
self.data_x = data[border1:border2]
if self.inverse:
self.data_y = df_data.values[border1:border2]
else:
self.data_y = data[border1:border2]
self.data_stamp = data_stamp
def __getitem__(self, index):
s_begin = index
s_end = s_begin + self.seq_len
r_begin = s_end
r_end = r_begin + self.pred_len
seq_x = self.data_x[s_begin:s_end]
seq_y = self.data_y[r_begin:r_end]
seq_x_mark = self.data_stamp[s_begin:s_end]
seq_y_mark = self.data_stamp[r_begin:r_end]
return seq_x, seq_y, seq_x_mark, seq_y_mark, self.scaler.mean, self.scaler.std
def __len__(self):
return len(self.data_x) - self.seq_len- self.pred_len + 1
def inverse_transform(self, data, seq_y, mean, std):
return self.scaler.inverse_transform(data), seq_y
class Dataset_ETT_minute(Dataset):
def __init__(self, root_path, flag='train', size=None, data_path='ETTm1.csv', dataset='ETTm1', inverse=False):
# size [seq_len, label_len, pred_len]
# info
if size == None:
self.seq_len = 24*4*4
self.pred_len = 24*4
else:
self.seq_len = size[0]
self.pred_len = size[1]
# init
assert flag in ['train', 'test', 'val']
type_map = {'train':0, 'val':1, 'test':2}
self.set_type = type_map[flag]
self.inverse = inverse
self.root_path = root_path
self.data_path = data_path
self.__read_data__()
def __read_data__(self):
self.scaler = StandardScaler()
df_raw = pd.read_csv(os.path.join(self.root_path,
self.data_path))
border1s = [0, 12*30*24*4 - self.seq_len, 12*30*24*4+4*30*24*4 - self.seq_len]
border2s = [12*30*24*4, 12*30*24*4+4*30*24*4, 12*30*24*4+8*30*24*4]
border1 = border1s[self.set_type]
border2 = border2s[self.set_type]
cols_data = df_raw.columns[1:]
df_data = df_raw[cols_data]
train_data = df_data[border1s[0]:border2s[0]]
self.scaler.fit(train_data.values)
data = self.scaler.transform(df_data.values)
df_stamp = df_raw[['date']][border1:border2]
df_stamp['date'] = pd.to_datetime(df_stamp.date)
data_stamp = time_features(df_stamp, timeenc=1, freq='h')
self.data_x = data[border1:border2]
if self.inverse:
self.data_y = df_data.values[border1:border2]
else:
self.data_y = data[border1:border2]
self.data_stamp = data_stamp
def __getitem__(self, index):
s_begin = index
s_end = s_begin + self.seq_len
r_begin = s_end
r_end = r_begin + self.pred_len
seq_x = self.data_x[s_begin:s_end]
seq_y = self.data_y[r_begin:r_end]
seq_x_mark = self.data_stamp[s_begin:s_end]
seq_y_mark = self.data_stamp[r_begin:r_end]
return seq_x, seq_y, seq_x_mark, seq_y_mark, self.scaler.mean, self.scaler.std
def __len__(self):
return len(self.data_x) - self.seq_len- self.pred_len + 1
def inverse_transform(self, data, seq_y, mean, std):
return self.scaler.inverse_transform(data), seq_y
class Dataset_Custom(Dataset):
def __init__(self, root_path, flag='train', size=None,
features='M', data_path='ETTh1.csv',
target='OT', scale=True, dataset='',timeenc=0, freq='h',inverse=False):
# size [seq_len, label_len, pred_len]
# info
if size == None:
self.seq_len = 24 * 4 * 4
self.pred_len = 24 * 4
else:
self.seq_len = size[0]
self.pred_len = size[1]
# init
assert flag in ['train', 'test', 'val']
type_map = {'train': 0, 'val': 1, 'test': 2}
self.set_type = type_map[flag]
self.features = features
self.target = target
self.scale = scale
self.timeenc = timeenc
self.freq = freq
self.root_path = root_path
self.data_path = data_path
self.__read_data__()
def __read_data__(self):
self.scaler = StandardScaler()
df_raw = pd.read_csv(os.path.join(self.root_path,
self.data_path))
'''
df_raw.columns: ['date', ...(other features), target feature]
'''
cols = list(df_raw.columns)
cols.remove(self.target)
cols.remove('date')
df_raw = df_raw[['date'] + cols + [self.target]]
# print(cols)
num_train = int(len(df_raw) * 0.7)
num_test = int(len(df_raw) * 0.2)
num_vali = len(df_raw) - num_train - num_test
border1s = [0, num_train - self.seq_len, len(df_raw) - num_test - self.seq_len]
border2s = [num_train, num_train + num_vali, len(df_raw)]
border1 = border1s[self.set_type]
border2 = border2s[self.set_type]
if self.features == 'M' or self.features == 'MS':
cols_data = df_raw.columns[1:]
df_data = df_raw[cols_data]
elif self.features == 'S':
df_data = df_raw[[self.target]]
if self.scale:
train_data = df_data[border1s[0]:border2s[0]]
self.scaler.fit(train_data.values)
# print(self.scaler.mean_)
# exit()
data = self.scaler.transform(df_data.values)
else:
data = df_data.values
df_stamp = df_raw[['date']][border1:border2]
df_stamp['date'] = pd.to_datetime(df_stamp.date)
if self.timeenc == 0:
df_stamp['month'] = df_stamp.date.apply(lambda row: row.month, 1)
df_stamp['day'] = df_stamp.date.apply(lambda row: row.day, 1)
df_stamp['weekday'] = df_stamp.date.apply(lambda row: row.weekday(), 1)
df_stamp['hour'] = df_stamp.date.apply(lambda row: row.hour, 1)
data_stamp = df_stamp.drop(['date'], axis=1).values
elif self.timeenc == 1:
data_stamp = time_features(pd.to_datetime(df_stamp['date'].values), freq=self.freq)
data_stamp = data_stamp.transpose(1, 0)
self.data_x = data[border1:border2]
self.data_y = data[border1:border2]
self.data_stamp = data_stamp
def __getitem__(self, index):
s_begin = index
s_end = s_begin + self.seq_len
r_begin = s_end
r_end = r_begin + self.pred_len
seq_x = self.data_x[s_begin:s_end]
seq_y = self.data_y[r_begin:r_end]
seq_x_mark = self.data_stamp[s_begin:s_end]
seq_y_mark = self.data_stamp[r_begin:r_end]
return seq_x, seq_y, seq_x_mark, seq_y_mark,self.scaler.mean, self.scaler.std
def __len__(self):
return len(self.data_x) - self.seq_len - self.pred_len + 1
def inverse_transform(self, data):
return self.scaler.inverse_transform(data)
# """Long range dataloader for dataset elect and app flow"""
class Dataset_Custom2(Dataset):
def __init__(self, root_path, flag='train', size=None, data_path='ETTh1.csv', dataset='elect',
inverse=False):
# size [seq_len, label_len, pred_len]
# info
self.seq_len = size[0]
self.pred_len = size[1]
# init
assert flag in ['train', 'test']
self.flag = flag
self.inverse = inverse
self.root_path = root_path
self.data_path = data_path
preprocess_path = os.path.join(self.root_path, self.data_path)
self.all_data, self.covariates, self.train_end = eval('preprocess_flow')(preprocess_path)
self.all_data = torch.from_numpy(self.all_data).transpose(0, 1)
self.covariates = torch.from_numpy(self.covariates)
self.test_start = self.train_end - self.seq_len + 1
self.window_stride = 24
self.seq_num = self.all_data.size(0)
def fit(self, data):
mean = data.mean()
std = data.std()
return mean, std
def inverse_transform(self, output, seq_y, mean, std):
output = output * (mean.unsqueeze(1).unsqueeze(1) + 1)
seq_y = seq_y * (mean.unsqueeze(1).unsqueeze(1) + 1)
return output, seq_y
def __len__(self):
if self.flag == 'train':
self.window_per_seq = (self.train_end - self.seq_len - self.pred_len) // self.window_stride
return self.window_per_seq * self.seq_num
else:
self.window_per_seq = (self.all_data.size(1) - self.test_start - self.seq_len - self.pred_len) // self.window_stride
return self.window_per_seq * self.seq_num
def __getitem__(self, index):
seq_idx = index // self.window_per_seq
window_idx = index % self.window_per_seq
if self.flag == 'train':
s_begin = window_idx * self.window_stride
else:
s_begin = self.test_start + window_idx * self.window_stride
s_end = s_begin + self.seq_len
r_begin = s_end
r_end = r_begin + self.pred_len
seq_x = self.all_data[seq_idx, s_begin:s_end].clone()
seq_y = self.all_data[seq_idx, r_begin:r_end].clone()
mean, std = self.fit(seq_x)
if mean > 0:
seq_x = seq_x / (mean + 1)
seq_y = seq_y / (mean + 1)
if len(self.covariates.size()) == 2:
seq_x_mark = self.covariates[s_begin:s_end]
seq_x_mark[:, -1] = int(seq_idx)
seq_y_mark = self.covariates[r_begin:r_end]
seq_y_mark[:, -1] = int(seq_idx)
else:
seq_x_mark = self.covariates[s_begin:s_end, seq_idx]
seq_x_mark[:, -1] = int(seq_idx)
seq_y_mark = self.covariates[r_begin:r_end, seq_idx]
seq_y_mark[:, -1] = int(seq_idx)
return seq_x.unsqueeze(1), seq_y.unsqueeze(1), seq_x_mark, seq_y_mark, mean, std
"""Long range dataloader for synthetic dataset"""
class Dataset_Synthetic(Dataset):
def __init__(self, root_path, flag='train', size=None, data_path='synthetic.npy', dataset='synthetic', inverse=False):
# size [seq_len, label_len, pred_len]
# info
self.seq_len = size[0]
self.pred_len = size[1]
# init
assert flag in ['train', 'test']
self.flag = flag
self.inverse = inverse
self.root_path = root_path
self.data_path = data_path
preprocess_path = os.path.join(self.root_path, self.data_path)
self.all_data =np.load(preprocess_path)
self.all_data = torch.from_numpy(self.all_data)
self.all_data, self.covariates = self.all_data[:, :, 0], self.all_data[:, :, 1:]
self.seq_num = self.all_data.size(0)
self.window_stride = 24
window_per_seq = (self.all_data.shape[1] - self.seq_len - self.pred_len) / self.window_stride
self.train_end = self.seq_len + self.pred_len + int(0.9 * window_per_seq) * self.window_stride
self.test_start = self.train_end - self.seq_len + 1
def fit(self, data):
mean = data.mean()
std = data.std()
return mean, std
def inverse_transform(self, output, seq_y, mean, std):
output = output * (mean.unsqueeze(1).unsqueeze(1) + 1)
seq_y = seq_y * (mean.unsqueeze(1).unsqueeze(1) + 1)
return output, seq_y
def __len__(self):
if self.flag == 'train':
self.window_per_seq = (self.train_end - self.seq_len - self.pred_len) // self.window_stride
return self.window_per_seq * self.seq_num
else:
self.window_per_seq = (self.all_data.size(1) - self.test_start - self.seq_len - self.pred_len) // self.window_stride
return self.window_per_seq * self.seq_num
def __getitem__(self, index):
seq_idx = index // self.window_per_seq
window_idx = index % self.window_per_seq
if self.flag == 'train':
s_begin = window_idx * self.window_stride
else:
s_begin = self.test_start + window_idx * self.window_stride
s_end = s_begin + self.seq_len
r_begin = s_end
r_end = r_begin + self.pred_len
seq_x = self.all_data[seq_idx, s_begin:s_end].clone()
seq_y = self.all_data[seq_idx, r_begin:r_end].clone()
mean, std = self.fit(seq_x)
if mean > 0:
seq_x = seq_x / (mean + 1)
seq_y = seq_y / (mean + 1)
seq_x_mark = self.covariates[seq_idx, s_begin:s_end]
seq_y_mark = self.covariates[seq_idx, r_begin:r_end]
return seq_x.unsqueeze(1), seq_y.unsqueeze(1), seq_x_mark, seq_y_mark, mean, std
def get_all_v(train_data, train_end, seq_len, pred_len, window_stride, type):
"""Get the normalization parameters of each sequence"""
seq_num = train_data.size(0)
window_per_seq = (train_end - seq_len - pred_len) // window_stride
window_number = seq_num * window_per_seq
v = torch.zeros(window_number, dtype=torch.float64)
for index in range(window_number):
seq_idx = index // window_per_seq
window_idx = index % window_per_seq
s_begin = window_idx * window_stride
s_end = s_begin + seq_len
seq_x = train_data[seq_idx, s_begin:s_end].clone()
if type == 'mean':
mean = seq_x.mean()
v[index] = mean + 1
else:
std = seq_x.std()
v[index] = std
return v
def gen_covariates(times, num_covariates):
"""Get covariates"""
covariates = np.zeros((times.shape[0], num_covariates))
for i, input_time in enumerate(times):
covariates[i, 0] = input_time.weekday() / 7
covariates[i, 1] = input_time.hour / 24
covariates[i, 2] = input_time.month / 12
return covariates
def preprocess_elect(csv_path):
"""preprocess the elect dataset for long range forecasting"""
num_covariates = 4
train_start = '2011-01-01 00:00:00'
train_end = '2014-04-01 23:00:00'
test_start = '2014-04-01 00:00:00'
test_end = '2014-09-07 23:00:00'
data_frame = pd.read_csv(csv_path, sep=";", index_col=0, parse_dates=True, decimal=',')
data_frame = data_frame.resample('1H',label = 'left',closed = 'right').sum()[train_start:test_end]
data_frame.fillna(0, inplace=True)
covariates = gen_covariates(data_frame[train_start:test_end].index, num_covariates)
all_data = data_frame[train_start:test_end].values
data_start = (all_data!=0).argmax(axis=0) #find first nonzero value in each time series
train_end = len(data_frame[train_start:train_end].values)
all_data = all_data[:, data_start < 10000]
data_start = data_start[data_start < 10000]
split_start = data_start.max()
all_data = all_data[split_start:]
covariates = covariates[split_start:]
train_end = train_end - split_start
return all_data.astype(np.float32), covariates.astype(np.float32), train_end
def preprocess_flow(csv_path):
"""preprocess the app flow dataset for long range forecasting"""
data_frame = pd.read_csv(csv_path, names=['app_name', 'zone', 'time', 'value'], parse_dates=True)
grouped_data = list(data_frame.groupby(["app_name", "zone"]))
# covariates = gen_covariates(data_frame.index, 3)
all_data = []
min_length = 10000
for i in range(len(grouped_data)):
single_df = grouped_data[i][1].drop(labels=['app_name', 'zone'], axis=1).sort_values(by="time", ascending=True)
times = pd.to_datetime(single_df.time)
single_df['weekday'] = times.dt.dayofweek / 7
single_df['hour'] = times.dt.hour / 24
single_df['month'] = times.dt.month / 12
temp_data = single_df.values[:, 1:]
if (temp_data[:, 0] == 0).sum() / len(temp_data) > 0.2 or len(temp_data) < 3000:
continue
if len(temp_data) < min_length:
min_length = len(temp_data)
all_data.append(temp_data)
all_data = np.array([data[len(data)-min_length:, :] for data in all_data]).transpose(1, 0, 2).astype(np.float32)
train_end = min(int(0.8 * min_length), min_length - 1000)
covariates = all_data.copy()
covariates[:, :, :-1] = covariates[:, :, 1:]
return all_data[:, :, 0], covariates, train_end
"""Single step dataloader"""
def split(split_start, label, cov, pred_length):
all_data = []
for batch_idx in range(len(label)):
batch_label = label[batch_idx]
for i in range(pred_length):
single_data = batch_label[i:(split_start+i)].clone().unsqueeze(1)
single_data[-1] = -1
single_cov = cov[batch_idx, i:(split_start+i), :].clone()
temp_data = [single_data, single_cov]
single_data = torch.cat(temp_data, dim=1)
all_data.append(single_data)
data = torch.stack(all_data, dim=0)
label = label[:, -pred_length:].reshape(pred_length*len(label))
return data, label
"""Single step training dataloader for the electricity dataset"""
class electTrainDataset(Dataset):
def __init__(self, data_path, data_name, predict_length, batch_size):
self.data = torch.from_numpy(np.load(os.path.join(data_path, f'train_data_{data_name}.npy')))
# Resample windows according to the average amplitude
v = np.load(os.path.join(data_path, f'train_v_{data_name}.npy'))
weights = torch.as_tensor(np.abs(v[:,0])/np.sum(np.abs(v[:,0])), dtype=torch.double)
num_samples = weights.size(0)
sample_index = torch.multinomial(weights, num_samples, True)
self.data = self.data[sample_index]
self.label = torch.from_numpy(np.load(os.path.join(data_path, f'train_label_{data_name}.npy')))
self.label = self.label[sample_index]
self.train_len = len(self.data) // batch_size
self.pred_length = predict_length
self.batch_size = batch_size
def __len__(self):
return self.train_len
def __getitem__(self, index):
if (index+1) <= self.train_len:
all_data = self.data[index*self.batch_size:(index+1)*self.batch_size].clone()
label = self.label[index*self.batch_size:(index+1)*self.batch_size].clone()
else:
all_data = self.data[index*self.batch_size:].clone()
label = self.label[index*self.batch_size:].clone()
cov = all_data[:, :, 2:]
split_start = len(label[0]) - self.pred_length + 1
data, label = split(split_start, label, cov, self.pred_length)
return data, label
"""Single step testing dataloader for the electricity dataset"""
class electTestDataset(Dataset):
def __init__(self, data_path, data_name, predict_length):
self.data = np.load(os.path.join(data_path, f'test_data_{data_name}.npy'))
self.v = np.load(os.path.join(data_path, f'test_v_{data_name}.npy'))
self.label = np.load(os.path.join(data_path, f'test_label_{data_name}.npy'))
self.test_len = self.data.shape[0]
self.pred_length = predict_length
def __len__(self):
return self.test_len
def __getitem__(self, index):
all_data = torch.from_numpy(self.data[index].copy())
cov = all_data[:, 2:]
label = torch.from_numpy(self.label[index].copy())
v = float(self.v[index][0])
if v > 0:
data = label / v
else:
data = label
split_start = len(label) - self.pred_length + 1
all_data = []
for i in range(self.pred_length):
single_data = data[i:(split_start+i)].clone().unsqueeze(1)
single_data[-1] = -1
single_cov = cov[i:(split_start+i), :].clone()
single_data = torch.cat([single_data, single_cov], dim=1)
all_data.append(single_data)
all_data = torch.stack(all_data, dim=0)
label = label[-self.pred_length:]
return all_data, label, v
"""Single step training dataloader for the app flow dataset"""
class flowTrainDataset(Dataset):
def __init__(self, data_path, data_name, predict_length, batch_size):
self.data = torch.from_numpy(np.load(os.path.join(data_path, f'train_data_{data_name}.npy')))
# Resample windows according to the average amplitude
v = np.load(os.path.join(data_path, f'train_v_{data_name}.npy'))
weights = torch.as_tensor(np.abs(v)/np.sum(np.abs(v)), dtype=torch.double)
num_samples = weights.size(0)
sample_index = torch.multinomial(weights, num_samples, True)
self.data = self.data[sample_index]
self.label = self.data[:, :, 0]
self.train_len = len(self.data) // batch_size
self.pred_length = predict_length
self.batch_size = batch_size
def __len__(self):
return self.train_len
def __getitem__(self, index):
if (index+1) <= self.train_len:
all_data = self.data[index*self.batch_size:(index+1)*self.batch_size].clone()
label = self.label[index*self.batch_size:(index+1)*self.batch_size].clone()
else:
all_data = self.data[index*self.batch_size:].clone()
label = self.label[index*self.batch_size:].clone()
cov = all_data[:, :, 1:]
split_start = len(label[0]) - self.pred_length + 1
data, label = split(split_start, label, cov, self.pred_length)
return data, label
"""Single step testing dataloader for the all flow dataset"""
class flowTestDataset(Dataset):
def __init__(self, data_path, data_name, predict_length):
self.data = np.load(os.path.join(data_path, f'test_data_{data_name}.npy'))
self.v = np.load(os.path.join(data_path, f'test_v_{data_name}.npy'))
self.label = self.data
self.test_len = self.data.shape[0]
self.pred_length = predict_length
def __len__(self):
return self.test_len
def __getitem__(self, index):
all_data = torch.from_numpy(self.data[index].copy())
cov = all_data[:, 1:]
data = all_data[:, 0]
label = torch.from_numpy(self.label[index, :, 0].copy())
v = float(self.v[index])
split_start = len(label) - self.pred_length + 1
all_data = []
for i in range(self.pred_length):
single_data = data[i:(split_start+i)].clone().unsqueeze(1)
single_data[-1] = -1
single_cov = cov[i:(split_start+i), :].clone()
single_data = torch.cat([single_data, single_cov], dim=1)
all_data.append(single_data)
all_data = torch.stack(all_data, dim=0)
label = label[-self.pred_length:] * v
return all_data, label, v
"""Single step training dataloader for the wind dataset"""
class windTrainDataset(Dataset):
def __init__(self, data_path, data_name, predict_length, batch_size):
self.data = torch.from_numpy(np.load(os.path.join(data_path, f'train_data_{data_name}.npy')))
# Resample windows according to the average amplitude
v = np.load(os.path.join(data_path, f'train_v_{data_name}.npy'))
weights = torch.as_tensor(np.abs(v)/np.sum(np.abs(v)), dtype=torch.double)
num_samples = weights.size(0)
sample_index = torch.multinomial(weights, num_samples, True)
self.data = self.data[sample_index]
self.train_len = len(self.data) // batch_size
self.pred_length = predict_length
self.batch_size = batch_size
def __len__(self):
return self.train_len
def __getitem__(self, index):
if (index+1) <= self.train_len:
all_data = self.data[index*self.batch_size:(index+1)*self.batch_size].clone()
else:
all_data = self.data[index*self.batch_size:].clone()
cov = all_data[:, :, 1:]
label = all_data[:, :, 0]
split_start = len(label[0]) - self.pred_length + 1
data, label = split(split_start, label, cov, self.pred_length)
return data, label
"""Single step testing dataloader for the wind dataset"""
class windTestDataset(Dataset):
def __init__(self, data_path, data_name, predict_length):
self.data = np.load(os.path.join(data_path, f'test_data_{data_name}.npy'))
self.v = np.load(os.path.join(data_path, f'test_v_{data_name}.npy'))
self.test_len = self.data.shape[0]
self.pred_length = predict_length
def __len__(self):
return self.test_len
def __getitem__(self, index):
all_data = torch.from_numpy(self.data[index].copy())
cov = all_data[:, 1:]
data = all_data[:, 0]
v = float(self.v[index])
label = data * v
split_start = len(label) - self.pred_length + 1
all_data = []
for i in range(self.pred_length):
single_data = data[i:(split_start+i)].clone().unsqueeze(1)
single_data[-1] = -1
single_cov = cov[i:(split_start+i), :].clone()
single_data = torch.cat([single_data, single_cov], dim=1)
all_data.append(single_data)
all_data = torch.stack(all_data, dim=0)
label = label[-self.pred_length:]
return all_data, label, v
@@ -0,0 +1,409 @@
import argparse
import numpy as np
import time
import torch
import torch.optim as optim
import pyraformer.Pyraformer_LR as Pyraformer
from tqdm import tqdm
from data_loader import *
from utils.tools import TopkMSELoss, metric
def prepare_dataloader(args):
""" Load data and prepare dataloader. """
data_dict = {
'ETTh1':Dataset_ETT_hour,
'ETTh2':Dataset_ETT_hour,
'ETTm1':Dataset_ETT_minute,
'ETTm2':Dataset_ETT_minute,
'electricity':Dataset_Custom,
'exchange':Dataset_Custom,
'traffic':Dataset_Custom,
'weather':Dataset_Custom,
'ili':Dataset_Custom,
# 'flow': Dataset_Custom2,
# 'synthetic': Dataset_Synthetic,
}
Data = data_dict[args.data]
# prepare training dataset and dataloader
shuffle_flag = True; drop_last = True; batch_size = args.batch_size
train_set = Data(
root_path=args.root_path,
data_path=args.data_path,
flag='train',
size=[args.input_size, args.predict_step],
inverse=args.inverse,
dataset=args.data
)
print('train', len(train_set))
train_loader = DataLoader(
train_set,
batch_size=batch_size,
shuffle=shuffle_flag,
num_workers=0,
drop_last=drop_last)
# prepare testing dataset and dataloader
shuffle_flag = False; drop_last = False; batch_size = args.batch_size
test_set = Data(
root_path=args.root_path,
data_path=args.data_path,
flag='test',
size=[args.input_size, args.predict_step],
inverse=args.inverse,
dataset=args.data
)
print('test', len(test_set))
test_loader = DataLoader(
test_set,
batch_size=batch_size,
shuffle=shuffle_flag,
num_workers=0,
drop_last=drop_last)
return train_loader, train_set, test_loader, test_set
def sample_mining_scheduler(epoch, batch_size):
if epoch < 2:
topk = batch_size
elif epoch < 4:
topk = int(batch_size * (5 - epoch) / (6 - epoch))
else:
topk = int(0.5 * batch_size)
return topk
def dataset_parameters(args, dataset):
"""Prepare specific parameters for different datasets"""
dataset2enc_in = {
'ETTh1':7,
'ETTh2':7,
'ETTm1':7,
'ETTm2':7,
'electricity':321,
'exchange':8,
'traffic':862,
'weather':21,
'ili':7,
'flow': 1,
'synthetic': 1
}
dataset2cov_size = {
'ETTh1':4,
'ETTh2':4,
'ETTm1':4,
'ETTm2':4,
'electricity':4,
'exchange':4,
'traffic':4,
'weather':4,
'ili':4,
'elect':3,
'flow': 3,
'synthetic': 3,
}
dataset2seq_num = {
'ETTh1':1,
'ETTh2':1,
'ETTm1':1,
'ETTm2':1,
'electricity':1,
'exchange':1,
'traffic':1,
'weather':1,
'ili':1,
'elect':321,
'flow': 1077,
'synthetic': 60
}
dataset2embed = {
'ETTh1':'DataEmbedding',
'ETTh2':'DataEmbedding',
'ETTm1':'DataEmbedding',
'ETTm2':'DataEmbedding',
'elect':'CustomEmbedding',
'electricity':'CustomEmbedding',
'exchange':'CustomEmbedding',
'traffic':'CustomEmbedding',
'weather':'CustomEmbedding',
'ili':'CustomEmbedding',
'flow': 'CustomEmbedding',
'synthetic': 'CustomEmbedding'
}
args.enc_in = dataset2enc_in[dataset]
args.dec_in = dataset2enc_in[dataset]
args.covariate_size = dataset2cov_size[dataset]
args.seq_num = dataset2seq_num[dataset]
args.embed_type = dataset2embed[dataset]
return args
def train_epoch(model, train_dataset, training_loader, optimizer, opt, epoch):
""" Epoch operation in training phase. """
model.train()
total_loss = 0
total_pred_number = 0
warm = False
for batch in tqdm(training_loader, mininterval=2,
desc=' - (Training) ', leave=False):
# prepare data
batch_x, batch_y, batch_x_mark, batch_y_mark, mean, std = map(lambda x: x.float().to(opt.device), batch)
dec_inp = torch.zeros_like(batch_y).float()
optimizer.zero_grad()
# forward
if opt.decoder == 'attention':
if opt.pretrain and epoch < 1:
outputs = model(batch_x, batch_x_mark, dec_inp, batch_y_mark, True)
batch_y = torch.cat([batch_x, batch_y], dim=1)
else:
outputs = model(batch_x, batch_x_mark, dec_inp, batch_y_mark, False)
elif opt.decoder == 'FC':
# Add a predict token into the history sequence
predict_token = torch.zeros(batch_x.size(0), 1, batch_x.size(-1), device=batch_x.device)
batch_x = torch.cat([batch_x, predict_token], dim=1)
batch_x_mark = torch.cat([batch_x_mark, batch_y_mark[:, 0:1, :]], dim=1)
outputs = model(batch_x, batch_x_mark, dec_inp, batch_y_mark, False)
# determine the loss function
if opt.hard_sample_mining and not (opt.pretrain and epoch < 1):
topk = sample_mining_scheduler(epoch, batch_x.size(0))
criterion = TopkMSELoss(topk)
else:
criterion = torch.nn.MSELoss(reduction='none')
# if inverse, both the output and the ground truth are denormalized.
if opt.inverse:
outputs, batch_y = train_dataset.inverse_transform(outputs, batch_y, mean, std)
# compute loss
losses = criterion(outputs, batch_y)
loss = losses.mean()
loss.backward()
""" update parameters """
optimizer.step()
total_loss += losses.sum().item()
total_pred_number += losses.numel()
return total_loss / total_pred_number
def eval_epoch(model, test_dataset, test_loader, opt, epoch):
""" Epoch operation in evaluation phase. """
model.eval()
preds = []
trues = []
warm = 0
with torch.no_grad():
for batch in tqdm(test_loader, mininterval=2,
desc=' - (Validation) ', leave=False):
""" prepare data """
batch_x, batch_y, batch_x_mark, batch_y_mark, mean, std = map(lambda x: x.float().to(opt.device), batch)
dec_inp = torch.zeros_like(batch_y).float()
# forward
if opt.decoder == 'FC':
# Add a predict token into the history sequence
predict_token = torch.zeros(batch_x.size(0), 1, batch_x.size(-1), device=batch_x.device)
batch_x = torch.cat([batch_x, predict_token], dim=1)
batch_x_mark = torch.cat([batch_x_mark, batch_y_mark[:, 0:1, :]], dim=1)
outputs = model(batch_x, batch_x_mark, dec_inp, batch_y_mark, False)
warm += 1
# if inverse, both the output and the ground truth are denormalized.
if opt.inverse:
outputs, batch_y = test_dataset.inverse_transform(outputs, batch_y, mean, std)
pred = outputs.detach().cpu().numpy()
true = batch_y.detach().cpu().numpy()
preds.append(pred)
trues.append(true)
preds = np.array(preds)
trues = np.array(trues)
# preds = preds.reshape(-1, preds.shape[-2], preds.shape[-1])
# trues = trues.reshape(-1, trues.shape[-2], trues.shape[-1])
preds = np.concatenate(preds, axis=0)
print(preds.shape)
trues = np.concatenate(trues, axis=0)
# np.save('./results/' + 'pred.npy', preds)
# np.save('./results/'+ 'true.npy', trues)
print('test shape:{}'.format(preds.shape))
mae, mse, rmse, mape, mspe = metric(preds, trues)
print('Epoch {}, mse:{}, mae:{}, rmse:{}, mape:{}, mspe:{}'.format(epoch, mse, mae, rmse, mape, mspe))
return mse, mae, rmse, mape, mspe
def train(model, optimizer, scheduler, opt, model_save_dir):
""" Start training. """
best_mse = 100000000
""" prepare dataloader """
training_dataloader, train_dataset, test_dataloader, test_dataset = prepare_dataloader(opt)
best_metrics = []
for epoch_i in range(opt.epoch):
epoch = epoch_i + 1
print('[ Epoch', epoch, ']')
start = time.time()
train_mse = train_epoch(model, train_dataset, training_dataloader, optimizer, opt, epoch_i)
print(' - (Training) '
'MSE: {mse: 8.5f}'
'elapse: {elapse:3.3f} min'
.format(mse=train_mse, elapse=(time.time() - start) / 60))
mse, mae, rmse, mape, mspe = eval_epoch(model, test_dataset, test_dataloader, opt, epoch_i)
scheduler.step()
current_metrics = [float(mse), float(mae), float(rmse), float(mape), float(mspe)]
if best_mse > mse:
best_mse = mse
best_metrics = current_metrics
torch.save(
{
"state_dict": model.state_dict(),
"metrics": best_metrics
},
model_save_dir
)
return best_metrics
def evaluate(model, opt, model_save_dir):
"""Evaluate preptrained models"""
best_mse = 100000000
""" prepare dataloader """
_, _, test_dataloader, test_dataset = prepare_dataloader(opt)
""" load pretrained model """
checkpoint = torch.load(model_save_dir)["state_dict"]
model.load_state_dict(checkpoint)
best_metrics = []
mse, mae, rmse, mape, mspe = eval_epoch(model, test_dataset, test_dataloader, opt, 0)
current_metrics = [float(mse), float(mae), float(rmse), float(mape), float(mspe)]
if best_mse > mse:
best_mse = mse
best_metrics = current_metrics
return best_metrics
def parse_args():
parser = argparse.ArgumentParser()
# running mode
parser.add_argument('-eval', action='store_true', default=False)
# Path parameters
parser.add_argument('-data', type=str, default='ETTh1')
parser.add_argument('-root_path', type=str, default='../dataset/', help='root path of the data file')
parser.add_argument('-data_path', type=str, default='ETTh1.csv', help='data file')
# Dataloader parameters.
parser.add_argument('-input_size', type=int, default=168)
parser.add_argument('-predict_step', type=int, default=168)
parser.add_argument('-inverse', action='store_true', help='denormalize output data', default=False)
# Architecture selection.
parser.add_argument('-model', type=str, default='Pyraformer')
parser.add_argument('-decoder', type=str, default='FC') # selection: [FC, attention]
# Training parameters.
parser.add_argument('-epoch', type=int, default=5)
parser.add_argument('-batch_size', type=int, default=32)
parser.add_argument('-pretrain', action='store_true', default=False)
parser.add_argument('-hard_sample_mining', action='store_true', default=False)
parser.add_argument('-dropout', type=float, default=0.05)
parser.add_argument('-lr', type=float, default=1e-4)
parser.add_argument('-lr_step', type=float, default=0.1)
# Common Model parameters.
parser.add_argument('-d_model', type=int, default=512)
parser.add_argument('-d_inner_hid', type=int, default=512)
parser.add_argument('-d_k', type=int, default=128)
parser.add_argument('-d_v', type=int, default=128)
parser.add_argument('-d_bottleneck', type=int, default=128)
parser.add_argument('-n_head', type=int, default=4)
parser.add_argument('-n_layer', type=int, default=4)
# Pyraformer parameters.
parser.add_argument('-window_size', type=str, default='[4, 4, 4]') # The number of children of a parent node.
parser.add_argument('-inner_size', type=int, default=3) # The number of ajacent nodes.
# CSCM structure. selection: [Bottleneck_Construct, Conv_Construct, MaxPooling_Construct, AvgPooling_Construct]
parser.add_argument('-CSCM', type=str, default='Bottleneck_Construct')
parser.add_argument('-truncate', action='store_true', default=False) # Whether to remove coarse-scale nodes from the attention structure
parser.add_argument('-use_tvm', action='store_true', default=False) # Whether to use TVM.
# Experiment repeat times.
parser.add_argument('-iter_num', type=int, default=1) # Repeat number.
opt = parser.parse_args()
return opt
def main(opt, iter_index):
""" Main function. """
print('[Info] parameters: {}'.format(opt))
if torch.cuda.is_available():
opt.device = torch.device("cuda")
else:
opt.device = torch.device('cpu')
""" prepare model """
model = eval(opt.model).Model(opt)
model.to(opt.device)
""" number of parameters """
num_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
print('[Info] Number of parameters: {}'.format(num_params))
""" train or evaluate the model """
model_save_dir = 'models/LongRange/{}/{}/'.format(opt.data, opt.predict_step)
os.makedirs(model_save_dir, exist_ok=True)
model_save_dir += 'best_iter{}.pth'.format(iter_index)
if opt.eval:
best_metrics = evaluate(model, opt, model_save_dir)
else:
""" optimizer and scheduler """
optimizer = optim.Adam(filter(lambda x: x.requires_grad, model.parameters()), opt.lr)
scheduler = optim.lr_scheduler.StepLR(optimizer, 1, gamma=opt.lr_step)
best_metrics = train(model, optimizer, scheduler, opt, model_save_dir)
print('Iteration best metrics: {}'.format(best_metrics))
return best_metrics
if __name__ == '__main__':
opt = parse_args()
opt = dataset_parameters(opt, opt.data)
opt.window_size = eval(opt.window_size)
iter_num = opt.iter_num
all_perf = []
for i in range(iter_num):
metrics = main(opt, i)
all_perf.append(metrics)
all_perf = np.array(all_perf)
all_perf = all_perf.mean(0)
print('Average Metrics: {}'.format(all_perf))
@@ -0,0 +1,121 @@
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
import os
from datetime import datetime, timedelta
import pandas as pd
import math
import numpy as np
import random
from tqdm import trange
from io import BytesIO
from urllib.request import urlopen
from zipfile import ZipFile
from math import sqrt
from pandas import read_csv, DataFrame
from scipy import stats
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
def prep_data(data, covariates, data_start, train = True):
"""Divide the training sequence into windows"""
time_len = data.shape[0]
input_size = window_size-stride_size
windows_per_series = np.full((num_series), (time_len-input_size) // stride_size)
if train: windows_per_series -= (data_start+stride_size-1) // stride_size
total_windows = np.sum(windows_per_series)
x_input = np.zeros((total_windows, window_size, 1 + num_covariates + 1), dtype='float32')
label = np.zeros((total_windows, window_size), dtype='float32')
v_input = np.zeros((total_windows, 2), dtype='float32')
count = 0
if not train:
covariates = covariates[-time_len:]
for series in trange(num_series):
cov_age = stats.zscore(np.arange(total_time-data_start[series])) # shape:(series_len,)
if train:
covariates[data_start[series]:time_len, 0] = cov_age[:time_len-data_start[series]]
else:
covariates[:, 0] = cov_age[-time_len:]
for i in range(windows_per_series[series]):
if train:
window_start = stride_size*i+data_start[series]
else:
window_start = stride_size*i
window_end = window_start+window_size
'''
print("x: ", x_input[count, 1:, 0].shape)
print("window start: ", window_start)
print("window end: ", window_end)
print("data: ", data.shape)
print("d: ", data[window_start:window_end-1, series].shape)
'''
x_input[count, 1:, 0] = data[window_start:window_end-1, series]
x_input[count, :, 1:1+num_covariates] = covariates[window_start:window_end, :]
x_input[count, :, -1] = series
label[count, :] = data[window_start:window_end, series]
nonzero_sum = (x_input[count, 1:input_size, 0]!=0).sum()
if nonzero_sum == 0:
v_input[count, 0] = 0
else:
v_input[count, 0] = np.true_divide(x_input[count, 1:input_size, 0].sum(),nonzero_sum)+1
x_input[count, :, 0] = x_input[count, :, 0]/v_input[count, 0]
if train:
label[count, :] = label[count, :]/v_input[count, 0]
count += 1
prefix = os.path.join(save_path, 'train_' if train else 'test_')
np.save(prefix+'data_'+save_name, x_input)
np.save(prefix+'v_'+save_name, v_input)
np.save(prefix+'label_'+save_name, label)
def gen_covariates(times, num_covariates):
"""Get covariates"""
covariates = np.zeros((times.shape[0], num_covariates))
for i, input_time in enumerate(times):
covariates[i, 1] = input_time.weekday()
covariates[i, 2] = input_time.hour
covariates[i, 3] = input_time.month
for i in range(1,num_covariates):
covariates[:,i] = stats.zscore(covariates[:,i])
return covariates[:, :num_covariates]
def visualize(data, week_start):
x = np.arange(window_size)
f = plt.figure()
plt.plot(x, data[week_start:week_start+window_size], color='b')
f.savefig("visual.png")
plt.close()
if __name__ == '__main__':
global save_path
csv_path = 'data/LD2011_2014.txt'
save_name = 'elect'
window_size = 192
stride_size = 24
num_covariates = 4
train_start = '2011-01-01 00:00:00'
train_end = '2014-08-31 23:00:00'
test_start = '2014-08-25 00:00:00' #need additional 7 days as given info
test_end = '2014-09-07 23:00:00'
pred_days = 7
given_days = 7
save_path = os.path.join('data', save_name)
data_frame = pd.read_csv(csv_path, sep=";", index_col=0, parse_dates=True, decimal=',')
data_frame = data_frame.resample('1H',label = 'left',closed = 'right').sum()[train_start:test_end]
data_frame.fillna(0, inplace=True)
covariates = gen_covariates(data_frame[train_start:test_end].index, num_covariates)
train_data = data_frame[train_start:train_end].values # shape: [seq_length, user_num]
test_data = data_frame[test_start:test_end].values
data_start = (train_data!=0).argmax(axis=0) #find first nonzero value in each time series
total_time = data_frame.shape[0] #32304
num_series = data_frame.shape[1] #370
prep_data(train_data, covariates, data_start)
prep_data(test_data, covariates, data_start, train=False)
@@ -0,0 +1,107 @@
from numpy.lib.npyio import save
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import os
from tqdm import trange
import zipfile
def load_data(filedir):
data_frame = pd.read_csv(filedir, header=0, parse_dates=True) #names=['app_name', 'zone', 'time', 'value']
data_frame = data_frame.drop(data_frame.columns[0], axis=1)
grouped_data = list(data_frame.groupby(["app_name", "zone"]))
# covariates = gen_covariates(data_frame.index, 3)
all_data = []
for i in range(len(grouped_data)):
single_df = grouped_data[i][1].drop(labels=['app_name', 'zone'], axis=1).sort_values(by="time", ascending=True)
times = pd.to_datetime(single_df.time)
single_df['weekday'] = times.dt.dayofweek / 6
single_df['hour'] = times.dt.hour / 23
single_df['month'] = times.dt.month / 12
temp_data = single_df.values[:, 1:]
if (temp_data[:, 0] == 0).sum() / len(temp_data) > 0.2:
continue
all_data.append(temp_data)
return all_data
def visualize(data, index, save_dir):
os.makedirs(save_dir, exist_ok=True)
for i in range(index):
x = np.arange(len(data[i]))
f = plt.figure()
plt.plot(x, data[i][:, 0])
f.savefig(os.path.join(save_dir, "visual_{}.png".format(i)))
plt.close()
def split_seq(sequences, seq_length, slide_step, predict_length, save_dir):
"""Divide the training sequence into windows"""
train_data = []
test_data = []
for seq_id in trange(len(sequences)):
split_start = 0
single_seq = sequences[seq_id][:, 0]
single_covariate = sequences[seq_id][:, 1:]
windows = (len(single_seq)-seq_length+slide_step) // slide_step
count = 0
train_count = int(0.97 * windows)
while len(single_seq[split_start:]) > (seq_length + predict_length):
seq_data = single_seq[split_start:(split_start+seq_length+predict_length-1)]
single_data = np.zeros((seq_length+predict_length-1, 5))
single_data[:, 0] = seq_data.copy()
single_data[:, 1:4] = single_covariate[split_start:(split_start+seq_length+predict_length-1)]
single_data[:, -1] = seq_id
count += 1
if count < train_count:
train_data.append(single_data)
else:
test_data.append(single_data)
split_start += slide_step
os.makedirs(save_dir, exist_ok=True)
train_data = np.array(train_data, dtype=np.float32)
train_data, v = normalize(train_data, seq_length)
save(train_data, v, save_dir + 'train')
test_data = np.array(test_data, dtype=np.float32)
test_data, v = normalize(test_data, seq_length)
save(test_data, v, save_dir + 'test')
def normalize(inputs, seq_length):
base_seq = inputs[:, :(seq_length-1), 0]
nonzeros = (base_seq > 0).sum(1)
v = base_seq.sum(1) / nonzeros
v[v == 0] = 1
inputs[:, :, 0] = inputs[:, :, 0] / v[:, None]
return inputs, v
def save(data, v, save_dir):
np.save(save_dir+'_data_flow.npy', data)
np.save(save_dir+'_v_flow.npy', v)
def dezip(filedir):
zip_file = zipfile.ZipFile(filedir)
zip_list = zip_file.namelist()
parent_dir = filedir.split('/')[0]
for f in zip_list:
zip_file.extract(f, parent_dir)
zip_file.close()
if __name__ == '__main__':
zip_dir = 'data/app_zone_rpc_hour_encrypted.zip'
dezip(zip_dir)
data_dir = 'data/app_zone_rpc_hour_encrypted.csv'
data = load_data(data_dir)
split_seq(data, 192, 24, 24, 'data/flow/')
@@ -0,0 +1,89 @@
import numpy as np
from numpy.core.defchararray import split
import pandas as pd
from datetime import datetime
from scipy import stats
import os
def load_data(datadir):
df = pd.read_csv(datadir)
data = (df.values).transpose(1, 0)
return data
def get_covariates(data_len, start_day):
"""Get covariates"""
start_timestamp = datetime.timestamp(datetime.strptime(start_day, '%Y-%m-%d %H:%M:%S'))
timestamps = np.arange(data_len) * 3600 + start_timestamp
timestamps = [datetime.fromtimestamp(i) for i in timestamps]
weekdays = stats.zscore(np.array([i.weekday() for i in timestamps]))
hours = stats.zscore(np.array([i.hour for i in timestamps]))
months = stats.zscore(np.array([i.month for i in timestamps]))
covariates = np.stack([weekdays, hours, months], axis=1)
return covariates
def split_seq(sequences, covariates, seq_length, slide_step, predict_length, save_dir):
"""Divide the training sequence into windows"""
data_length = len(sequences[0])
windows = (data_length-seq_length+slide_step) // slide_step
train_windows = int(0.97 * windows)
test_windows = windows - train_windows
train_data = np.zeros((train_windows*len(sequences), seq_length+predict_length-1, 5), dtype=np.float32)
test_data = np.zeros((test_windows*len(sequences), seq_length+predict_length-1, 5), dtype=np.float32)
count = 0
split_start = 0
seq_ids = np.arange(len(sequences))[:, None]
end = split_start + seq_length + predict_length - 1
while end <= data_length:
if count < train_windows:
train_data[count*len(sequences):(count+1)*len(sequences), :, 0] = sequences[:, split_start:end]
train_data[count*len(sequences):(count+1)*len(sequences), :, 1:4] = covariates[split_start:end, :]
train_data[count*len(sequences):(count+1)*len(sequences), :, -1] = seq_ids
else:
test_data[(count-train_windows)*len(sequences):(count-train_windows+1)*len(sequences), :, 0] = sequences[:, split_start:end]
test_data[(count-train_windows)*len(sequences):(count-train_windows+1)*len(sequences), :, 1:4] = covariates[split_start:end, :]
test_data[(count-train_windows)*len(sequences):(count-train_windows+1)*len(sequences), :, -1] = seq_ids
count += 1
split_start += slide_step
end = split_start + seq_length + predict_length - 1
os.makedirs(save_dir, exist_ok=True)
train_data, v = normalize(train_data, seq_length)
save(train_data, v, save_dir + 'train')
test_data, v = normalize(test_data, seq_length)
save(test_data, v, save_dir + 'test')
def normalize(inputs, seq_length):
base_seq = inputs[:, :seq_length, 0]
nonzeros = (base_seq > 0).sum(1)
inputs = inputs[nonzeros > 0]
base_seq = inputs[:, :seq_length, 0]
nonzeros = nonzeros[nonzeros > 0]
v = base_seq.sum(1) / nonzeros
v[v == 0] = 1
inputs[:, :, 0] = inputs[:, :, 0] / v[:, None]
return inputs, v
def save(data, v, save_dir):
np.save(save_dir+'_data_wind.npy', data)
np.save(save_dir+'_v_wind.npy', v)
if __name__ == '__main__':
datadir = 'data/EMHIRESPV_TSh_CF_Country_19862015.csv'
all_data = load_data(datadir)
covariates = get_covariates(len(all_data[0]), '1986-01-01 00:00:00')
split_seq(all_data, covariates, 192, 24, 24, 'data/wind/')
@@ -0,0 +1,408 @@
from torch.functional import align_tensors
import torch.nn as nn
from torch.nn.modules.linear import Linear
from .SubLayers import MultiHeadAttention, PositionwiseFeedForward
import torch
from .embed import DataEmbedding, CustomEmbedding
import math
def get_mask(input_size, window_size, inner_size, device):
"""Get the attention mask of PAM-Naive"""
# Get the size of all layers
all_size = []
all_size.append(input_size)
for i in range(len(window_size)):
layer_size = math.floor(all_size[i] / window_size[i])
all_size.append(layer_size)
seq_length = sum(all_size)
mask = torch.zeros(seq_length, seq_length, device=device)
# get intra-scale mask
inner_window = inner_size // 2
for layer_idx in range(len(all_size)):
start = sum(all_size[:layer_idx])
for i in range(start, start + all_size[layer_idx]):
left_side = max(i - inner_window, start)
right_side = min(i + inner_window + 1, start + all_size[layer_idx])
mask[i, left_side:right_side] = 1
# get inter-scale mask
for layer_idx in range(1, len(all_size)):
start = sum(all_size[:layer_idx])
for i in range(start, start + all_size[layer_idx]):
left_side = (start - all_size[layer_idx - 1]) + (i - start) * window_size[layer_idx - 1]
if i == ( start + all_size[layer_idx] - 1):
right_side = start
else:
right_side = (start - all_size[layer_idx - 1]) + (i - start + 1) * window_size[layer_idx - 1]
mask[i, left_side:right_side] = 1
mask[left_side:right_side, i] = 1
mask = (1 - mask).bool()
return mask, all_size
def refer_points(all_sizes, window_size, device):
"""Gather features from PAM's pyramid sequences"""
input_size = all_sizes[0]
indexes = torch.zeros(input_size, len(all_sizes), device=device)
for i in range(input_size):
indexes[i][0] = i
former_index = i
for j in range(1, len(all_sizes)):
start = sum(all_sizes[:j])
inner_layer_idx = former_index - (start - all_sizes[j - 1])
former_index = start + min(inner_layer_idx // window_size[j - 1], all_sizes[j] - 1)
indexes[i][j] = former_index
indexes = indexes.unsqueeze(0).unsqueeze(3)
return indexes.long()
def get_subsequent_mask(input_size, window_size, predict_step, truncate):
"""Get causal attention mask for decoder."""
if truncate:
mask = torch.zeros(predict_step, input_size + predict_step)
for i in range(predict_step):
mask[i][:input_size+i+1] = 1
mask = (1 - mask).bool().unsqueeze(0)
else:
all_size = []
all_size.append(input_size)
for i in range(len(window_size)):
layer_size = math.floor(all_size[i] / window_size[i])
all_size.append(layer_size)
all_size = sum(all_size)
mask = torch.zeros(predict_step, all_size + predict_step)
for i in range(predict_step):
mask[i][:all_size+i+1] = 1
mask = (1 - mask).bool().unsqueeze(0)
return mask
def get_q_k(input_size, window_size, stride, device):
"""
Get the index of the key that a given query needs to attend to.
"""
second_length = input_size // stride
second_last = input_size - (second_length - 1) * stride
third_start = input_size + second_length
third_length = second_length // stride
third_last = second_length - (third_length - 1) * stride
max_attn = max(second_last, third_last)
fourth_start = third_start + third_length
fourth_length = third_length // stride
full_length = fourth_start + fourth_length
fourth_last = third_length - (fourth_length - 1) * stride
max_attn = max(third_last, fourth_last)
max_attn += window_size + 1
mask = torch.zeros(full_length, max_attn, dtype=torch.int32, device=device) - 1
for i in range(input_size):
mask[i, 0:window_size] = i + torch.arange(window_size) - window_size // 2
mask[i, mask[i] > input_size - 1] = -1
mask[i, -1] = i // stride + input_size
mask[i][mask[i] > third_start - 1] = third_start - 1
for i in range(second_length):
mask[input_size+i, 0:window_size] = input_size + i + torch.arange(window_size) - window_size // 2
mask[input_size+i, mask[input_size+i] < input_size] = -1
mask[input_size+i, mask[input_size+i] > third_start - 1] = -1
if i < second_length - 1:
mask[input_size+i, window_size:(window_size+stride)] = torch.arange(stride) + i * stride
else:
mask[input_size+i, window_size:(window_size+second_last)] = torch.arange(second_last) + i * stride
mask[input_size+i, -1] = i // stride + third_start
mask[input_size+i, mask[input_size+i] > fourth_start - 1] = fourth_start - 1
for i in range(third_length):
mask[third_start+i, 0:window_size] = third_start + i + torch.arange(window_size) - window_size // 2
mask[third_start+i, mask[third_start+i] < third_start] = -1
mask[third_start+i, mask[third_start+i] > fourth_start - 1] = -1
if i < third_length - 1:
mask[third_start+i, window_size:(window_size+stride)] = input_size + torch.arange(stride) + i * stride
else:
mask[third_start+i, window_size:(window_size+third_last)] = input_size + torch.arange(third_last) + i * stride
mask[third_start+i, -1] = i // stride + fourth_start
mask[third_start+i, mask[third_start+i] > full_length - 1] = full_length - 1
for i in range(fourth_length):
mask[fourth_start+i, 0:window_size] = fourth_start + i + torch.arange(window_size) - window_size // 2
mask[fourth_start+i, mask[fourth_start+i] < fourth_start] = -1
mask[fourth_start+i, mask[fourth_start+i] > full_length - 1] = -1
if i < fourth_length - 1:
mask[fourth_start+i, window_size:(window_size+stride)] = third_start + torch.arange(stride) + i * stride
else:
mask[fourth_start+i, window_size:(window_size+fourth_last)] = third_start + torch.arange(fourth_last) + i * stride
return mask
def get_k_q(q_k_mask):
"""
Get the index of the query that can attend to the given key.
"""
k_q_mask = q_k_mask.clone()
for i in range(len(q_k_mask)):
for j in range(len(q_k_mask[0])):
if q_k_mask[i, j] >= 0:
k_q_mask[i, j] = torch.where(q_k_mask[q_k_mask[i, j]] ==i )[0]
return k_q_mask
class EncoderLayer(nn.Module):
""" Compose with two layers """
def __init__(self, d_model, d_inner, n_head, d_k, d_v, dropout=0.1, normalize_before=True, use_tvm=False, q_k_mask=None, k_q_mask=None):
super(EncoderLayer, self).__init__()
self.use_tvm = use_tvm
if use_tvm:
from .PAM_TVM import PyramidalAttention
self.slf_attn = PyramidalAttention(n_head, d_model, d_k, d_v, dropout=dropout, normalize_before=normalize_before, q_k_mask=q_k_mask, k_q_mask=k_q_mask)
else:
self.slf_attn = MultiHeadAttention(n_head, d_model, d_k, d_v, dropout=dropout, normalize_before=normalize_before)
self.pos_ffn = PositionwiseFeedForward(
d_model, d_inner, dropout=dropout, normalize_before=normalize_before)
def forward(self, enc_input, slf_attn_mask=None):
if self.use_tvm:
enc_output = self.slf_attn(enc_input)
enc_slf_attn = None
else:
enc_output, enc_slf_attn = self.slf_attn(enc_input, enc_input, enc_input, mask=slf_attn_mask)
enc_output = self.pos_ffn(enc_output)
return enc_output, enc_slf_attn
class DecoderLayer(nn.Module):
""" Compose with two layers """
def __init__(self, d_model, d_inner, n_head, d_k, d_v, dropout=0.1, normalize_before=True):
super(DecoderLayer, self).__init__()
self.slf_attn = MultiHeadAttention(
n_head, d_model, d_k, d_v, dropout=dropout, normalize_before=normalize_before)
self.pos_ffn = PositionwiseFeedForward(
d_model, d_inner, dropout=dropout, normalize_before=normalize_before)
def forward(self, Q, K, V, slf_attn_mask=None):
enc_output, enc_slf_attn = self.slf_attn(
Q, K, V, mask=slf_attn_mask)
enc_output = self.pos_ffn(enc_output)
return enc_output, enc_slf_attn
class ConvLayer(nn.Module):
def __init__(self, c_in, window_size):
super(ConvLayer, self).__init__()
self.downConv = nn.Conv1d(in_channels=c_in,
out_channels=c_in,
kernel_size=window_size,
stride=window_size)
self.norm = nn.BatchNorm1d(c_in)
self.activation = nn.ELU()
def forward(self, x):
x = self.downConv(x)
x = self.norm(x)
x = self.activation(x)
return x
class Conv_Construct(nn.Module):
"""Convolution CSCM"""
def __init__(self, d_model, window_size, d_inner):
super(Conv_Construct, self).__init__()
if not isinstance(window_size, list):
self.conv_layers = nn.ModuleList([
ConvLayer(d_model, window_size),
ConvLayer(d_model, window_size),
ConvLayer(d_model, window_size)
])
else:
self.conv_layers = nn.ModuleList([
ConvLayer(d_model, window_size[0]),
ConvLayer(d_model, window_size[1]),
ConvLayer(d_model, window_size[2])
])
self.norm = nn.LayerNorm(d_model)
def forward(self, enc_input):
all_inputs = []
enc_input = enc_input.permute(0, 2, 1)
all_inputs.append(enc_input)
for i in range(len(self.conv_layers)):
enc_input = self.conv_layers[i](enc_input)
all_inputs.append(enc_input)
all_inputs = torch.cat(all_inputs, dim=2).transpose(1, 2)
all_inputs = self.norm(all_inputs)
return all_inputs
class Bottleneck_Construct(nn.Module):
"""Bottleneck convolution CSCM"""
def __init__(self, d_model, window_size, d_inner):
super(Bottleneck_Construct, self).__init__()
if not isinstance(window_size, list):
self.conv_layers = nn.ModuleList([
ConvLayer(d_inner, window_size),
ConvLayer(d_inner, window_size),
ConvLayer(d_inner, window_size)
])
else:
self.conv_layers = []
for i in range(len(window_size)):
self.conv_layers.append(ConvLayer(d_inner, window_size[i]))
self.conv_layers = nn.ModuleList(self.conv_layers)
self.up = Linear(d_inner, d_model)
self.down = Linear(d_model, d_inner)
self.norm = nn.LayerNorm(d_model)
def forward(self, enc_input):
temp_input = self.down(enc_input).permute(0, 2, 1)
all_inputs = []
for i in range(len(self.conv_layers)):
temp_input = self.conv_layers[i](temp_input)
all_inputs.append(temp_input)
all_inputs = torch.cat(all_inputs, dim=2).transpose(1, 2)
all_inputs = self.up(all_inputs)
all_inputs = torch.cat([enc_input, all_inputs], dim=1)
all_inputs = self.norm(all_inputs)
return all_inputs
class MaxPooling_Construct(nn.Module):
"""Max pooling CSCM"""
def __init__(self, d_model, window_size, d_inner):
super(MaxPooling_Construct, self).__init__()
if not isinstance(window_size, list):
self.pooling_layers = nn.ModuleList([
nn.MaxPool1d(kernel_size=window_size),
nn.MaxPool1d(kernel_size=window_size),
nn.MaxPool1d(kernel_size=window_size)
])
else:
self.pooling_layers = nn.ModuleList([
nn.MaxPool1d(kernel_size=window_size[0]),
nn.MaxPool1d(kernel_size=window_size[1]),
nn.MaxPool1d(kernel_size=window_size[2])
])
self.norm = nn.LayerNorm(d_model)
def forward(self, enc_input):
all_inputs = []
enc_input = enc_input.transpose(1, 2).contiguous()
all_inputs.append(enc_input)
for layer in self.pooling_layers:
enc_input = layer(enc_input)
all_inputs.append(enc_input)
all_inputs = torch.cat(all_inputs, dim=2).transpose(1, 2)
all_inputs = self.norm(all_inputs)
return all_inputs
class AvgPooling_Construct(nn.Module):
"""Average pooling CSCM"""
def __init__(self, d_model, window_size, d_inner):
super(AvgPooling_Construct, self).__init__()
if not isinstance(window_size, list):
self.pooling_layers = nn.ModuleList([
nn.AvgPool1d(kernel_size=window_size),
nn.AvgPool1d(kernel_size=window_size),
nn.AvgPool1d(kernel_size=window_size)
])
else:
self.pooling_layers = nn.ModuleList([
nn.AvgPool1d(kernel_size=window_size[0]),
nn.AvgPool1d(kernel_size=window_size[1]),
nn.AvgPool1d(kernel_size=window_size[2])
])
self.norm = nn.LayerNorm(d_model)
def forward(self, enc_input):
all_inputs = []
enc_input = enc_input.transpose(1, 2).contiguous()
all_inputs.append(enc_input)
for layer in self.pooling_layers:
enc_input = layer(enc_input)
all_inputs.append(enc_input)
all_inputs = torch.cat(all_inputs, dim=2).transpose(1, 2)
all_inputs = self.norm(all_inputs)
return all_inputs
class Predictor(nn.Module):
def __init__(self, dim, num_types):
super().__init__()
self.linear = nn.Linear(dim, num_types, bias=False)
nn.init.xavier_normal_(self.linear.weight)
def forward(self, data):
out = self.linear(data)
out = out
return out
class Decoder(nn.Module):
""" A encoder model with self attention mechanism. """
def __init__(self, opt, mask):
super().__init__()
self.model_type = opt.model
self.mask = mask
self.layers = nn.ModuleList([
DecoderLayer(opt.d_model, opt.d_inner_hid, opt.n_head, opt.d_k, opt.d_v, dropout=opt.dropout, \
normalize_before=False),
DecoderLayer(opt.d_model, opt.d_inner_hid, opt.n_head, opt.d_k, opt.d_v, dropout=opt.dropout, \
normalize_before=False)
])
if opt.embed_type == 'CustomEmbedding':
self.dec_embedding = CustomEmbedding(opt.enc_in, opt.d_model, opt.covariate_size, opt.seq_num, opt.dropout)
else:
self.dec_embedding = DataEmbedding(opt.enc_in, opt.d_model, opt.dropout)
def forward(self, x_dec, x_mark_dec, refer):
dec_enc = self.dec_embedding(x_dec, x_mark_dec)
dec_enc, _ = self.layers[0](dec_enc, refer, refer)
refer_enc = torch.cat([refer, dec_enc], dim=1)
mask = self.mask.repeat(len(dec_enc), 1, 1).to(dec_enc.device)
dec_enc, _ = self.layers[1](dec_enc, refer_enc, refer_enc, slf_attn_mask=mask)
return dec_enc
@@ -0,0 +1,25 @@
import torch
import torch.nn as nn
import torch.nn.functional as F
class ScaledDotProductAttention(nn.Module):
""" Scaled Dot-Product Attention """
def __init__(self, temperature, attn_dropout=0.2):
super().__init__()
self.temperature = temperature
self.dropout = nn.Dropout(attn_dropout)
def forward(self, q, k, v, mask=None):
attn = torch.matmul(q / self.temperature, k.transpose(2, 3))
if mask is not None:
attn = attn.masked_fill(mask, -1e9)
attn = self.dropout(F.softmax(attn, dim=-1))
output = torch.matmul(attn, v)
return output, attn
@@ -0,0 +1,65 @@
import torch.nn as nn
import torch.nn.functional as F
import math
from .hierarchical_mm_tvm import graph_mm as graph_mm_tvm
class PyramidalAttention(nn.Module):
def __init__(self, n_head, d_model, d_k, d_v, dropout, normalize_before, q_k_mask, k_q_mask):
super(PyramidalAttention, self).__init__()
self.normalize_before = normalize_before
self.n_head = n_head
self.d_k = d_k
self.w_qs = nn.Linear(d_model, n_head * d_k, bias=False)
self.w_ks = nn.Linear(d_model, n_head * d_k, bias=False)
self.w_vs = nn.Linear(d_model, n_head * d_k, bias=False)
nn.init.xavier_uniform_(self.w_qs.weight)
nn.init.xavier_uniform_(self.w_ks.weight)
nn.init.xavier_uniform_(self.w_vs.weight)
self.fc = nn.Linear(d_k * n_head, d_model)
nn.init.xavier_uniform_(self.fc.weight)
self.layer_norm = nn.LayerNorm(d_model, eps=1e-6)
self.dropout_attn = nn.Dropout(dropout)
self.dropout_fc = nn.Dropout(dropout)
self.q_k_mask = q_k_mask
self.k_q_mask = k_q_mask
def forward(self, hidden_states):
residual = hidden_states
hidden_states = hidden_states
bsz, seq_len, _ = hidden_states.size()
q = hidden_states
if self.normalize_before:
q = self.layer_norm(q)
q = self.w_qs(q)
k = self.w_ks(hidden_states)
v = self.w_vs(hidden_states)
q /= math.sqrt(self.d_k)
q = q.view(bsz, seq_len, self.n_head, self.d_k)
k = k.view(bsz, seq_len, self.n_head, self.d_k)
q = q.float().contiguous()
k = k.float().contiguous()
# attn_weights.size(): (batch_size, L, num_heads, 11)
attn_weights = graph_mm_tvm(q, k, self.q_k_mask, self.k_q_mask, False, 0)
attn_weights = self.dropout_attn(F.softmax(attn_weights, dim=-1))
v = v.view(bsz, seq_len, self.n_head, self.d_k)
v = v.float().contiguous()
# is_t1_diagonaled=True
attn = graph_mm_tvm(attn_weights, v, self.q_k_mask, self.k_q_mask, True, 0)
attn = attn.reshape(bsz, seq_len, self.n_head * self.d_k).contiguous()
context = self.dropout_fc(self.fc(attn))
context += residual
if not self.normalize_before:
context = self.layer_norm(context)
return context

Some files were not shown because too many files have changed in this diff Show More