ATC_DDPTP_DC_submodel_trainer.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "14c9891f",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.optim as optim\n",
    "import os\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "import numpy as np\n",
    "from torch.utils.data import TensorDataset, DataLoader\n",
    "import pickle\n",
    "from model_definitions import GRU_submodel\n",
    "from utils import get_ade,get_fde\n",
    "import datetime\n",
    "def seed_worker(worker_id):\n",
    "    worker_seed = torch.initial_seed() % 2**32\n",
    "    numpy.random.seed(worker_seed)\n",
    "    random.seed(worker_seed)\n",
    "print(datetime.datetime.now())\n",
    "torch.use_deterministic_algorithms(True)\n",
    "os.environ[\"CUBLAS_WORKSPACE_CONFIG\"] = \":4096:8\"\n",
    "torch.manual_seed(0)\n",
    "torch.autograd.set_detect_anomaly(True)\n",
    "device = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dd7cdad9",
   "metadata": {},
   "outputs": [],
   "source": [
    "epoch_num = 3000\n",
    "learning_rate = 0.001\n",
    "embed_dim = 128\n",
    "hidden_dim = 128\n",
    "dropout_rate = 0.5\n",
    "pred_len = 20\n",
    "batch_size = 512\n",
    "numrc = 4"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1d70fd92",
   "metadata": {},
   "outputs": [],
   "source": [
    "with open('datas/ATC_DDPTP_DC_sub_x.pickle', 'rb') as f:\n",
    "    x = pickle.load(f)\n",
    "with open('datas/ATC_DDPTP_DC_sub_y.pickle', 'rb') as f:\n",
    "    y = pickle.load(f)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c9bd5c6c",
   "metadata": {},
   "outputs": [],
   "source": [
    "for rc in range(numrc):\n",
    "    tensor_x = torch.Tensor(x[rc])/10/100\n",
    "    tensor_y = torch.tensor(y[rc])/10/100\n",
    "    datas = TensorDataset(tensor_x,tensor_y)\n",
    "    train = DataLoader(datas,batch_size=512,shuffle=False)\n",
    "    train_size = int(0.8 * len(datas))\n",
    "    test_size = len(datas) - train_size\n",
    "    \n",
    "    net = GRU_submodel(embed_dim,hidden_dim,dropout_rate,pred_len)\n",
    "    net.to(device)\n",
    "    lossfn = nn.MSELoss()\n",
    "    optimizer = optim.RMSprop(net.parameters(), lr=learning_rate)\n",
    "    train = DataLoader(datas,batch_size=batch_size,shuffle=False)\n",
    "    for epoch in range(epoch_num):  # loop over the dataset multiple times\n",
    "        running_loss = []\n",
    "        net.train()\n",
    "        for i, data in enumerate(train, 0):\n",
    "            inputs, labels = data[0].to(device), data[1].to(device)\n",
    "            optimizer.zero_grad()\n",
    "            output1 = net(inputs)\n",
    "            loss = lossfn(output1, labels.float())\n",
    "            loss.backward()\n",
    "            optimizer.step()\n",
    "            running_loss.append(loss.item())\n",
    "        print('[%d] loss: %.3f' %(epoch + 1, np.mean(running_loss)))\n",
    "        net.eval()\n",
    "        accuADE = []\n",
    "        if epoch % 50 == 0:\n",
    "            _, test_dataset = torch.utils.data.random_split(datas, [train_size, test_size], generator=torch.Generator().manual_seed(epoch))\n",
    "            test = DataLoader(test_dataset,batch_size=batch_size,shuffle=False)\n",
    "            for i, data in enumerate(test, 0):\n",
    "                inputs, labels = data[0].to(device), data[1].to(device)\n",
    "                output1 = net(inputs)\n",
    "                accuADE.append(get_ade(output1.cpu().detach().numpy(),labels.cpu().detach().numpy()))\n",
    "            print('[%d, %5d] ADE: %.3f' %(epoch + 1, i + 1, np.mean(accuADE)))\n",
    "\n",
    "    torch.save(net.state_dict(), \"./subnets/ATC_DDPTP_DC_submodel_\"+str(rc)+\".pth\")\n",
    "print(datetime.datetime.now())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dec77770",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python [conda env:gpu6]",
   "language": "python",
   "name": "conda-env-gpu6-py"
  },
  "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.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}