TFAS_Example_Demo.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "33c3457c-5ba5-4017-a3f8-734bd88075ba",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import argparse, os, logging, warnings, json, uuid, random\n",
    "\n",
    "import time\n",
    "import torch\n",
    "import search_space\n",
    "import searcher\n",
    "import trainer\n",
    "import numpy as np\n",
    "import experiments as exp_setup\n",
    "\n",
    "from torch import nn\n",
    "from exp.exp_long_term_forecasting import Exp_Long_Term_Forecast\n",
    "from exp.exp_imputation import Exp_Imputation\n",
    "from exp.exp_short_term_forecasting import Exp_Short_Term_Forecast\n",
    "from exp.exp_anomaly_detection import Exp_Anomaly_Detection\n",
    "from exp.exp_classification import Exp_Classification\n",
    "\n",
    "from zc_proxies.TimeZC.main_proxy import compute_zc_score\n",
    "\n",
    "from statsmodels.tsa.seasonal import STL\n",
    "from scipy.fft import fft, dct\n",
    "from pywt import dwt\n",
    "\n",
    "from data_provider.m4 import M4Meta\n",
    "from data_provider.data_factory import data_provider\n",
    "\n",
    "logging.disable(logging.WARNING)\n",
    "warnings.filterwarnings('ignore')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "12e706cf-69ca-42af-b00d-95ae26ebeae8",
   "metadata": {
    "tags": []
   },
   "source": [
    "# Search Phase"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "63b41375-c0d1-4910-97ce-ca9aa8fbdcea",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "parser = argparse.ArgumentParser(description='TimesNAS')\n",
    "# ZiCo\n",
    "parser.add_argument('--population_size', type=int, default=100, help='population size of evolution.')\n",
    "parser.add_argument('--save_dir', type=str, default=None, help='output directory')\n",
    "parser.add_argument('--maxbatch', type=int, default=2, help='N in Eq. (15)')\n",
    "# TimesNAS\n",
    "parser.add_argument('--dft_size', type=int, default=None, help='size of auxiliary task infomartion')\n",
    "parser.add_argument('--dwt_size', type=int, default=None, help='size of auxiliary frequency-domain decomposition infomartion')\n",
    "parser.add_argument('--td_size', type=int, default=None, help='size of auxiliary time-domain decomposition infomartion')\n",
    "parser.add_argument('--method_name', type=str, default='TimesNAS', help='searcher method')    \n",
    "parser.add_argument('--data_name', type=str, help='semantic data name')\n",
    "# basic config\n",
    "parser.add_argument('--task_name', type=str, required=True, default='long_term_forecast', help='task name, options:[long_term_forecast, short_term_forecast, imputation, classification, anomaly_detection]')\n",
    "parser.add_argument('--is_training', type=int, default=1, help='status')\n",
    "parser.add_argument('--model_id', type=str, default='test', help='model id')\n",
    "parser.add_argument('--model', type=str, default='Autoformer', help='model name, options: [Autoformer, Transformer, TimesNet]') \n",
    "# GPU\n",
    "parser.add_argument('--use_gpu', type=bool, default=True, help='use gpu')\n",
    "parser.add_argument('--gpu', type=int, default=1, help='gpu')\n",
    "parser.add_argument('--use_multi_gpu', action='store_true', help='use multiple gpus', default=False)\n",
    "parser.add_argument('--devices', type=str, default='0,1,2,3', help='device ids of multile gpus')\n",
    "# data loader\n",
    "parser.add_argument('--data', type=str, default='ETTm1', help='dataset type')\n",
    "parser.add_argument('--root_path', type=str, default='./data/ETT/', help='root path of the data file')\n",
    "parser.add_argument('--data_path', type=str, default='ETTh1.csv', help='data file')\n",
    "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')\n",
    "parser.add_argument('--target', type=str, default='OT', help='target feature in S or MS task')\n",
    "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')\n",
    "parser.add_argument('--checkpoints', type=str, default='./checkpoints/', help='location of model checkpoints')\n",
    "# forecasting task\n",
    "parser.add_argument('--seq_len', type=int, default=96, help='input sequence length')\n",
    "parser.add_argument('--label_len', type=int, default=48, help='start token length')\n",
    "parser.add_argument('--pred_len', type=int, default=96, help='prediction sequence length')\n",
    "parser.add_argument('--seasonal_patterns', type=str, default='Monthly', help='subset for M4')\n",
    "# inputation task\n",
    "parser.add_argument('--mask_rate', type=float, default=0.25, help='mask ratio')\n",
    "# anomaly detection task\n",
    "parser.add_argument('--anomaly_ratio', type=float, default=0.25, help='prior anomaly ratio (%)')\n",
    "# model define\n",
    "parser.add_argument('--top_k', type=int, default=5, help='for TimesBlock')\n",
    "parser.add_argument('--num_kernels', type=int, default=6, help='for Inception')\n",
    "parser.add_argument('--enc_in', type=int, default=7, help='encoder input size')\n",
    "parser.add_argument('--dec_in', type=int, default=7, help='decoder input size')\n",
    "parser.add_argument('--c_out', type=int, default=7, help='output size')\n",
    "parser.add_argument('--d_model', type=int, default=512, help='dimension of model')\n",
    "parser.add_argument('--n_heads', type=int, default=8, help='num of heads')\n",
    "parser.add_argument('--e_layers', type=int, default=2, help='num of encoder layers')\n",
    "parser.add_argument('--d_layers', type=int, default=1, help='num of decoder layers')\n",
    "parser.add_argument('--d_ff', type=int, default=2048, help='dimension of fcn')\n",
    "parser.add_argument('--moving_avg', type=int, default=25, help='window size of moving average')\n",
    "parser.add_argument('--factor', type=int, default=1, help='attn factor')\n",
    "parser.add_argument('--distil', action='store_false', help='whether to use distilling in encoder, using this argument means not using distilling', default=True)\n",
    "parser.add_argument('--dropout', type=float, default=0.1, help='dropout')\n",
    "parser.add_argument('--embed', type=str, default='timeF', help='time features encoding, options:[timeF, fixed, learned]')\n",
    "parser.add_argument('--activation', type=str, default='gelu', help='activation')\n",
    "parser.add_argument('--output_attention', action='store_true', help='whether to output attention in ecoder')\n",
    "# optimization\n",
    "parser.add_argument('--num_workers', type=int, default=40, help='data loader num workers')\n",
    "parser.add_argument('--itr', type=int, default=1, help='experiments times')\n",
    "parser.add_argument('--train_epochs', type=int, default=10, help='train epochs')\n",
    "parser.add_argument('--batch_size', type=int, default=32, help='batch size of train input data')\n",
    "parser.add_argument('--patience', type=int, default=3, help='early stopping patience')\n",
    "parser.add_argument('--learning_rate', type=float, default=0.0001, help='optimizer learning rate')\n",
    "parser.add_argument('--des', type=str, default='test', help='exp description')\n",
    "parser.add_argument('--loss', type=str, default='MSE', help='loss function')\n",
    "parser.add_argument('--lradj', type=str, default='type1', help='adjust learning rate')\n",
    "parser.add_argument('--use_amp', action='store_true', help='use automatic mixed precision training', default=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6a3ce5a5-5361-43e0-9078-49077ab5ee9a",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "task_name = 'long_term_forecast'\n",
    "method_name = f'TimesNAS'\n",
    "\n",
    "fix_seed = 2021\n",
    "random.seed(fix_seed)\n",
    "torch.manual_seed(fix_seed)\n",
    "np.random.seed(fix_seed)\n",
    "\n",
    "args = parser.parse_args(['--task_name', task_name])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6a2e92df-fd69-4100-8d9d-f97087096f4a",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "args.use_gpu = True if torch.cuda.is_available() and args.use_gpu else False\n",
    "\n",
    "gpu = args.gpu\n",
    "if gpu is not None:\n",
    "    torch.cuda.set_device('cuda:{}'.format(gpu))\n",
    "    torch.backends.cudnn.benchmark = True"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d60fec85-efca-4a6c-8996-327800b733e1",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "data_name = 'ETTh1'\n",
    "args.data_name = data_name\n",
    "args.is_training = 1\n",
    "args.model = 'TimesNAS'\n",
    "args.des = 'Exp'\n",
    "args.itr = 1\n",
    "\n",
    "args.features = 'M'\n",
    "args.seq_len = 96\n",
    "args.label_len = 48\n",
    "\n",
    "args.train_epochs = exp_setup.TRAINING_CONFIGS[args.task_name]['train_epochs']\n",
    "args.learning_rate = exp_setup.TRAINING_CONFIGS[args.task_name]['lr']\n",
    "\n",
    "Exp = Exp_Long_Term_Forecast\n",
    "\n",
    "args.data = exp_setup.DATA_CONFIGS[data_name]['data']\n",
    "args.enc_in = exp_setup.DATA_CONFIGS[data_name]['enc_in']\n",
    "args.dec_in = exp_setup.DATA_CONFIGS[data_name]['dec_in']\n",
    "args.c_out = exp_setup.DATA_CONFIGS[data_name]['c_out']\n",
    "args.root_path = exp_setup.DATA_CONFIGS[data_name]['root_path']\n",
    "args.data_path = exp_setup.DATA_CONFIGS[data_name]['data_path']\n",
    "\n",
    "pred_len = 96\n",
    "args.model_id = f'{data_name}_96_{pred_len}'\n",
    "args.pred_len = pred_len"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "10f39fdf-5d03-4713-bc7d-ce9ae0a32fb3",
   "metadata": {
    "scrolled": true,
    "tags": []
   },
   "outputs": [],
   "source": [
    "searcher.main(Exp, args)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6d1bd49c-7b8e-46e7-806d-84cf8948e033",
   "metadata": {},
   "source": [
    "# Training Phase"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b64188fb-ca85-44ad-9c63-cf452cc70a8a",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "args.batch_size = exp_setup.TRAINING_CONFIGS[args.task_name]['batch_size']\n",
    "args.train_epochs = exp_setup.TRAINING_CONFIGS[args.task_name]['train_epochs']\n",
    "args.learning_rate = exp_setup.TRAINING_CONFIGS[args.task_name]['lr']\n",
    "\n",
    "args.features = 'M'\n",
    "args.seq_len = 96\n",
    "args.label_len = 48\n",
    "\n",
    "args.data = exp_setup.DATA_CONFIGS[data_name]['data']\n",
    "args.enc_in = exp_setup.DATA_CONFIGS[data_name]['enc_in']\n",
    "args.dec_in = exp_setup.DATA_CONFIGS[data_name]['dec_in']\n",
    "args.c_out = exp_setup.DATA_CONFIGS[data_name]['c_out']\n",
    "args.root_path = exp_setup.DATA_CONFIGS[data_name]['root_path']\n",
    "args.data_path = exp_setup.DATA_CONFIGS[data_name]['data_path'] \n",
    "\n",
    "args.model_id = f'{args.data_name}_96_{args.pred_len}' \n",
    "\n",
    "Exp = Exp_Long_Term_Forecast"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "db684b98-01c4-4aea-9510-9d26008f6807",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "for ii in range(args.itr):\n",
    "    # setting record of experiments\n",
    "    arch_configs = np.load(f'zc_results/{args.task_name}/{args.data_name}.npy', allow_pickle=True).item()['best']['arch_configs']\n",
    "    args.d_model =  arch_configs['d_model']\n",
    "    args.d_ff = arch_configs['d_ff']\n",
    "    args.num_kernels = arch_configs['num_kernels']\n",
    "    args.top_k = arch_configs['top_k']\n",
    "    args.e_layers = arch_configs['e_layers']\n",
    "    args.dropout = arch_configs['dropout']\n",
    "    args.embed = arch_configs['embed']\n",
    "\n",
    "    setting = '{}_{}_{}_{}_ft{}_sl{}_ll{}_pl{}_dm{}_nh{}_el{}_dl{}_df{}_fc{}_eb{}_dt{}_{}_{}'.format(\n",
    "        args.task_name,\n",
    "        args.model_id,\n",
    "        args.model,\n",
    "        args.data,\n",
    "        args.features,\n",
    "        args.seq_len,\n",
    "        args.label_len,\n",
    "        args.pred_len,\n",
    "        args.d_model,\n",
    "        args.n_heads,\n",
    "        args.e_layers,\n",
    "        args.d_layers,\n",
    "        args.d_ff,\n",
    "        args.factor,\n",
    "        args.embed,\n",
    "        args.distil,\n",
    "        args.des, ii)            \n",
    "\n",
    "    print('Args in experiment:')\n",
    "    print(args)\n",
    "\n",
    "    exp = Exp(args)  # set experiments\n",
    "    print('>>>>>>>start training : {}>>>>>>>>>>>>>>>>>>>>>>>>>>'.format(setting))\n",
    "    _, train_stats = exp.train(setting)                \n",
    "    print('>>>>>>>testing : {}<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<'.format(setting))\n",
    "    mae, mse, rmse, mape, mspe = exp.test(setting)\n",
    "    torch.cuda.empty_cache()\n",
    "\n",
    "    arch_perf = {\n",
    "        'key': setting,\n",
    "        'arch_configs': arch_configs,\n",
    "        'train_stats': train_stats,\n",
    "        'final_mae': mae,\n",
    "        'final_mse': mse,\n",
    "        'final_rmse': rmse,\n",
    "        'final_mape': mape,\n",
    "        'final_mspe': mspe,\n",
    "    }\n",
    "    np.save(f\"arch_results/{args.task_name}/{args.method_name}_{setting}_Final.npy\", arch_perf)\n",
    "    print(f'SAVED {setting} !!!')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cc58ade8-3c22-46ba-99eb-9c4711f0770a",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "TimeZone",
   "language": "python",
   "name": "timezone"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}