PASTA_Example_Demo.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f96d0502-d0bc-44ef-ad64-78663ed242f2",
   "metadata": {},
   "outputs": [],
   "source": [
    "import random, os, itertools, sys, warnings, logging, json, multiprocessing, time, glob\n",
    "\n",
    "# FOR GPU SELECTION\n",
    "os.environ[\"CUDA_DEVICE_ORDER\"] = \"PCI_BUS_ID\"\n",
    "os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"3\" # -1 = not using gpu\n",
    "os.environ[\"TF_FORCE_GPU_ALLOW_GROWTH\"] = \"true\"\n",
    "\n",
    "logging.disable(logging.WARNING)\n",
    "warnings.filterwarnings('ignore')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e89ed58c-8f50-4086-9d4f-6f076bc4e8e3",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import pandas as pd\n",
    "import tensorflow as tf\n",
    "import ngboost as ngb\n",
    "\n",
    "from utils.evaluator import compute_anomaly_scores, compute_new_metrics\n",
    "from utils.experiment import core_seed, running_seeds, save_model_configs, save_results\n",
    "from utils.experiment import data_loaders, seq_lengths, strides, n_epochs, batch_size, clear_memory\n",
    "\n",
    "from tensorflow import keras\n",
    "from tensorflow.keras import layers, optimizers, callbacks, Input, Model\n",
    "from tensorflow.keras.losses import MSE, MAE, logcosh \n",
    "\n",
    "from PASTA.graph_builder import build_graph\n",
    "from PASTA.search_space import SearchSpace\n",
    "\n",
    "from tqdm.notebook import tqdm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6183ce26-90ca-4db0-9347-948dd96e9f91",
   "metadata": {},
   "outputs": [],
   "source": [
    "data_name = 'psm' # ['tods', 'asd', 'psm'] for SWaT, please request the dataset first\n",
    "budget = 100\n",
    "z_dim = 32 # default value"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "04c8f1f0-e9d3-47d6-a91d-36a620a08d34",
   "metadata": {},
   "outputs": [],
   "source": [
    "# THESE LINES ARE FOR REPRODUCIBILITY\n",
    "random.seed(core_seed)\n",
    "np.random.seed(core_seed)\n",
    "tf.random.set_seed(core_seed)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "21e0ce9b-71a6-42ad-aad9-ebcfc7da801b",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Data Loaders \n",
    "seq_length, stride = seq_lengths[data_name], strides[data_name]\n",
    "data = data_loaders[data_name](seq_length=seq_length, stride=stride)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1b527ce9-622f-4f2f-a0bc-e097c101b10c",
   "metadata": {},
   "outputs": [],
   "source": [
    "def valid_topK(idx): # arch_matrices, graph_configs, arch_connections\n",
    "    \n",
    "    scores = []     \n",
    "    main_configs, layer_configs = graph_configs[idx]\n",
    "    reverse_output = main_configs[\"reverse_output\"]\n",
    "\n",
    "    selected_idx = list(range(len(data['x_train']))) # dataset index in a benchmark\n",
    "        \n",
    "    for i in tqdm(selected_idx):\n",
    "        x_train, x_valid, x_test = data['x_train'][i], data['x_valid'][i], data['x_test'][i]\n",
    "        y_valid, y_test = data['y_valid'][i], data['y_test'][i]\n",
    "        y_segment_valid, y_segment_test = data['y_segment_valid'][i], data['y_segment_test'][i]\n",
    "\n",
    "        start_time = time.time()\n",
    "        model = build_graph(x_train.shape, graph_configs[idx], arch_connections[idx])\n",
    "        build_time = time.time() - start_time\n",
    "        print(f'graph built: {build_time:.2f} seconds.')\n",
    "\n",
    "        if reverse_output:\n",
    "            x_train_reverse = np.flip(x_train, axis=1)\n",
    "            x_valid_reverse = np.flip(x_valid, axis=1)\n",
    "            x_test_reverse = np.flip(x_test, axis=1)\n",
    "\n",
    "            # Prepare the training dataset.\n",
    "            train_dataset = tf.data.Dataset.from_tensor_slices((x_train, x_train_reverse))\n",
    "            train_dataset = train_dataset.batch(batch_size).prefetch(tf.data.AUTOTUNE)\n",
    "\n",
    "            # Prepare the validation dataset.\n",
    "            valid_dataset = tf.data.Dataset.from_tensor_slices((x_valid, x_valid_reverse))\n",
    "            valid_dataset = valid_dataset.batch(batch_size).prefetch(tf.data.AUTOTUNE)\n",
    "        else:\n",
    "            # Prepare the training dataset.\n",
    "            train_dataset = tf.data.Dataset.from_tensor_slices((x_train, x_train))\n",
    "            train_dataset = train_dataset.batch(batch_size).prefetch(tf.data.AUTOTUNE)\n",
    "\n",
    "            # Prepare the validation dataset.\n",
    "            valid_dataset = tf.data.Dataset.from_tensor_slices((x_valid, x_valid))\n",
    "            valid_dataset = valid_dataset.batch(batch_size).prefetch(tf.data.AUTOTUNE)\n",
    "\n",
    "        patience = 5\n",
    "        losses = {\n",
    "            'mse': tf.keras.losses.MSE,\n",
    "            'mae': tf.keras.losses.MAE,\n",
    "            'logcosh': tf.keras.losses.logcosh\n",
    "        }\n",
    "        optimizer = optimizers.Adam()\n",
    "        loss_fn = main_configs[\"loss_fn\"]\n",
    "        es = callbacks.EarlyStopping(monitor=\"val_loss\", patience=patience, mode=\"min\", restore_best_weights=True)\n",
    "\n",
    "        start_time = time.time()\n",
    "        model.compile(loss=loss_fn, optimizer=optimizer)\n",
    "        logs = model.fit(train_dataset, validation_data=valid_dataset, epochs=n_epochs, callbacks=[es], verbose=2)\n",
    "        train_time = time.time() - start_time\n",
    "        print(f'Train Time: {train_time:.2f}s')\n",
    "\n",
    "        if reverse_output:\n",
    "            train_pred = [np.flip(rec, axis=1) for rec in model.predict(x_train)]\n",
    "            valid_pred = [np.flip(rec, axis=1) for rec in model.predict(x_valid)]\n",
    "        else:\n",
    "            train_pred = model.predict(x_train)\n",
    "            valid_pred = model.predict(x_valid)\n",
    "\n",
    "        train_pred = np.array(train_pred)\n",
    "        valid_pred = np.array(valid_pred)  \n",
    "\n",
    "        valid_rec = compute_anomaly_scores(x_valid, valid_pred, scoring=main_configs[\"scoring\"], x_val = x_train, rec_val = train_pred)\n",
    "        valid_scores = compute_new_metrics(valid_rec, y_valid, stride=stride)\n",
    "\n",
    "        scores.append(max(valid_scores['eTaF1'])) # search by best-F1\n",
    "\n",
    "    return {'perf': np.average(scores), 'idx': idx}\n",
    "\n",
    "def run_full_train(rid:int, seed:int, arch_matrices, graph_configs, arch_connections):\n",
    "    \n",
    "    gpus = tf.config.list_physical_devices('GPU')\n",
    "    if gpus:\n",
    "        # Restrict TensorFlow to only allocate 1GB of memory on the first GPU\n",
    "        try:\n",
    "            tf.config.set_logical_device_configuration(\n",
    "                gpus[0],\n",
    "                [tf.config.LogicalDeviceConfiguration(memory_limit=10000)])\n",
    "            logical_gpus = tf.config.list_logical_devices('GPU')\n",
    "            print(len(gpus), \"Physical GPUs,\", len(logical_gpus), \"Logical GPUs\")\n",
    "        except RuntimeError as e:\n",
    "            # Virtual devices must be set before GPUs have been initialized\n",
    "            print(e)    \n",
    "    \n",
    "    tf.keras.backend.clear_session()\n",
    "    \n",
    "    # setting for each run\n",
    "    random.seed(seed)\n",
    "    np.random.seed(seed)\n",
    "    tf.random.set_seed(seed)    \n",
    "    \n",
    "    main_configs, layer_configs = graph_configs\n",
    "    reverse_output = main_configs[\"reverse_output\"]   \n",
    "    \n",
    "    selected_idx = list(range(len(data['x_train']))) # dataset index in a benchmark\n",
    "        \n",
    "    for i in tqdm(selected_idx):\n",
    "        x_train, x_valid, x_test = data['x_train'][i], data['x_valid'][i], data['x_test'][i]\n",
    "        y_valid, y_test = data['y_valid'][i], data['y_test'][i]\n",
    "        y_segment_valid, y_segment_test = data['y_segment_valid'][i], data['y_segment_test'][i]\n",
    "        \n",
    "        start_time = time.time()\n",
    "        model = build_graph(x_train.shape, graph_configs, arch_connections)\n",
    "        build_time = time.time() - start_time\n",
    "        print(f'graph built: {build_time:.2f} seconds.')\n",
    "        \n",
    "        if reverse_output:\n",
    "            x_train_reverse = np.flip(x_train, axis=1)\n",
    "            x_valid_reverse = np.flip(x_valid, axis=1)\n",
    "            x_test_reverse = np.flip(x_test, axis=1)\n",
    "            \n",
    "            # Prepare the training dataset.\n",
    "            train_dataset = tf.data.Dataset.from_tensor_slices((x_train, x_train_reverse))\n",
    "            train_dataset = train_dataset.batch(batch_size).prefetch(tf.data.AUTOTUNE)\n",
    "\n",
    "            # Prepare the validation dataset.\n",
    "            valid_dataset = tf.data.Dataset.from_tensor_slices((x_valid, x_valid_reverse))\n",
    "            valid_dataset = valid_dataset.batch(batch_size).prefetch(tf.data.AUTOTUNE)\n",
    "        else:\n",
    "            # Prepare the training dataset.\n",
    "            train_dataset = tf.data.Dataset.from_tensor_slices((x_train, x_train))\n",
    "            train_dataset = train_dataset.batch(batch_size).prefetch(tf.data.AUTOTUNE)\n",
    "\n",
    "            # Prepare the validation dataset.\n",
    "            valid_dataset = tf.data.Dataset.from_tensor_slices((x_valid, x_valid))\n",
    "            valid_dataset = valid_dataset.batch(batch_size).prefetch(tf.data.AUTOTUNE)\n",
    "        \n",
    "        patience = 5\n",
    "        losses = {\n",
    "            'mse': tf.keras.losses.MSE,\n",
    "            'mae': tf.keras.losses.MAE,\n",
    "            'logcosh': tf.keras.losses.logcosh\n",
    "        }\n",
    "        optimizer = optimizers.Adam()\n",
    "        loss_fn = main_configs[\"loss_fn\"]\n",
    "        es = callbacks.EarlyStopping(monitor=\"val_loss\", patience=patience, mode=\"min\", restore_best_weights=True)\n",
    "        \n",
    "        start_time = time.time()\n",
    "        model.compile(loss=loss_fn, optimizer=optimizer)\n",
    "        logs = model.fit(train_dataset, validation_data=valid_dataset, epochs=n_epochs, callbacks=[es], verbose=2)\n",
    "        train_time = time.time() - start_time\n",
    "        print(f'Train Time: {train_time:.2f}s')\n",
    "        \n",
    "        if reverse_output:\n",
    "            train_pred = [np.flip(rec, axis=1) for rec in model.predict(x_train)]\n",
    "            valid_pred = [np.flip(rec, axis=1) for rec in model.predict(x_valid)]\n",
    "            test_pred = [np.flip(rec, axis=1) for rec in model.predict(x_test)]\n",
    "            train_errors = logs.history['loss'] # model.evaluate(x_train, x_train_reverse)\n",
    "            valid_errors = logs.history['val_loss'] # model.evaluate(x_valid, x_valid_reverse)\n",
    "            test_errors = model.evaluate(x_test, x_test_reverse, verbose=0)            \n",
    "        else:\n",
    "            train_pred = model.predict(x_train)\n",
    "            valid_pred = model.predict(x_valid)\n",
    "            test_pred = model.predict(x_test)\n",
    "            train_errors = logs.history['loss'] # model.evaluate(x_train, x_train)\n",
    "            valid_errors = logs.history['val_loss'] # model.evaluate(x_valid, x_valid)\n",
    "            test_errors = model.evaluate(x_test, x_test, verbose=0)     \n",
    "            \n",
    "        train_pred = np.array(train_pred)\n",
    "        valid_pred = np.array(valid_pred)  \n",
    "        test_pred = np.array(test_pred)\n",
    "        \n",
    "        start_time = time.time()\n",
    "        test_rec = compute_anomaly_scores(x_test, test_pred, scoring=main_configs[\"scoring\"], x_val = x_valid, rec_val = valid_pred)\n",
    "        test_scores = compute_new_metrics(test_rec, y_test, stride=stride)\n",
    "        test_time = time.time() - start_time\n",
    "        print(f'Test Time: {test_time:.2f}s, {max(test_scores[\"eTaF1\"])}') # selecting best F1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8bc96e86-154f-4d59-bab7-42f32289bf89",
   "metadata": {},
   "outputs": [],
   "source": [
    "def predict_z(X, return_dict):\n",
    "    gpus = tf.config.list_physical_devices('GPU')\n",
    "    if gpus:\n",
    "        # Restrict TensorFlow to only allocate 1GB of memory on the first GPU\n",
    "        try:\n",
    "            tf.config.set_logical_device_configuration(gpus[0], [tf.config.LogicalDeviceConfiguration(memory_limit=10000)])\n",
    "            logical_gpus = tf.config.list_logical_devices('GPU')\n",
    "        except RuntimeError as e:\n",
    "            # Virtual devices must be set before GPUs have been initialized\n",
    "            print(e)    \n",
    "    \n",
    "    tf.keras.backend.clear_session()\n",
    "    \n",
    "    model = tf.keras.models.load_model(f'results/pretrained_models/PASTA_CTAE/model_{seq_length}_{z_dim}_all/')\n",
    "    encoder = tf.keras.Model(inputs=model.inputs, outputs=model.get_layer('encoder_output').output)\n",
    "    \n",
    "    return_dict[0] = np.array(encoder.predict(X))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e1d0c675-8ebc-4e90-a6ed-448203e7357c",
   "metadata": {},
   "outputs": [],
   "source": [
    "# search configuration\n",
    "# N = evaluation pool size, (randomly select) subset of top-k, sieze of top-k architecture subset\n",
    "N, Ni, TopK = int(1e3), 20, 100"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bf3731c6-681b-4bf0-bee2-351155da32c8",
   "metadata": {},
   "outputs": [],
   "source": [
    "used_budget = 0\n",
    "start_time = time.time()\n",
    "f_names = glob.glob(f'results/architectures/{data_name.upper()}/Full/*.npy') + glob.glob(f'results/architectures/{data_name.upper()}/Reduced/*.npy')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bc6ad274-56af-487c-adab-23d83f5ef97c",
   "metadata": {},
   "outputs": [],
   "source": [
    "y_true_f1 = []\n",
    "x_setting_test = []\n",
    "x_network_test = []\n",
    "x_temp_test = []\n",
    "x_encoder_test = []\n",
    "x_decoder_test = []\n",
    "\n",
    "for f_name in tqdm(f_names):\n",
    "    test_arch = np.load(f_name, allow_pickle=True).item()\n",
    "    xs_test = test_arch['onehot']\n",
    "    xc_test = test_arch['connection']\n",
    "\n",
    "    x_setting_test.append(xs_test[0])\n",
    "    x_network_test.append(xs_test[1])\n",
    "    x_temp_test.append(xs_test[2])\n",
    "    x_encoder_test.append(xc_test[0])\n",
    "    x_decoder_test.append(xc_test[1])\n",
    "\n",
    "    y_true_f1.append(np.average(test_arch['scores']['valid']['eTaF1']))\n",
    "\n",
    "x_test = [np.array(x_setting_test), np.array(x_network_test), np.array(x_temp_test), np.array(x_encoder_test), np.array(x_decoder_test)]\n",
    "y = np.array(y_true_f1)\n",
    "\n",
    "# to enable clearing GPU memory usage after prediction\n",
    "manager = multiprocessing.Manager()\n",
    "return_dict = manager.dict()\n",
    "\n",
    "p = multiprocessing.Process(target=predict_z, args=(x_test, return_dict,))\n",
    "p.start()\n",
    "p.join()\n",
    "\n",
    "Z = return_dict.values()[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4e5c2c98-2055-47cd-8127-462e67d94611",
   "metadata": {},
   "outputs": [],
   "source": [
    "N0 = int(budget * 0.50) # use 50% of budget to initialize the performace predictor\n",
    "predictor = ngb.NGBRegressor(random_state=core_seed, verbose=0)\n",
    "train_idx = np.random.choice(len(Z), size=N0, replace=False)\n",
    "Z0, y0 = Z[train_idx], y[train_idx]\n",
    "predictor.fit(Z0, y0)\n",
    "used_budget += N0\n",
    "print(f'{used_budget}/{budget}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "63e7d471-3e8a-4ba7-9a68-855f7246932b",
   "metadata": {},
   "outputs": [],
   "source": [
    "# search on evaluation pool\n",
    "search_space = SearchSpace()\n",
    "search_space.build_search_space(N)\n",
    "arch_matrices = search_space.get_random_architectures(N, with_adj = True)\n",
    "graph_configs = search_space.get_architecture_configs(arch_matrices[\"onehot\"])\n",
    "arch_connections =  []\n",
    "for config in tqdm(graph_configs):\n",
    "    arch_connections.append(search_space.get_architecture_connections(config, seq_length))    \n",
    "xs_test = arch_matrices['onehot']\n",
    "xc_test = arch_connections\n",
    "\n",
    "x_setting_test = []\n",
    "x_network_test = []\n",
    "x_temp_test = []\n",
    "x_encoder_test = []\n",
    "x_decoder_test = []\n",
    "\n",
    "for i in range(N):\n",
    "    x_setting_test.append(xs_test[i][0])\n",
    "    x_network_test.append(xs_test[i][1])\n",
    "    x_temp_test.append(xs_test[i][2])\n",
    "    x_encoder_test.append(xc_test[i][0])\n",
    "    x_decoder_test.append(xc_test[i][1])\n",
    "\n",
    "x_test = [np.array(x_setting_test), np.array(x_network_test), np.array(x_temp_test), np.array(x_encoder_test), np.array(x_decoder_test)]\n",
    "\n",
    "manager = multiprocessing.Manager()\n",
    "return_dict = manager.dict()\n",
    "\n",
    "p = multiprocessing.Process(target=predict_z, args=(x_test, return_dict,))\n",
    "p.start()\n",
    "p.join()\n",
    "\n",
    "Z = return_dict.values()[0]\n",
    "Z_Ni = Z0\n",
    "y_Ni = y0\n",
    "selected_idx = []\n",
    "while used_budget < budget:\n",
    "    y_pred = predictor.predict(Z).ravel()\n",
    "    topK_idx = np.argpartition(y_pred, -TopK)[-TopK:]\n",
    "    topK_Ni_idx = np.random.choice(topK_idx, size=Ni, replace=False)\n",
    "\n",
    "    Z_Ni = np.concatenate([Z[topK_Ni_idx], Z_Ni], axis=0)\n",
    "    y_Ni = np.concatenate([y_pred[topK_Ni_idx], y_Ni], axis=0)\n",
    "    predictor.fit(Z_Ni, y_Ni) # P_i+1\n",
    "    Z = np.delete(Z, topK_Ni_idx, axis=0) # exclude those obsereved!\n",
    "\n",
    "    used_budget += Ni\n",
    "    print(f'{used_budget}/{budget}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7b12f0ea-c8dd-459e-8d2c-88efabd03d33",
   "metadata": {},
   "outputs": [],
   "source": [
    "K = 3 # number of predicted top-k architectures for final validation\n",
    "y_pred = predictor.predict(Z).ravel()\n",
    "topK_idx = np.argpartition(y_pred, -K)[-K:] # get predicted best arch index\n",
    "print(y_pred[topK_idx])\n",
    "search_time = time.time() - start_time\n",
    "print(f'Search Time: {search_time}')\n",
    "\n",
    "with multiprocessing.Pool(1) as pool:\n",
    "    results = pool.map(valid_topK, topK_idx)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "46242919-0649-401b-823f-051dbe6b6079",
   "metadata": {},
   "outputs": [],
   "source": [
    "best_idx = None\n",
    "best_valid = -np.inf\n",
    "for result in results:\n",
    "    if result['perf'] > best_valid:\n",
    "        best_idx = result['idx']\n",
    "        best_valid = result['perf']\n",
    "\n",
    "print(f'Best Valid: {best_valid}')\n",
    "\n",
    "# full train the best validated architecture\n",
    "for rid, seed in enumerate(running_seeds(1)):\n",
    "    p = multiprocessing.Process(target=run_full_train, args=(rid, seed, arch_matrices[\"onehot\"][best_idx], graph_configs[best_idx], arch_connections[best_idx],))\n",
    "    p.start()\n",
    "    p.join()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "automl",
   "language": "python",
   "name": "automl"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}