{ "cells": [ { "cell_type": "code", "execution_count": 1, "id": "6570fadd-4e69-40e8-92c6-ff4acaabafaa", "metadata": {}, "outputs": [ { "name": "stderr", "output_type": "stream", "text": [ "/opt/anaconda3/envs/yt38/lib/python3.8/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n", " from .autonotebook import tqdm as notebook_tqdm\n" ] } ], "source": [ "import argparse\n", "import gc\n", "import logging\n", "import os\n", "import sys\n", "import time\n", "\n", "from collections import defaultdict\n", "\n", "import torch\n", "import torch.nn as nn\n", "import torch.optim as optim\n", "\n", "from sgan.data.loader import data_loader\n", "from sgan.losses import gan_g_loss, gan_d_loss, l2_loss\n", "from sgan.losses import displacement_error, final_displacement_error\n", "\n", "from sgan.models import TrajectoryGenerator, TrajectoryDiscriminator\n", "from sgan.utils import int_tuple, bool_flag, get_total_norm\n", "from sgan.utils import relative_to_abs, get_dset_path" ] }, { "cell_type": "code", "execution_count": 11, "id": "612c7867-94b7-4e70-8d62-7843cb28467f", "metadata": {}, "outputs": [], "source": [ "torch.backends.cudnn.benchmark = True\n", "\n", "parser = argparse.ArgumentParser()\n", "FORMAT = '[%(levelname)s: %(filename)s: %(lineno)4d]: %(message)s'\n", "logging.basicConfig(level=logging.INFO, format=FORMAT, stream=sys.stdout)\n", "logger = logging.getLogger(__name__)\n", "\n", "def make_parser():\n", " \n", " # Dataset options\n", " parser.add_argument('--dataset_name', default='zara1', type=str)\n", " parser.add_argument('--delim', default=' ')\n", " parser.add_argument('--loader_num_workers', default=4, type=int)\n", " parser.add_argument('--obs_len', default=8, type=int)\n", " parser.add_argument('--pred_len', default=8, type=int)\n", " parser.add_argument('--skip', default=1, type=int)\n", " \n", " # Optimization\n", " parser.add_argument('--batch_size', default=64, type=int)\n", " parser.add_argument('--num_iterations', default=10000, type=int)\n", " parser.add_argument('--num_epochs', default=200, type=int)\n", " \n", " # Model Options\n", " parser.add_argument('--embedding_dim', default=64, type=int)\n", " parser.add_argument('--num_layers', default=1, type=int)\n", " parser.add_argument('--dropout', default=0, type=float)\n", " parser.add_argument('--batch_norm', default=0, type=bool_flag)\n", " parser.add_argument('--mlp_dim', default=1024, type=int)\n", " \n", " # Generator Options\n", " parser.add_argument('--encoder_h_dim_g', default=64, type=int)\n", " parser.add_argument('--decoder_h_dim_g', default=128, type=int)\n", " parser.add_argument('--noise_dim', default=\"0,8\", type=int_tuple)\n", " parser.add_argument('--noise_type', default='gaussian')\n", " parser.add_argument('--noise_mix_type', default='ped')\n", " parser.add_argument('--clipping_threshold_g', default=0, type=float)\n", " parser.add_argument('--g_learning_rate', default=5e-4, type=float)\n", " parser.add_argument('--g_steps', default=1, type=int)\n", " \n", " # Pooling Options\n", " parser.add_argument('--pooling_type', default='pool_net')\n", " parser.add_argument('--pool_every_timestep', default=1, type=bool_flag)\n", " \n", " # Pool Net Option\n", " parser.add_argument('--bottleneck_dim', default=1024, type=int)\n", " \n", " # Social Pooling Options\n", " parser.add_argument('--neighborhood_size', default=2.0, type=float)\n", " parser.add_argument('--grid_size', default=8, type=int)\n", " \n", " # Discriminator Options\n", " parser.add_argument('--d_type', default='local', type=str)\n", " parser.add_argument('--encoder_h_dim_d', default=64, type=int)\n", " parser.add_argument('--d_learning_rate', default=5e-4, type=float)\n", " parser.add_argument('--d_steps', default=2, type=int)\n", " parser.add_argument('--clipping_threshold_d', default=0, type=float)\n", " \n", " # Loss Options\n", " parser.add_argument('--l2_loss_weight', default=0, type=float)\n", " parser.add_argument('--best_k', default=1, type=int)\n", " \n", " # Output\n", " parser.add_argument('--output_dir', default=os.getcwd())\n", " parser.add_argument('--print_every', default=5, type=int)\n", " parser.add_argument('--checkpoint_every', default=100, type=int)\n", " parser.add_argument('--checkpoint_name', default='checkpoint')\n", " parser.add_argument('--checkpoint_start_from', default=None)\n", " parser.add_argument('--restore_from_checkpoint', default=1, type=int)\n", " parser.add_argument('--num_samples_check', default=5000, type=int)\n", " \n", " # Misc\n", " parser.add_argument('--use_gpu', default=1, type=int)\n", " parser.add_argument('--timing', default=0, type=int)\n", " parser.add_argument('--gpu_num', default=\"0\", type=str)\n", "\n", " return parser.parse_args(args=[]) " ] }, { "cell_type": "code", "execution_count": 12, "id": "27104b1f-ae84-4b40-b704-2d47bba0888d", "metadata": {}, "outputs": [], "source": [ "def init_weights(m):\n", " classname = m.__class__.__name__\n", " if classname.find('Linear') != -1:\n", " nn.init.kaiming_normal_(m.weight)\n", "\n", "\n", "def get_dtypes(args):\n", " long_dtype = torch.LongTensor\n", " float_dtype = torch.FloatTensor\n", " if args.use_gpu == 1:\n", " long_dtype = torch.cuda.LongTensor\n", " float_dtype = torch.cuda.FloatTensor\n", " return long_dtype, float_dtype" ] }, { "cell_type": "code", "execution_count": 13, "id": "ab327578-4e55-4d08-bcba-d1de60e0293a", "metadata": {}, "outputs": [ { "data": { "text/plain": [ "Namespace(batch_norm=0, batch_size=64, best_k=1, bottleneck_dim=1024, checkpoint_every=100, checkpoint_name='checkpoint', checkpoint_start_from=None, clipping_threshold_d=0, clipping_threshold_g=0, d_learning_rate=0.0005, d_steps=2, d_type='local', dataset_name='zara1', decoder_h_dim_g=128, delim=' ', dropout=0, embedding_dim=64, encoder_h_dim_d=64, encoder_h_dim_g=64, g_learning_rate=0.0005, g_steps=1, gpu_num='0', grid_size=8, l2_loss_weight=0, loader_num_workers=4, mlp_dim=1024, neighborhood_size=2.0, noise_dim=(0, 8), noise_mix_type='ped', noise_type='gaussian', num_epochs=200, num_iterations=10000, num_layers=1, num_samples_check=5000, obs_len=8, output_dir='/data/autocon/experiments/sgan', pool_every_timestep=1, pooling_type='pool_net', pred_len=8, print_every=5, restore_from_checkpoint=1, skip=1, timing=0, use_gpu=1)" ] }, "execution_count": 13, "metadata": {}, "output_type": "execute_result" } ], "source": [ "args = make_parser()\n", "args" ] }, { "cell_type": "code", "execution_count": 15, "id": "a1fbc59f-639b-4c28-84ab-a85f2c4a8c7f", "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "[INFO: 182503375.py: 29]: Here is the generator:\n", "[INFO: 182503375.py: 30]: TrajectoryGenerator(\n", " (encoder): Encoder(\n", " (encoder): LSTM(64, 64)\n", " (spatial_embedding): Linear(in_features=2, out_features=64, bias=True)\n", " )\n", " (decoder): Decoder(\n", " (decoder): LSTM(64, 128)\n", " (pool_net): PoolHiddenNet(\n", " (spatial_embedding): Linear(in_features=2, out_features=64, bias=True)\n", " (mlp_pre_pool): Sequential(\n", " (0): Linear(in_features=192, out_features=512, bias=True)\n", " (1): ReLU()\n", " (2): Linear(in_features=512, out_features=1024, bias=True)\n", " (3): ReLU()\n", " )\n", " )\n", " (mlp): Sequential(\n", " (0): Linear(in_features=1152, out_features=1024, bias=True)\n", " (1): ReLU()\n", " (2): Linear(in_features=1024, out_features=128, bias=True)\n", " (3): ReLU()\n", " )\n", " (spatial_embedding): Linear(in_features=2, out_features=64, bias=True)\n", " (hidden2pos): Linear(in_features=128, out_features=2, bias=True)\n", " )\n", " (pool_net): PoolHiddenNet(\n", " (spatial_embedding): Linear(in_features=2, out_features=64, bias=True)\n", " (mlp_pre_pool): Sequential(\n", " (0): Linear(in_features=128, out_features=512, bias=True)\n", " (1): ReLU()\n", " (2): Linear(in_features=512, out_features=1024, bias=True)\n", " (3): ReLU()\n", " )\n", " )\n", " (mlp_decoder_context): Sequential(\n", " (0): Linear(in_features=1088, out_features=1024, bias=True)\n", " (1): ReLU()\n", " (2): Linear(in_features=1024, out_features=128, bias=True)\n", " (3): ReLU()\n", " )\n", ")\n", "[INFO: 182503375.py: 45]: Here is the discriminator:\n", "[INFO: 182503375.py: 46]: TrajectoryDiscriminator(\n", " (encoder): Encoder(\n", " (encoder): LSTM(64, 64)\n", " (spatial_embedding): Linear(in_features=2, out_features=64, bias=True)\n", " )\n", " (real_classifier): Sequential(\n", " (0): Linear(in_features=64, out_features=1024, bias=True)\n", " (1): ReLU()\n", " (2): Linear(in_features=1024, out_features=1, bias=True)\n", " (3): ReLU()\n", " )\n", ")\n" ] } ], "source": [ "os.environ[\"CUDA_VISIBLE_DEVICES\"] = args.gpu_num\n", "train_path = get_dset_path(args.dataset_name, 'train')\n", "val_path = get_dset_path(args.dataset_name, 'val')\n", "\n", "long_dtype, float_dtype = get_dtypes(args)\n", "\n", "\n", "generator = TrajectoryGenerator(\n", " obs_len=args.obs_len,\n", " pred_len=args.pred_len,\n", " embedding_dim=args.embedding_dim,\n", " encoder_h_dim=args.encoder_h_dim_g,\n", " decoder_h_dim=args.decoder_h_dim_g,\n", " mlp_dim=args.mlp_dim,\n", " num_layers=args.num_layers,\n", " noise_dim=args.noise_dim,\n", " noise_type=args.noise_type,\n", " noise_mix_type=args.noise_mix_type,\n", " pooling_type=args.pooling_type,\n", " pool_every_timestep=args.pool_every_timestep,\n", " dropout=args.dropout,\n", " bottleneck_dim=args.bottleneck_dim,\n", " neighborhood_size=args.neighborhood_size,\n", " grid_size=args.grid_size,\n", " batch_norm=args.batch_norm)\n", "\n", "generator.apply(init_weights)\n", "generator.type(float_dtype).train()\n", "logger.info('Here is the generator:')\n", "logger.info(generator)\n", "\n", "discriminator = TrajectoryDiscriminator(\n", " obs_len=args.obs_len,\n", " pred_len=args.pred_len,\n", " embedding_dim=args.embedding_dim,\n", " h_dim=args.encoder_h_dim_d,\n", " mlp_dim=args.mlp_dim,\n", " num_layers=args.num_layers,\n", " dropout=args.dropout,\n", " batch_norm=args.batch_norm,\n", " d_type=args.d_type)\n", "\n", "discriminator.apply(init_weights)\n", "discriminator.type(float_dtype).train()\n", "logger.info('Here is the discriminator:')\n", "logger.info(discriminator)\n", "\n", "g_loss_fn = gan_g_loss\n", "d_loss_fn = gan_d_loss\n", "\n", "optimizer_g = optim.Adam(generator.parameters(), lr=args.g_learning_rate)\n", "optimizer_d = optim.Adam(\n", " discriminator.parameters(), lr=args.d_learning_rate\n", ")\n", "\n", "# Maybe restore from checkpoint\n", "restore_path = None\n", "if args.checkpoint_start_from is not None:\n", " restore_path = args.checkpoint_start_from\n", "elif args.restore_from_checkpoint == 1:\n", " restore_path = os.path.join(args.output_dir,\n", " '%s_with_model.pt' % args.checkpoint_name)\n", "\n", "if restore_path is not None and os.path.isfile(restore_path):\n", " logger.info('Restoring from checkpoint {}'.format(restore_path))\n", " checkpoint = torch.load(restore_path)\n", " generator.load_state_dict(checkpoint['g_state'])\n", " discriminator.load_state_dict(checkpoint['d_state'])\n", " optimizer_g.load_state_dict(checkpoint['g_optim_state'])\n", " optimizer_d.load_state_dict(checkpoint['d_optim_state'])" ] }, { "cell_type": "code", "execution_count": null, "id": "6556a81c-25eb-4aeb-a1fa-f33bf9dd7db1", "metadata": {}, "outputs": [], "source": [] }, { "cell_type": "code", "execution_count": null, "id": "0cfa4ef1-a9b5-45c9-a50c-f96a856adc7d", "metadata": {}, "outputs": [], "source": [] }, { "cell_type": "code", "execution_count": 16, "id": "365ca282-afd2-43b5-b829-c3f900447d19", "metadata": {}, "outputs": [], "source": [ "def count_parameters(model):\n", " return sum(p.numel() for p in model.parameters() if p.requires_grad)\n", "\n", "generator_params = count_parameters(generator)\n", "discriminator_params = count_parameters(discriminator)\n", "\n", "total_params = generator_params + discriminator_params" ] }, { "cell_type": "code", "execution_count": 17, "id": "20b7b330-7657-419f-99cb-9690cd150a5b", "metadata": {}, "outputs": [ { "data": { "text/plain": [ "4008387" ] }, "execution_count": 17, "metadata": {}, "output_type": "execute_result" } ], "source": [ "total_params" ] }, { "cell_type": "code", "execution_count": null, "id": "1a70f47c-bb07-48ba-b05a-65dec2468086", "metadata": {}, "outputs": [], "source": [] } ], "metadata": { "kernelspec": { "display_name": "yt38", "language": "python", "name": "yt38" }, "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.8.0" } }, "nbformat": 4, "nbformat_minor": 5 }