example_train_sinusoid.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "db4cf137",
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib notebook"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7b37faa1",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import pprint\n",
    "import copy\n",
    "import hydra\n",
    "import numpy as np\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "from hydra import initialize, initialize_config_module, initialize_config_dir, compose\n",
    "from typing import Dict, Tuple\n",
    "\n",
    "from omegaconf import OmegaConf\n",
    "from torch.utils.data.dataloader import DataLoader\n",
    "\n",
    "from research.mtm.models.mtm_model import MTM, make_plots_with_masks\n",
    "from research.mtm.tokenizers.base import Tokenizer, TokenizerManager\n",
    "from research.mtm.train import main, create_eval_logs_states"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "05956960",
   "metadata": {},
   "outputs": [],
   "source": [
    "with initialize(version_base=None, config_path=\"research/mtm\"):\n",
    "    cfg = compose(config_name=\"config.yaml\", overrides=[\"+exp_mtm=sinusoid_cont\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d0073764",
   "metadata": {},
   "outputs": [],
   "source": [
    "os.makedirs(\"sinusoid_exp\", exist_ok=True)\n",
    "os.chdir(\"sinusoid_exp\")\n",
    "print(os.getcwd())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4c631c68",
   "metadata": {},
   "outputs": [],
   "source": [
    "# train model\n",
    "main(cfg)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "971cab90",
   "metadata": {},
   "outputs": [],
   "source": [
    "# load model\n",
    "def get_mtm_model(\n",
    "    path: str,\n",
    ") -> Tuple[MTM, TokenizerManager, Dict[str, Tuple[int, int]]]:\n",
    "    def _get_dataset(dataset, traj_length):\n",
    "        return hydra.utils.call(dataset, seq_steps=traj_length)\n",
    "\n",
    "    # find checkpoints in the directory\n",
    "    steps = []\n",
    "    names = []\n",
    "    paths_ = os.listdir(path)\n",
    "    for name in [os.path.join(path, n) for n in paths_ if \"pt\" in n]:\n",
    "        step = os.path.basename(name).split(\"_\")[-1].split(\".\")[0]\n",
    "        steps.append(int(step))\n",
    "        names.append(name)\n",
    "    ckpt_path = names[np.argmax(steps)]\n",
    "\n",
    "    hydra_cfg = OmegaConf.load(os.path.join(path, \"config.yaml\"))\n",
    "    cfg = hydra.utils.instantiate(hydra_cfg.args)\n",
    "    train_dataset, val_dataset = _get_dataset(hydra_cfg.dataset, cfg.traj_length)\n",
    "    tokenizers: Dict[str, Tokenizer] = {\n",
    "        k: hydra.utils.call(v, key=k, train_dataset=train_dataset)\n",
    "        for k, v in hydra_cfg.tokenizers.items()\n",
    "    }\n",
    "    tokenizer_manager = TokenizerManager(tokenizers)\n",
    "    discrete_map: Dict[str, bool] = {}\n",
    "    for k, v in tokenizers.items():\n",
    "        discrete_map[k] = v.discrete\n",
    "    train_loader = DataLoader(\n",
    "        train_dataset,\n",
    "        # shuffle=True,\n",
    "        pin_memory=True,\n",
    "        batch_size=cfg.batch_size,\n",
    "        num_workers=cfg.n_workers,\n",
    "    )\n",
    "    train_batch = next(iter(train_loader))\n",
    "    tokenized = tokenizer_manager.encode(train_batch)\n",
    "    data_shapes = {}\n",
    "    for k, v in tokenized.items():\n",
    "        data_shapes[k] = v.shape[-2:]\n",
    "\n",
    "    model_config = hydra.utils.instantiate(hydra_cfg.model_config)\n",
    "    model = MTM(data_shapes, cfg.traj_length, model_config)\n",
    "    model.load_state_dict(torch.load(ckpt_path)[\"model\"])\n",
    "    model.eval()\n",
    "\n",
    "    # freeze the model\n",
    "    for param in model.parameters():\n",
    "        param.requires_grad = False\n",
    "\n",
    "    return model, tokenizer_manager, data_shapes, val_dataset\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e5177b20",
   "metadata": {},
   "outputs": [],
   "source": [
    "model, tokenizer_manager, data_shapes, val_dataset = get_mtm_model(\".\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b27913aa",
   "metadata": {},
   "outputs": [],
   "source": [
    "val_sampler = torch.utils.data.SequentialSampler(val_dataset)\n",
    "val_loader = DataLoader(\n",
    "    val_dataset,\n",
    "    batch_size=32,\n",
    "    num_workers=0,\n",
    "    sampler=val_sampler,\n",
    ")\n",
    "val_batch = next(iter(val_loader))\n",
    "\n",
    "# visualize the data\n",
    "L = val_batch[\"states\"].shape[1]\n",
    "for states in val_batch[\"states\"][:4]:\n",
    "    plt.plot(np.arange(L), states, \"-\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2202e67a",
   "metadata": {},
   "outputs": [],
   "source": [
    "val_batch = {\n",
    "    k: v.to(\"cpu\", non_blocking=True) for k, v in val_batch.items()\n",
    "}\n",
    "device = val_batch[\"states\"].device\n",
    "seq_len = val_batch[\"states\"].shape[1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9c613877",
   "metadata": {},
   "outputs": [],
   "source": [
    "# generate masks\n",
    "obs_mask = np.ones(seq_len)\n",
    "obs_mask[seq_len // 2 + 2 :] = 0 # mask out future observations\n",
    "obs_use_mask_list = [obs_mask]\n",
    "\n",
    "masks_list = []\n",
    "for obs_mask in obs_use_mask_list:\n",
    "    masks_list.append({\"states\": torch.from_numpy(obs_mask).to(device)})\n",
    "\n",
    "prefixs = [\"prediction\"]\n",
    "logs = make_plots_with_masks(model, val_batch, tokenizer_manager, masks_list, prefixs, batch_idxs = (0, 1))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a892983c",
   "metadata": {},
   "outputs": [],
   "source": [
    "# visualize prediction\n",
    "logs[\"prediction_eval/batch=0|0_states\"].image"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1f295dae",
   "metadata": {},
   "outputs": [],
   "source": [
    "# visualize prediction\n",
    "logs[\"prediction_eval/batch=1|0_states\"].image"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}