Loading model weights.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "5384e773",
   "metadata": {},
   "source": [
    "## Loading pretrained model weights\n",
    "\n",
    "If the loading of parameters was successful, a message should be printed out saying `Loaded parameters from /PATH/TO/WEIGHTS`.\n",
    "\n",
    "Note that you will have to change the directories: `ROOT_DR` in `mpmodels/core/default_dirs.py` if loading the models outside of your home directory."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "c3a05184",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2\n",
    "\n",
    "import os\n",
    "from ptutils.models.utils import load_model\n",
    "import mpmodels\n",
    "from mpmodels.models.paths import PATHS\n",
    "from mpmodels.models.layers import LAYERS\n",
    "from mpmodels.core.feature_extractor import get_model_kwargs\n",
    "from brainmodel_utils.models.utils import get_base_model_name, get_model_func_from_name, get_model_path_from_name, get_model_layers_from_name"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "445ca01e",
   "metadata": {},
   "outputs": [],
   "source": [
    "def load_pretrained_model(model_name):\n",
    "    model_name, trained = get_base_model_name(model_name)\n",
    "    model_path = get_model_path_from_name(\n",
    "                model_paths_dict=PATHS, model_name=model_name\n",
    "            )\n",
    "    if trained:\n",
    "        # sometimes model functions automatically load the model ckpt, so no model_path needs to be passed in\n",
    "        # in which case, we alert the user\n",
    "        if model_path is None:\n",
    "            print(\n",
    "                \"No model path was explictly passed in to this trained model. If you meant to do this and didn't load the model checkpoint some other way (e.g. in the model func), then please cancel and retry!\"\n",
    "            )\n",
    "        else:\n",
    "            assert os.path.isfile(model_path)\n",
    "    else:\n",
    "        assert model_path is None\n",
    "    \n",
    "    model_layers = get_model_layers_from_name(\n",
    "            model_layers_dict=LAYERS, model_name=model_name\n",
    "        )\n",
    "    model = get_model_func_from_name(\n",
    "            model_func_dict=mpmodels.models.__dict__,\n",
    "            model_name=model_name,\n",
    "            model_kwargs=get_model_kwargs(),\n",
    "        )\n",
    "    model = load_model(\n",
    "        model=model, \n",
    "        trained=trained, \n",
    "        model_path=model_path,\n",
    "    )\n",
    "    return model, model_layers"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "187da6a4",
   "metadata": {},
   "source": [
    "# 1. VC-1 + Dynamics"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8867e87c",
   "metadata": {},
   "source": [
    "## 1a. VC-1 + CTRNN, dynamics module pretrained on Physion"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "bf1c09c4",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Loaded parameters from /om/weka/yanglab/anayebi/mental-sim/trained_models/VC-1+CTRNN_physion.pt\n",
      "======= Model architecture =======\n",
      " FrozenPretrainedEncoder(\n",
      "  (encoder): VC1_pretrained(\n",
      "    (vc1): VisionTransformer(\n",
      "      (patch_embed): PatchEmbed(\n",
      "        (proj): Conv2d(3, 1024, kernel_size=(16, 16), stride=(16, 16))\n",
      "        (norm): Identity()\n",
      "      )\n",
      "      (pos_drop): Dropout(p=0.0, inplace=False)\n",
      "      (norm_pre): Identity()\n",
      "      (blocks): Sequential(\n",
      "        (0): Block(\n",
      "          (norm1): LayerNorm((1024,), eps=1e-06, elementwise_affine=True)\n",
      "          (attn): Attention(\n",
      "            (qkv): Linear(in_features=1024, out_features=3072, bias=True)\n",
      "            (attn_drop): Dropout(p=0.0, inplace=False)\n",
      "            (proj): Linear(in_features=1024, out_features=1024, bias=True)\n",
      "            (proj_drop): Dropout(p=0.0, inplace=False)\n",
      "          )\n",
      "          (ls1): Identity()\n",
      "          (drop_path1): Identity()\n",
      "          (norm2): LayerNorm((1024,), eps=1e-06, elementwise_affine=True)\n",
      "          (mlp): Mlp(\n",
      "            (fc1): Linear(in_features=1024, out_features=4096, bias=True)\n",
      "            (act): GELU(approximate='none')\n",
      "            (drop1): Dropout(p=0.0, inplace=False)\n",
      "            (fc2): Linear(in_features=4096, out_features=1024, bias=True)\n",
      "            (drop2): Dropout(p=0.0, inplace=False)\n",
      "          )\n",
      "          (ls2): Identity()\n",
      "          (drop_path2): Identity()\n",
      "        )\n",
      "        (1): Block(\n",
      "          (norm1): LayerNorm((1024,), eps=1e-06, elementwise_affine=True)\n",
      "          (attn): Attention(\n",
      "            (qkv): Linear(in_features=1024, out_features=3072, bias=True)\n",
      "            (attn_drop): Dropout(p=0.0, inplace=False)\n",
      "            (proj): Linear(in_features=1024, out_features=1024, bias=True)\n",
      "            (proj_drop): Dropout(p=0.0, inplace=False)\n",
      "          )\n",
      "          (ls1): Identity()\n",
      "          (drop_path1): Identity()\n",
      "          (norm2): LayerNorm((1024,), eps=1e-06, elementwise_affine=True)\n",
      "          (mlp): Mlp(\n",
      "            (fc1): Linear(in_features=1024, out_features=4096, bias=True)\n",
      "            (act): GELU(approximate='none')\n",
      "            (drop1): Dropout(p=0.0, inplace=False)\n",
      "            (fc2): Linear(in_features=4096, out_features=1024, bias=True)\n",
      "            (drop2): Dropout(p=0.0, inplace=False)\n",
      "          )\n",
      "          (ls2): Identity()\n",
      "          (drop_path2): Identity()\n",
      "        )\n",
      "        (2): Block(\n",
      "          (norm1): LayerNorm((1024,), eps=1e-06, elementwise_affine=True)\n",
      "          (attn): Attention(\n",
      "            (qkv): Linear(in_features=1024, out_features=3072, bias=True)\n",
      "            (attn_drop): Dropout(p=0.0, inplace=False)\n",
      "            (proj): Linear(in_features=1024, out_features=1024, bias=True)\n",
      "            (proj_drop): Dropout(p=0.0, inplace=False)\n",
      "          )\n",
      "          (ls1): Identity()\n",
      "          (drop_path1): Identity()\n",
      "          (norm2): LayerNorm((1024,), eps=1e-06, elementwise_affine=True)\n",
      "          (mlp): Mlp(\n",
      "            (fc1): Linear(in_features=1024, out_features=4096, bias=True)\n",
      "            (act): GELU(approximate='none')\n",
      "            (drop1): Dropout(p=0.0, inplace=False)\n",
      "            (fc2): Linear(in_features=4096, out_features=1024, bias=True)\n",
      "            (drop2): Dropout(p=0.0, inplace=False)\n",
      "          )\n",
      "          (ls2): Identity()\n",
      "          (drop_path2): Identity()\n",
      "        )\n",
      "        (3): Block(\n",
      "          (norm1): LayerNorm((1024,), eps=1e-06, elementwise_affine=True)\n",
      "          (attn): Attention(\n",
      "            (qkv): Linear(in_features=1024, out_features=3072, bias=True)\n",
      "            (attn_drop): Dropout(p=0.0, inplace=False)\n",
      "            (proj): Linear(in_features=1024, out_features=1024, bias=True)\n",
      "            (proj_drop): Dropout(p=0.0, inplace=False)\n",
      "          )\n",
      "          (ls1): Identity()\n",
      "          (drop_path1): Identity()\n",
      "          (norm2): LayerNorm((1024,), eps=1e-06, elementwise_affine=True)\n",
      "          (mlp): Mlp(\n",
      "            (fc1): Linear(in_features=1024, out_features=4096, bias=True)\n",
      "            (act): GELU(approximate='none')\n",
      "            (drop1): Dropout(p=0.0, inplace=False)\n",
      "            (fc2): Linear(in_features=4096, out_features=1024, bias=True)\n",
      "            (drop2): Dropout(p=0.0, inplace=False)\n",
      "          )\n",
      "          (ls2): Identity()\n",
      "          (drop_path2): Identity()\n",
      "        )\n",
      "        (4): Block(\n",
      "          (norm1): LayerNorm((1024,), eps=1e-06, elementwise_affine=True)\n",
      "          (attn): Attention(\n",
      "            (qkv): Linear(in_features=1024, out_features=3072, bias=True)\n",
      "            (attn_drop): Dropout(p=0.0, inplace=False)\n",
      "            (proj): Linear(in_features=1024, out_features=1024, bias=True)\n",
      "            (proj_drop): Dropout(p=0.0, inplace=False)\n",
      "          )\n",
      "          (ls1): Identity()\n",
      "          (drop_path1): Identity()\n",
      "          (norm2): LayerNorm((1024,), eps=1e-06, elementwise_affine=True)\n",
      "          (mlp): Mlp(\n",
      "            (fc1): Linear(in_features=1024, out_features=4096, bias=True)\n",
      "            (act): GELU(approximate='none')\n",
      "            (drop1): Dropout(p=0.0, inplace=False)\n",
      "            (fc2): Linear(in_features=4096, out_features=1024, bias=True)\n",
      "            (drop2): Dropout(p=0.0, inplace=False)\n",
      "          )\n",
      "          (ls2): Identity()\n",
      "          (drop_path2): Identity()\n",
      "        )\n",
      "        (5): Block(\n",
      "          (norm1): LayerNorm((1024,), eps=1e-06, elementwise_affine=True)\n",
      "          (attn): Attention(\n",
      "            (qkv): Linear(in_features=1024, out_features=3072, bias=True)\n",
      "            (attn_drop): Dropout(p=0.0, inplace=False)\n",
      "            (proj): Linear(in_features=1024, out_features=1024, bias=True)\n",
      "            (proj_drop): Dropout(p=0.0, inplace=False)\n",
      "          )\n",
      "          (ls1): Identity()\n",
      "          (drop_path1): Identity()\n",
      "          (norm2): LayerNorm((1024,), eps=1e-06, elementwise_affine=True)\n",
      "          (mlp): Mlp(\n",
      "            (fc1): Linear(in_features=1024, out_features=4096, bias=True)\n",
      "            (act): GELU(approximate='none')\n",
      "            (drop1): Dropout(p=0.0, inplace=False)\n",
      "            (fc2): Linear(in_features=4096, out_features=1024, bias=True)\n",
      "            (drop2): Dropout(p=0.0, inplace=False)\n",
      "          )\n",
      "          (ls2): Identity()\n",
      "          (drop_path2): Identity()\n",
      "        )\n",
      "        (6): Block(\n",
      "          (norm1): LayerNorm((1024,), eps=1e-06, elementwise_affine=True)\n",
      "          (attn): Attention(\n",
      "            (qkv): Linear(in_features=1024, out_features=3072, bias=True)\n",
      "            (attn_drop): Dropout(p=0.0, inplace=False)\n",
      "            (proj): Linear(in_features=1024, out_features=1024, bias=True)\n",
      "            (proj_drop): Dropout(p=0.0, inplace=False)\n",
      "          )\n",
      "          (ls1): Identity()\n",
      "          (drop_path1): Identity()\n",
      "          (norm2): LayerNorm((1024,), eps=1e-06, elementwise_affine=True)\n",
      "          (mlp): Mlp(\n",
      "            (fc1): Linear(in_features=1024, out_features=4096, bias=True)\n",
      "            (act): GELU(approximate='none')\n",
      "            (drop1): Dropout(p=0.0, inplace=False)\n",
      "            (fc2): Linear(in_features=4096, out_features=1024, bias=True)\n",
      "            (drop2): Dropout(p=0.0, inplace=False)\n",
      "          )\n",
      "          (ls2): Identity()\n",
      "          (drop_path2): Identity()\n",
      "        )\n",
      "        (7): Block(\n",
      "          (norm1): LayerNorm((1024,), eps=1e-06, elementwise_affine=True)\n",
      "          (attn): Attention(\n",
      "            (qkv): Linear(in_features=1024, out_features=3072, bias=True)\n",
      "            (attn_drop): Dropout(p=0.0, inplace=False)\n",
      "            (proj): Linear(in_features=1024, out_features=1024, bias=True)\n",
      "            (proj_drop): Dropout(p=0.0, inplace=False)\n",
      "          )\n",
      "          (ls1): Identity()\n",
      "          (drop_path1): Identity()\n",
      "          (norm2): LayerNorm((1024,), eps=1e-06, elementwise_affine=True)\n",
      "          (mlp): Mlp(\n",
      "            (fc1): Linear(in_features=1024, out_features=4096, bias=True)\n",
      "            (act): GELU(approximate='none')\n",
      "            (drop1): Dropout(p=0.0, inplace=False)\n",
      "            (fc2): Linear(in_features=4096, out_features=1024, bias=True)\n",
      "            (drop2): Dropout(p=0.0, inplace=False)\n",
      "          )\n",
      "          (ls2): Identity()\n",
      "          (drop_path2): Identity()\n",
      "        )\n",
      "        (8): Block(\n",
      "          (norm1): LayerNorm((1024,), eps=1e-06, elementwise_affine=True)\n",
      "          (attn): Attention(\n",
      "            (qkv): Linear(in_features=1024, out_features=3072, bias=True)\n",
      "            (attn_drop): Dropout(p=0.0, inplace=False)\n",
      "            (proj): Linear(in_features=1024, out_features=1024, bias=True)\n",
      "            (proj_drop): Dropout(p=0.0, inplace=False)\n",
      "          )\n",
      "          (ls1): Identity()\n",
      "          (drop_path1): Identity()\n",
      "          (norm2): LayerNorm((1024,), eps=1e-06, elementwise_affine=True)\n",
      "          (mlp): Mlp(\n",
      "            (fc1): Linear(in_features=1024, out_features=4096, bias=True)\n",
      "            (act): GELU(approximate='none')\n",
      "            (drop1): Dropout(p=0.0, inplace=False)\n",
      "            (fc2): Linear(in_features=4096, out_features=1024, bias=True)\n",
      "            (drop2): Dropout(p=0.0, inplace=False)\n",
      "          )\n",
      "          (ls2): Identity()\n",
      "          (drop_path2): Identity()\n",
      "        )\n",
      "        (9): Block(\n",
      "          (norm1): LayerNorm((1024,), eps=1e-06, elementwise_affine=True)\n",
      "          (attn): Attention(\n",
      "            (qkv): Linear(in_features=1024, out_features=3072, bias=True)\n",
      "            (attn_drop): Dropout(p=0.0, inplace=False)\n",
      "            (proj): Linear(in_features=1024, out_features=1024, bias=True)\n",
      "            (proj_drop): Dropout(p=0.0, inplace=False)\n",
      "          )\n",
      "          (ls1): Identity()\n",
      "          (drop_path1): Identity()\n",
      "          (norm2): LayerNorm((1024,), eps=1e-06, elementwise_affine=True)\n",
      "          (mlp): Mlp(\n",
      "            (fc1): Linear(in_features=1024, out_features=4096, bias=True)\n",
      "            (act): GELU(approximate='none')\n",
      "            (drop1): Dropout(p=0.0, inplace=False)\n",
      "            (fc2): Linear(in_features=4096, out_features=1024, bias=True)\n",
      "            (drop2): Dropout(p=0.0, inplace=False)\n",
      "          )\n",
      "          (ls2): Identity()\n",
      "          (drop_path2): Identity()\n",
      "        )\n",
      "        (10): Block(\n",
      "          (norm1): LayerNorm((1024,), eps=1e-06, elementwise_affine=True)\n",
      "          (attn): Attention(\n",
      "            (qkv): Linear(in_features=1024, out_features=3072, bias=True)\n",
      "            (attn_drop): Dropout(p=0.0, inplace=False)\n",
      "            (proj): Linear(in_features=1024, out_features=1024, bias=True)\n",
      "            (proj_drop): Dropout(p=0.0, inplace=False)\n",
      "          )\n",
      "          (ls1): Identity()\n",
      "          (drop_path1): Identity()\n",
      "          (norm2): LayerNorm((1024,), eps=1e-06, elementwise_affine=True)\n",
      "          (mlp): Mlp(\n",
      "            (fc1): Linear(in_features=1024, out_features=4096, bias=True)\n",
      "            (act): GELU(approximate='none')\n",
      "            (drop1): Dropout(p=0.0, inplace=False)\n",
      "            (fc2): Linear(in_features=4096, out_features=1024, bias=True)\n",
      "            (drop2): Dropout(p=0.0, inplace=False)\n",
      "          )\n",
      "          (ls2): Identity()\n",
      "          (drop_path2): Identity()\n",
      "        )\n",
      "        (11): Block(\n",
      "          (norm1): LayerNorm((1024,), eps=1e-06, elementwise_affine=True)\n",
      "          (attn): Attention(\n",
      "            (qkv): Linear(in_features=1024, out_features=3072, bias=True)\n",
      "            (attn_drop): Dropout(p=0.0, inplace=False)\n",
      "            (proj): Linear(in_features=1024, out_features=1024, bias=True)\n",
      "            (proj_drop): Dropout(p=0.0, inplace=False)\n",
      "          )\n",
      "          (ls1): Identity()\n",
      "          (drop_path1): Identity()\n",
      "          (norm2): LayerNorm((1024,), eps=1e-06, elementwise_affine=True)\n",
      "          (mlp): Mlp(\n",
      "            (fc1): Linear(in_features=1024, out_features=4096, bias=True)\n",
      "            (act): GELU(approximate='none')\n",
      "            (drop1): Dropout(p=0.0, inplace=False)\n",
      "            (fc2): Linear(in_features=4096, out_features=1024, bias=True)\n",
      "            (drop2): Dropout(p=0.0, inplace=False)\n",
      "          )\n",
      "          (ls2): Identity()\n",
      "          (drop_path2): Identity()\n",
      "        )\n",
      "        (12): Block(\n",
      "          (norm1): LayerNorm((1024,), eps=1e-06, elementwise_affine=True)\n",
      "          (attn): Attention(\n",
      "            (qkv): Linear(in_features=1024, out_features=3072, bias=True)\n",
      "            (attn_drop): Dropout(p=0.0, inplace=False)\n",
      "            (proj): Linear(in_features=1024, out_features=1024, bias=True)\n",
      "            (proj_drop): Dropout(p=0.0, inplace=False)\n",
      "          )\n",
      "          (ls1): Identity()\n",
      "          (drop_path1): Identity()\n",
      "          (norm2): LayerNorm((1024,), eps=1e-06, elementwise_affine=True)\n",
      "          (mlp): Mlp(\n",
      "            (fc1): Linear(in_features=1024, out_features=4096, bias=True)\n",
      "            (act): GELU(approximate='none')\n",
      "            (drop1): Dropout(p=0.0, inplace=False)\n",
      "            (fc2): Linear(in_features=4096, out_features=1024, bias=True)\n",
      "            (drop2): Dropout(p=0.0, inplace=False)\n",
      "          )\n",
      "          (ls2): Identity()\n",
      "          (drop_path2): Identity()\n",
      "        )\n",
      "        (13): Block(\n",
      "          (norm1): LayerNorm((1024,), eps=1e-06, elementwise_affine=True)\n",
      "          (attn): Attention(\n",
      "            (qkv): Linear(in_features=1024, out_features=3072, bias=True)\n",
      "            (attn_drop): Dropout(p=0.0, inplace=False)\n",
      "            (proj): Linear(in_features=1024, out_features=1024, bias=True)\n",
      "            (proj_drop): Dropout(p=0.0, inplace=False)\n",
      "          )\n",
      "          (ls1): Identity()\n",
      "          (drop_path1): Identity()\n",
      "          (norm2): LayerNorm((1024,), eps=1e-06, elementwise_affine=True)\n",
      "          (mlp): Mlp(\n",
      "            (fc1): Linear(in_features=1024, out_features=4096, bias=True)\n",
      "            (act): GELU(approximate='none')\n",
      "            (drop1): Dropout(p=0.0, inplace=False)\n",
      "            (fc2): Linear(in_features=4096, out_features=1024, bias=True)\n",
      "            (drop2): Dropout(p=0.0, inplace=False)\n",
      "          )\n",
      "          (ls2): Identity()\n",
      "          (drop_path2): Identity()\n",
      "        )\n",
      "        (14): Block(\n",
      "          (norm1): LayerNorm((1024,), eps=1e-06, elementwise_affine=True)\n",
      "          (attn): Attention(\n",
      "            (qkv): Linear(in_features=1024, out_features=3072, bias=True)\n",
      "            (attn_drop): Dropout(p=0.0, inplace=False)\n",
      "            (proj): Linear(in_features=1024, out_features=1024, bias=True)\n",
      "            (proj_drop): Dropout(p=0.0, inplace=False)\n",
      "          )\n",
      "          (ls1): Identity()\n",
      "          (drop_path1): Identity()\n",
      "          (norm2): LayerNorm((1024,), eps=1e-06, elementwise_affine=True)\n",
      "          (mlp): Mlp(\n",
      "            (fc1): Linear(in_features=1024, out_features=4096, bias=True)\n",
      "            (act): GELU(approximate='none')\n",
      "            (drop1): Dropout(p=0.0, inplace=False)\n",
      "            (fc2): Linear(in_features=4096, out_features=1024, bias=True)\n",
      "            (drop2): Dropout(p=0.0, inplace=False)\n",
      "          )\n",
      "          (ls2): Identity()\n",
      "          (drop_path2): Identity()\n",
      "        )\n",
      "        (15): Block(\n",
      "          (norm1): LayerNorm((1024,), eps=1e-06, elementwise_affine=True)\n",
      "          (attn): Attention(\n",
      "            (qkv): Linear(in_features=1024, out_features=3072, bias=True)\n",
      "            (attn_drop): Dropout(p=0.0, inplace=False)\n",
      "            (proj): Linear(in_features=1024, out_features=1024, bias=True)\n",
      "            (proj_drop): Dropout(p=0.0, inplace=False)\n",
      "          )\n",
      "          (ls1): Identity()\n",
      "          (drop_path1): Identity()\n",
      "          (norm2): LayerNorm((1024,), eps=1e-06, elementwise_affine=True)\n",
      "          (mlp): Mlp(\n",
      "            (fc1): Linear(in_features=1024, out_features=4096, bias=True)\n",
      "            (act): GELU(approximate='none')\n",
      "            (drop1): Dropout(p=0.0, inplace=False)\n",
      "            (fc2): Linear(in_features=4096, out_features=1024, bias=True)\n",
      "            (drop2): Dropout(p=0.0, inplace=False)\n",
      "          )\n",
      "          (ls2): Identity()\n",
      "          (drop_path2): Identity()\n",
      "        )\n",
      "        (16): Block(\n",
      "          (norm1): LayerNorm((1024,), eps=1e-06, elementwise_affine=True)\n",
      "          (attn): Attention(\n",
      "            (qkv): Linear(in_features=1024, out_features=3072, bias=True)\n",
      "            (attn_drop): Dropout(p=0.0, inplace=False)\n",
      "            (proj): Linear(in_features=1024, out_features=1024, bias=True)\n",
      "            (proj_drop): Dropout(p=0.0, inplace=False)\n",
      "          )\n",
      "          (ls1): Identity()\n",
      "          (drop_path1): Identity()\n",
      "          (norm2): LayerNorm((1024,), eps=1e-06, elementwise_affine=True)\n",
      "          (mlp): Mlp(\n",
      "            (fc1): Linear(in_features=1024, out_features=4096, bias=True)\n",
      "            (act): GELU(approximate='none')\n",
      "            (drop1): Dropout(p=0.0, inplace=False)\n",
      "            (fc2): Linear(in_features=4096, out_features=1024, bias=True)\n",
      "            (drop2): Dropout(p=0.0, inplace=False)\n",
      "          )\n",
      "          (ls2): Identity()\n",
      "          (drop_path2): Identity()\n",
      "        )\n",
      "        (17): Block(\n",
      "          (norm1): LayerNorm((1024,), eps=1e-06, elementwise_affine=True)\n",
      "          (attn): Attention(\n",
      "            (qkv): Linear(in_features=1024, out_features=3072, bias=True)\n",
      "            (attn_drop): Dropout(p=0.0, inplace=False)\n",
      "            (proj): Linear(in_features=1024, out_features=1024, bias=True)\n",
      "            (proj_drop): Dropout(p=0.0, inplace=False)\n",
      "          )\n",
      "          (ls1): Identity()\n",
      "          (drop_path1): Identity()\n",
      "          (norm2): LayerNorm((1024,), eps=1e-06, elementwise_affine=True)\n",
      "          (mlp): Mlp(\n",
      "            (fc1): Linear(in_features=1024, out_features=4096, bias=True)\n",
      "            (act): GELU(approximate='none')\n",
      "            (drop1): Dropout(p=0.0, inplace=False)\n",
      "            (fc2): Linear(in_features=4096, out_features=1024, bias=True)\n",
      "            (drop2): Dropout(p=0.0, inplace=False)\n",
      "          )\n",
      "          (ls2): Identity()\n",
      "          (drop_path2): Identity()\n",
      "        )\n",
      "        (18): Block(\n",
      "          (norm1): LayerNorm((1024,), eps=1e-06, elementwise_affine=True)\n",
      "          (attn): Attention(\n",
      "            (qkv): Linear(in_features=1024, out_features=3072, bias=True)\n",
      "            (attn_drop): Dropout(p=0.0, inplace=False)\n",
      "            (proj): Linear(in_features=1024, out_features=1024, bias=True)\n",
      "            (proj_drop): Dropout(p=0.0, inplace=False)\n",
      "          )\n",
      "          (ls1): Identity()\n",
      "          (drop_path1): Identity()\n",
      "          (norm2): LayerNorm((1024,), eps=1e-06, elementwise_affine=True)\n",
      "          (mlp): Mlp(\n",
      "            (fc1): Linear(in_features=1024, out_features=4096, bias=True)\n",
      "            (act): GELU(approximate='none')\n",
      "            (drop1): Dropout(p=0.0, inplace=False)\n",
      "            (fc2): Linear(in_features=4096, out_features=1024, bias=True)\n",
      "            (drop2): Dropout(p=0.0, inplace=False)\n",
      "          )\n",
      "          (ls2): Identity()\n",
      "          (drop_path2): Identity()\n",
      "        )\n",
      "        (19): Block(\n",
      "          (norm1): LayerNorm((1024,), eps=1e-06, elementwise_affine=True)\n",
      "          (attn): Attention(\n",
      "            (qkv): Linear(in_features=1024, out_features=3072, bias=True)\n",
      "            (attn_drop): Dropout(p=0.0, inplace=False)\n",
      "            (proj): Linear(in_features=1024, out_features=1024, bias=True)\n",
      "            (proj_drop): Dropout(p=0.0, inplace=False)\n",
      "          )\n",
      "          (ls1): Identity()\n",
      "          (drop_path1): Identity()\n",
      "          (norm2): LayerNorm((1024,), eps=1e-06, elementwise_affine=True)\n",
      "          (mlp): Mlp(\n",
      "            (fc1): Linear(in_features=1024, out_features=4096, bias=True)\n",
      "            (act): GELU(approximate='none')\n",
      "            (drop1): Dropout(p=0.0, inplace=False)\n",
      "            (fc2): Linear(in_features=4096, out_features=1024, bias=True)\n",
      "            (drop2): Dropout(p=0.0, inplace=False)\n",
      "          )\n",
      "          (ls2): Identity()\n",
      "          (drop_path2): Identity()\n",
      "        )\n",
      "        (20): Block(\n",
      "          (norm1): LayerNorm((1024,), eps=1e-06, elementwise_affine=True)\n",
      "          (attn): Attention(\n",
      "            (qkv): Linear(in_features=1024, out_features=3072, bias=True)\n",
      "            (attn_drop): Dropout(p=0.0, inplace=False)\n",
      "            (proj): Linear(in_features=1024, out_features=1024, bias=True)\n",
      "            (proj_drop): Dropout(p=0.0, inplace=False)\n",
      "          )\n",
      "          (ls1): Identity()\n",
      "          (drop_path1): Identity()\n",
      "          (norm2): LayerNorm((1024,), eps=1e-06, elementwise_affine=True)\n",
      "          (mlp): Mlp(\n",
      "            (fc1): Linear(in_features=1024, out_features=4096, bias=True)\n",
      "            (act): GELU(approximate='none')\n",
      "            (drop1): Dropout(p=0.0, inplace=False)\n",
      "            (fc2): Linear(in_features=4096, out_features=1024, bias=True)\n",
      "            (drop2): Dropout(p=0.0, inplace=False)\n",
      "          )\n",
      "          (ls2): Identity()\n",
      "          (drop_path2): Identity()\n",
      "        )\n",
      "        (21): Block(\n",
      "          (norm1): LayerNorm((1024,), eps=1e-06, elementwise_affine=True)\n",
      "          (attn): Attention(\n",
      "            (qkv): Linear(in_features=1024, out_features=3072, bias=True)\n",
      "            (attn_drop): Dropout(p=0.0, inplace=False)\n",
      "            (proj): Linear(in_features=1024, out_features=1024, bias=True)\n",
      "            (proj_drop): Dropout(p=0.0, inplace=False)\n",
      "          )\n",
      "          (ls1): Identity()\n",
      "          (drop_path1): Identity()\n",
      "          (norm2): LayerNorm((1024,), eps=1e-06, elementwise_affine=True)\n",
      "          (mlp): Mlp(\n",
      "            (fc1): Linear(in_features=1024, out_features=4096, bias=True)\n",
      "            (act): GELU(approximate='none')\n",
      "            (drop1): Dropout(p=0.0, inplace=False)\n",
      "            (fc2): Linear(in_features=4096, out_features=1024, bias=True)\n",
      "            (drop2): Dropout(p=0.0, inplace=False)\n",
      "          )\n",
      "          (ls2): Identity()\n",
      "          (drop_path2): Identity()\n",
      "        )\n",
      "        (22): Block(\n",
      "          (norm1): LayerNorm((1024,), eps=1e-06, elementwise_affine=True)\n",
      "          (attn): Attention(\n",
      "            (qkv): Linear(in_features=1024, out_features=3072, bias=True)\n",
      "            (attn_drop): Dropout(p=0.0, inplace=False)\n",
      "            (proj): Linear(in_features=1024, out_features=1024, bias=True)\n",
      "            (proj_drop): Dropout(p=0.0, inplace=False)\n",
      "          )\n",
      "          (ls1): Identity()\n",
      "          (drop_path1): Identity()\n",
      "          (norm2): LayerNorm((1024,), eps=1e-06, elementwise_affine=True)\n",
      "          (mlp): Mlp(\n",
      "            (fc1): Linear(in_features=1024, out_features=4096, bias=True)\n",
      "            (act): GELU(approximate='none')\n",
      "            (drop1): Dropout(p=0.0, inplace=False)\n",
      "            (fc2): Linear(in_features=4096, out_features=1024, bias=True)\n",
      "            (drop2): Dropout(p=0.0, inplace=False)\n",
      "          )\n",
      "          (ls2): Identity()\n",
      "          (drop_path2): Identity()\n",
      "        )\n",
      "        (23): Block(\n",
      "          (norm1): LayerNorm((1024,), eps=1e-06, elementwise_affine=True)\n",
      "          (attn): Attention(\n",
      "            (qkv): Linear(in_features=1024, out_features=3072, bias=True)\n",
      "            (attn_drop): Dropout(p=0.0, inplace=False)\n",
      "            (proj): Linear(in_features=1024, out_features=1024, bias=True)\n",
      "            (proj_drop): Dropout(p=0.0, inplace=False)\n",
      "          )\n",
      "          (ls1): Identity()\n",
      "          (drop_path1): Identity()\n",
      "          (norm2): LayerNorm((1024,), eps=1e-06, elementwise_affine=True)\n",
      "          (mlp): Mlp(\n",
      "            (fc1): Linear(in_features=1024, out_features=4096, bias=True)\n",
      "            (act): GELU(approximate='none')\n",
      "            (drop1): Dropout(p=0.0, inplace=False)\n",
      "            (fc2): Linear(in_features=4096, out_features=1024, bias=True)\n",
      "            (drop2): Dropout(p=0.0, inplace=False)\n",
      "          )\n",
      "          (ls2): Identity()\n",
      "          (drop_path2): Identity()\n",
      "        )\n",
      "      )\n",
      "      (norm): LayerNorm((1024,), eps=1e-06, elementwise_affine=True)\n",
      "      (fc_norm): Identity()\n",
      "    )\n",
      "  )\n",
      "  (dynamics): CTRNN(\n",
      "    (ctrnn): CTRNN(\n",
      "      (fc_x2ah): Linear(in_features=1024, out_features=1024, bias=True)\n",
      "      (fc_h2ah): Linear(in_features=1024, out_features=1024, bias=False)\n",
      "      (fc_h2y): Linear(in_features=1024, out_features=1024, bias=True)\n",
      "    )\n",
      "    (regressor): Linear(in_features=1024, out_features=1024, bias=True)\n",
      "  )\n",
      ")\n",
      "======= Model layers =======\n",
      "['encoder', 'dynamics']\n"
     ]
    }
   ],
   "source": [
    "name = \"pfVC1_CTRNN_physion\"\n",
    "model, model_layers = load_pretrained_model(name)\n",
    "print(\"======= Model architecture =======\\n\", model)\n",
    "print(f\"======= Model layers =======\\n{model_layers}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7d43f8e1",
   "metadata": {},
   "source": [
    "## 1b. VC-1 + CTRNN, dynamics module pretrained on Kinetics-700"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "e2d2b423",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Loaded parameters from /om/weka/yanglab/anayebi/mental-sim/trained_models/VC-1+CTRNN_k700.pt\n",
      "======= Model architecture =======\n",
      " FrozenPretrainedEncoder(\n",
      "  (encoder): VC1_pretrained(\n",
      "    (vc1): VisionTransformer(\n",
      "      (patch_embed): PatchEmbed(\n",
      "        (proj): Conv2d(3, 1024, kernel_size=(16, 16), stride=(16, 16))\n",
      "        (norm): Identity()\n",
      "      )\n",
      "      (pos_drop): Dropout(p=0.0, inplace=False)\n",
      "      (norm_pre): Identity()\n",
      "      (blocks): Sequential(\n",
      "        (0): Block(\n",
      "          (norm1): LayerNorm((1024,), eps=1e-06, elementwise_affine=True)\n",
      "          (attn): Attention(\n",
      "            (qkv): Linear(in_features=1024, out_features=3072, bias=True)\n",
      "            (attn_drop): Dropout(p=0.0, inplace=False)\n",
      "            (proj): Linear(in_features=1024, out_features=1024, bias=True)\n",
      "            (proj_drop): Dropout(p=0.0, inplace=False)\n",
      "          )\n",
      "          (ls1): Identity()\n",
      "          (drop_path1): Identity()\n",
      "          (norm2): LayerNorm((1024,), eps=1e-06, elementwise_affine=True)\n",
      "          (mlp): Mlp(\n",
      "            (fc1): Linear(in_features=1024, out_features=4096, bias=True)\n",
      "            (act): GELU(approximate='none')\n",
      "            (drop1): Dropout(p=0.0, inplace=False)\n",
      "            (fc2): Linear(in_features=4096, out_features=1024, bias=True)\n",
      "            (drop2): Dropout(p=0.0, inplace=False)\n",
      "          )\n",
      "          (ls2): Identity()\n",
      "          (drop_path2): Identity()\n",
      "        )\n",
      "        (1): Block(\n",
      "          (norm1): LayerNorm((1024,), eps=1e-06, elementwise_affine=True)\n",
      "          (attn): Attention(\n",
      "            (qkv): Linear(in_features=1024, out_features=3072, bias=True)\n",
      "            (attn_drop): Dropout(p=0.0, inplace=False)\n",
      "            (proj): Linear(in_features=1024, out_features=1024, bias=True)\n",
      "            (proj_drop): Dropout(p=0.0, inplace=False)\n",
      "          )\n",
      "          (ls1): Identity()\n",
      "          (drop_path1): Identity()\n",
      "          (norm2): LayerNorm((1024,), eps=1e-06, elementwise_affine=True)\n",
      "          (mlp): Mlp(\n",
      "            (fc1): Linear(in_features=1024, out_features=4096, bias=True)\n",
      "            (act): GELU(approximate='none')\n",
      "            (drop1): Dropout(p=0.0, inplace=False)\n",
      "            (fc2): Linear(in_features=4096, out_features=1024, bias=True)\n",
      "            (drop2): Dropout(p=0.0, inplace=False)\n",
      "          )\n",
      "          (ls2): Identity()\n",
      "          (drop_path2): Identity()\n",
      "        )\n",
      "        (2): Block(\n",
      "          (norm1): LayerNorm((1024,), eps=1e-06, elementwise_affine=True)\n",
      "          (attn): Attention(\n",
      "            (qkv): Linear(in_features=1024, out_features=3072, bias=True)\n",
      "            (attn_drop): Dropout(p=0.0, inplace=False)\n",
      "            (proj): Linear(in_features=1024, out_features=1024, bias=True)\n",
      "            (proj_drop): Dropout(p=0.0, inplace=False)\n",
      "          )\n",
      "          (ls1): Identity()\n",
      "          (drop_path1): Identity()\n",
      "          (norm2): LayerNorm((1024,), eps=1e-06, elementwise_affine=True)\n",
      "          (mlp): Mlp(\n",
      "            (fc1): Linear(in_features=1024, out_features=4096, bias=True)\n",
      "            (act): GELU(approximate='none')\n",
      "            (drop1): Dropout(p=0.0, inplace=False)\n",
      "            (fc2): Linear(in_features=4096, out_features=1024, bias=True)\n",
      "            (drop2): Dropout(p=0.0, inplace=False)\n",
      "          )\n",
      "          (ls2): Identity()\n",
      "          (drop_path2): Identity()\n",
      "        )\n",
      "        (3): Block(\n",
      "          (norm1): LayerNorm((1024,), eps=1e-06, elementwise_affine=True)\n",
      "          (attn): Attention(\n",
      "            (qkv): Linear(in_features=1024, out_features=3072, bias=True)\n",
      "            (attn_drop): Dropout(p=0.0, inplace=False)\n",
      "            (proj): Linear(in_features=1024, out_features=1024, bias=True)\n",
      "            (proj_drop): Dropout(p=0.0, inplace=False)\n",
      "          )\n",
      "          (ls1): Identity()\n",
      "          (drop_path1): Identity()\n",
      "          (norm2): LayerNorm((1024,), eps=1e-06, elementwise_affine=True)\n",
      "          (mlp): Mlp(\n",
      "            (fc1): Linear(in_features=1024, out_features=4096, bias=True)\n",
      "            (act): GELU(approximate='none')\n",
      "            (drop1): Dropout(p=0.0, inplace=False)\n",
      "            (fc2): Linear(in_features=4096, out_features=1024, bias=True)\n",
      "            (drop2): Dropout(p=0.0, inplace=False)\n",
      "          )\n",
      "          (ls2): Identity()\n",
      "          (drop_path2): Identity()\n",
      "        )\n",
      "        (4): Block(\n",
      "          (norm1): LayerNorm((1024,), eps=1e-06, elementwise_affine=True)\n",
      "          (attn): Attention(\n",
      "            (qkv): Linear(in_features=1024, out_features=3072, bias=True)\n",
      "            (attn_drop): Dropout(p=0.0, inplace=False)\n",
      "            (proj): Linear(in_features=1024, out_features=1024, bias=True)\n",
      "            (proj_drop): Dropout(p=0.0, inplace=False)\n",
      "          )\n",
      "          (ls1): Identity()\n",
      "          (drop_path1): Identity()\n",
      "          (norm2): LayerNorm((1024,), eps=1e-06, elementwise_affine=True)\n",
      "          (mlp): Mlp(\n",
      "            (fc1): Linear(in_features=1024, out_features=4096, bias=True)\n",
      "            (act): GELU(approximate='none')\n",
      "            (drop1): Dropout(p=0.0, inplace=False)\n",
      "            (fc2): Linear(in_features=4096, out_features=1024, bias=True)\n",
      "            (drop2): Dropout(p=0.0, inplace=False)\n",
      "          )\n",
      "          (ls2): Identity()\n",
      "          (drop_path2): Identity()\n",
      "        )\n",
      "        (5): Block(\n",
      "          (norm1): LayerNorm((1024,), eps=1e-06, elementwise_affine=True)\n",
      "          (attn): Attention(\n",
      "            (qkv): Linear(in_features=1024, out_features=3072, bias=True)\n",
      "            (attn_drop): Dropout(p=0.0, inplace=False)\n",
      "            (proj): Linear(in_features=1024, out_features=1024, bias=True)\n",
      "            (proj_drop): Dropout(p=0.0, inplace=False)\n",
      "          )\n",
      "          (ls1): Identity()\n",
      "          (drop_path1): Identity()\n",
      "          (norm2): LayerNorm((1024,), eps=1e-06, elementwise_affine=True)\n",
      "          (mlp): Mlp(\n",
      "            (fc1): Linear(in_features=1024, out_features=4096, bias=True)\n",
      "            (act): GELU(approximate='none')\n",
      "            (drop1): Dropout(p=0.0, inplace=False)\n",
      "            (fc2): Linear(in_features=4096, out_features=1024, bias=True)\n",
      "            (drop2): Dropout(p=0.0, inplace=False)\n",
      "          )\n",
      "          (ls2): Identity()\n",
      "          (drop_path2): Identity()\n",
      "        )\n",
      "        (6): Block(\n",
      "          (norm1): LayerNorm((1024,), eps=1e-06, elementwise_affine=True)\n",
      "          (attn): Attention(\n",
      "            (qkv): Linear(in_features=1024, out_features=3072, bias=True)\n",
      "            (attn_drop): Dropout(p=0.0, inplace=False)\n",
      "            (proj): Linear(in_features=1024, out_features=1024, bias=True)\n",
      "            (proj_drop): Dropout(p=0.0, inplace=False)\n",
      "          )\n",
      "          (ls1): Identity()\n",
      "          (drop_path1): Identity()\n",
      "          (norm2): LayerNorm((1024,), eps=1e-06, elementwise_affine=True)\n",
      "          (mlp): Mlp(\n",
      "            (fc1): Linear(in_features=1024, out_features=4096, bias=True)\n",
      "            (act): GELU(approximate='none')\n",
      "            (drop1): Dropout(p=0.0, inplace=False)\n",
      "            (fc2): Linear(in_features=4096, out_features=1024, bias=True)\n",
      "            (drop2): Dropout(p=0.0, inplace=False)\n",
      "          )\n",
      "          (ls2): Identity()\n",
      "          (drop_path2): Identity()\n",
      "        )\n",
      "        (7): Block(\n",
      "          (norm1): LayerNorm((1024,), eps=1e-06, elementwise_affine=True)\n",
      "          (attn): Attention(\n",
      "            (qkv): Linear(in_features=1024, out_features=3072, bias=True)\n",
      "            (attn_drop): Dropout(p=0.0, inplace=False)\n",
      "            (proj): Linear(in_features=1024, out_features=1024, bias=True)\n",
      "            (proj_drop): Dropout(p=0.0, inplace=False)\n",
      "          )\n",
      "          (ls1): Identity()\n",
      "          (drop_path1): Identity()\n",
      "          (norm2): LayerNorm((1024,), eps=1e-06, elementwise_affine=True)\n",
      "          (mlp): Mlp(\n",
      "            (fc1): Linear(in_features=1024, out_features=4096, bias=True)\n",
      "            (act): GELU(approximate='none')\n",
      "            (drop1): Dropout(p=0.0, inplace=False)\n",
      "            (fc2): Linear(in_features=4096, out_features=1024, bias=True)\n",
      "            (drop2): Dropout(p=0.0, inplace=False)\n",
      "          )\n",
      "          (ls2): Identity()\n",
      "          (drop_path2): Identity()\n",
      "        )\n",
      "        (8): Block(\n",
      "          (norm1): LayerNorm((1024,), eps=1e-06, elementwise_affine=True)\n",
      "          (attn): Attention(\n",
      "            (qkv): Linear(in_features=1024, out_features=3072, bias=True)\n",
      "            (attn_drop): Dropout(p=0.0, inplace=False)\n",
      "            (proj): Linear(in_features=1024, out_features=1024, bias=True)\n",
      "            (proj_drop): Dropout(p=0.0, inplace=False)\n",
      "          )\n",
      "          (ls1): Identity()\n",
      "          (drop_path1): Identity()\n",
      "          (norm2): LayerNorm((1024,), eps=1e-06, elementwise_affine=True)\n",
      "          (mlp): Mlp(\n",
      "            (fc1): Linear(in_features=1024, out_features=4096, bias=True)\n",
      "            (act): GELU(approximate='none')\n",
      "            (drop1): Dropout(p=0.0, inplace=False)\n",
      "            (fc2): Linear(in_features=4096, out_features=1024, bias=True)\n",
      "            (drop2): Dropout(p=0.0, inplace=False)\n",
      "          )\n",
      "          (ls2): Identity()\n",
      "          (drop_path2): Identity()\n",
      "        )\n",
      "        (9): Block(\n",
      "          (norm1): LayerNorm((1024,), eps=1e-06, elementwise_affine=True)\n",
      "          (attn): Attention(\n",
      "            (qkv): Linear(in_features=1024, out_features=3072, bias=True)\n",
      "            (attn_drop): Dropout(p=0.0, inplace=False)\n",
      "            (proj): Linear(in_features=1024, out_features=1024, bias=True)\n",
      "            (proj_drop): Dropout(p=0.0, inplace=False)\n",
      "          )\n",
      "          (ls1): Identity()\n",
      "          (drop_path1): Identity()\n",
      "          (norm2): LayerNorm((1024,), eps=1e-06, elementwise_affine=True)\n",
      "          (mlp): Mlp(\n",
      "            (fc1): Linear(in_features=1024, out_features=4096, bias=True)\n",
      "            (act): GELU(approximate='none')\n",
      "            (drop1): Dropout(p=0.0, inplace=False)\n",
      "            (fc2): Linear(in_features=4096, out_features=1024, bias=True)\n",
      "            (drop2): Dropout(p=0.0, inplace=False)\n",
      "          )\n",
      "          (ls2): Identity()\n",
      "          (drop_path2): Identity()\n",
      "        )\n",
      "        (10): Block(\n",
      "          (norm1): LayerNorm((1024,), eps=1e-06, elementwise_affine=True)\n",
      "          (attn): Attention(\n",
      "            (qkv): Linear(in_features=1024, out_features=3072, bias=True)\n",
      "            (attn_drop): Dropout(p=0.0, inplace=False)\n",
      "            (proj): Linear(in_features=1024, out_features=1024, bias=True)\n",
      "            (proj_drop): Dropout(p=0.0, inplace=False)\n",
      "          )\n",
      "          (ls1): Identity()\n",
      "          (drop_path1): Identity()\n",
      "          (norm2): LayerNorm((1024,), eps=1e-06, elementwise_affine=True)\n",
      "          (mlp): Mlp(\n",
      "            (fc1): Linear(in_features=1024, out_features=4096, bias=True)\n",
      "            (act): GELU(approximate='none')\n",
      "            (drop1): Dropout(p=0.0, inplace=False)\n",
      "            (fc2): Linear(in_features=4096, out_features=1024, bias=True)\n",
      "            (drop2): Dropout(p=0.0, inplace=False)\n",
      "          )\n",
      "          (ls2): Identity()\n",
      "          (drop_path2): Identity()\n",
      "        )\n",
      "        (11): Block(\n",
      "          (norm1): LayerNorm((1024,), eps=1e-06, elementwise_affine=True)\n",
      "          (attn): Attention(\n",
      "            (qkv): Linear(in_features=1024, out_features=3072, bias=True)\n",
      "            (attn_drop): Dropout(p=0.0, inplace=False)\n",
      "            (proj): Linear(in_features=1024, out_features=1024, bias=True)\n",
      "            (proj_drop): Dropout(p=0.0, inplace=False)\n",
      "          )\n",
      "          (ls1): Identity()\n",
      "          (drop_path1): Identity()\n",
      "          (norm2): LayerNorm((1024,), eps=1e-06, elementwise_affine=True)\n",
      "          (mlp): Mlp(\n",
      "            (fc1): Linear(in_features=1024, out_features=4096, bias=True)\n",
      "            (act): GELU(approximate='none')\n",
      "            (drop1): Dropout(p=0.0, inplace=False)\n",
      "            (fc2): Linear(in_features=4096, out_features=1024, bias=True)\n",
      "            (drop2): Dropout(p=0.0, inplace=False)\n",
      "          )\n",
      "          (ls2): Identity()\n",
      "          (drop_path2): Identity()\n",
      "        )\n",
      "        (12): Block(\n",
      "          (norm1): LayerNorm((1024,), eps=1e-06, elementwise_affine=True)\n",
      "          (attn): Attention(\n",
      "            (qkv): Linear(in_features=1024, out_features=3072, bias=True)\n",
      "            (attn_drop): Dropout(p=0.0, inplace=False)\n",
      "            (proj): Linear(in_features=1024, out_features=1024, bias=True)\n",
      "            (proj_drop): Dropout(p=0.0, inplace=False)\n",
      "          )\n",
      "          (ls1): Identity()\n",
      "          (drop_path1): Identity()\n",
      "          (norm2): LayerNorm((1024,), eps=1e-06, elementwise_affine=True)\n",
      "          (mlp): Mlp(\n",
      "            (fc1): Linear(in_features=1024, out_features=4096, bias=True)\n",
      "            (act): GELU(approximate='none')\n",
      "            (drop1): Dropout(p=0.0, inplace=False)\n",
      "            (fc2): Linear(in_features=4096, out_features=1024, bias=True)\n",
      "            (drop2): Dropout(p=0.0, inplace=False)\n",
      "          )\n",
      "          (ls2): Identity()\n",
      "          (drop_path2): Identity()\n",
      "        )\n",
      "        (13): Block(\n",
      "          (norm1): LayerNorm((1024,), eps=1e-06, elementwise_affine=True)\n",
      "          (attn): Attention(\n",
      "            (qkv): Linear(in_features=1024, out_features=3072, bias=True)\n",
      "            (attn_drop): Dropout(p=0.0, inplace=False)\n",
      "            (proj): Linear(in_features=1024, out_features=1024, bias=True)\n",
      "            (proj_drop): Dropout(p=0.0, inplace=False)\n",
      "          )\n",
      "          (ls1): Identity()\n",
      "          (drop_path1): Identity()\n",
      "          (norm2): LayerNorm((1024,), eps=1e-06, elementwise_affine=True)\n",
      "          (mlp): Mlp(\n",
      "            (fc1): Linear(in_features=1024, out_features=4096, bias=True)\n",
      "            (act): GELU(approximate='none')\n",
      "            (drop1): Dropout(p=0.0, inplace=False)\n",
      "            (fc2): Linear(in_features=4096, out_features=1024, bias=True)\n",
      "            (drop2): Dropout(p=0.0, inplace=False)\n",
      "          )\n",
      "          (ls2): Identity()\n",
      "          (drop_path2): Identity()\n",
      "        )\n",
      "        (14): Block(\n",
      "          (norm1): LayerNorm((1024,), eps=1e-06, elementwise_affine=True)\n",
      "          (attn): Attention(\n",
      "            (qkv): Linear(in_features=1024, out_features=3072, bias=True)\n",
      "            (attn_drop): Dropout(p=0.0, inplace=False)\n",
      "            (proj): Linear(in_features=1024, out_features=1024, bias=True)\n",
      "            (proj_drop): Dropout(p=0.0, inplace=False)\n",
      "          )\n",
      "          (ls1): Identity()\n",
      "          (drop_path1): Identity()\n",
      "          (norm2): LayerNorm((1024,), eps=1e-06, elementwise_affine=True)\n",
      "          (mlp): Mlp(\n",
      "            (fc1): Linear(in_features=1024, out_features=4096, bias=True)\n",
      "            (act): GELU(approximate='none')\n",
      "            (drop1): Dropout(p=0.0, inplace=False)\n",
      "            (fc2): Linear(in_features=4096, out_features=1024, bias=True)\n",
      "            (drop2): Dropout(p=0.0, inplace=False)\n",
      "          )\n",
      "          (ls2): Identity()\n",
      "          (drop_path2): Identity()\n",
      "        )\n",
      "        (15): Block(\n",
      "          (norm1): LayerNorm((1024,), eps=1e-06, elementwise_affine=True)\n",
      "          (attn): Attention(\n",
      "            (qkv): Linear(in_features=1024, out_features=3072, bias=True)\n",
      "            (attn_drop): Dropout(p=0.0, inplace=False)\n",
      "            (proj): Linear(in_features=1024, out_features=1024, bias=True)\n",
      "            (proj_drop): Dropout(p=0.0, inplace=False)\n",
      "          )\n",
      "          (ls1): Identity()\n",
      "          (drop_path1): Identity()\n",
      "          (norm2): LayerNorm((1024,), eps=1e-06, elementwise_affine=True)\n",
      "          (mlp): Mlp(\n",
      "            (fc1): Linear(in_features=1024, out_features=4096, bias=True)\n",
      "            (act): GELU(approximate='none')\n",
      "            (drop1): Dropout(p=0.0, inplace=False)\n",
      "            (fc2): Linear(in_features=4096, out_features=1024, bias=True)\n",
      "            (drop2): Dropout(p=0.0, inplace=False)\n",
      "          )\n",
      "          (ls2): Identity()\n",
      "          (drop_path2): Identity()\n",
      "        )\n",
      "        (16): Block(\n",
      "          (norm1): LayerNorm((1024,), eps=1e-06, elementwise_affine=True)\n",
      "          (attn): Attention(\n",
      "            (qkv): Linear(in_features=1024, out_features=3072, bias=True)\n",
      "            (attn_drop): Dropout(p=0.0, inplace=False)\n",
      "            (proj): Linear(in_features=1024, out_features=1024, bias=True)\n",
      "            (proj_drop): Dropout(p=0.0, inplace=False)\n",
      "          )\n",
      "          (ls1): Identity()\n",
      "          (drop_path1): Identity()\n",
      "          (norm2): LayerNorm((1024,), eps=1e-06, elementwise_affine=True)\n",
      "          (mlp): Mlp(\n",
      "            (fc1): Linear(in_features=1024, out_features=4096, bias=True)\n",
      "            (act): GELU(approximate='none')\n",
      "            (drop1): Dropout(p=0.0, inplace=False)\n",
      "            (fc2): Linear(in_features=4096, out_features=1024, bias=True)\n",
      "            (drop2): Dropout(p=0.0, inplace=False)\n",
      "          )\n",
      "          (ls2): Identity()\n",
      "          (drop_path2): Identity()\n",
      "        )\n",
      "        (17): Block(\n",
      "          (norm1): LayerNorm((1024,), eps=1e-06, elementwise_affine=True)\n",
      "          (attn): Attention(\n",
      "            (qkv): Linear(in_features=1024, out_features=3072, bias=True)\n",
      "            (attn_drop): Dropout(p=0.0, inplace=False)\n",
      "            (proj): Linear(in_features=1024, out_features=1024, bias=True)\n",
      "            (proj_drop): Dropout(p=0.0, inplace=False)\n",
      "          )\n",
      "          (ls1): Identity()\n",
      "          (drop_path1): Identity()\n",
      "          (norm2): LayerNorm((1024,), eps=1e-06, elementwise_affine=True)\n",
      "          (mlp): Mlp(\n",
      "            (fc1): Linear(in_features=1024, out_features=4096, bias=True)\n",
      "            (act): GELU(approximate='none')\n",
      "            (drop1): Dropout(p=0.0, inplace=False)\n",
      "            (fc2): Linear(in_features=4096, out_features=1024, bias=True)\n",
      "            (drop2): Dropout(p=0.0, inplace=False)\n",
      "          )\n",
      "          (ls2): Identity()\n",
      "          (drop_path2): Identity()\n",
      "        )\n",
      "        (18): Block(\n",
      "          (norm1): LayerNorm((1024,), eps=1e-06, elementwise_affine=True)\n",
      "          (attn): Attention(\n",
      "            (qkv): Linear(in_features=1024, out_features=3072, bias=True)\n",
      "            (attn_drop): Dropout(p=0.0, inplace=False)\n",
      "            (proj): Linear(in_features=1024, out_features=1024, bias=True)\n",
      "            (proj_drop): Dropout(p=0.0, inplace=False)\n",
      "          )\n",
      "          (ls1): Identity()\n",
      "          (drop_path1): Identity()\n",
      "          (norm2): LayerNorm((1024,), eps=1e-06, elementwise_affine=True)\n",
      "          (mlp): Mlp(\n",
      "            (fc1): Linear(in_features=1024, out_features=4096, bias=True)\n",
      "            (act): GELU(approximate='none')\n",
      "            (drop1): Dropout(p=0.0, inplace=False)\n",
      "            (fc2): Linear(in_features=4096, out_features=1024, bias=True)\n",
      "            (drop2): Dropout(p=0.0, inplace=False)\n",
      "          )\n",
      "          (ls2): Identity()\n",
      "          (drop_path2): Identity()\n",
      "        )\n",
      "        (19): Block(\n",
      "          (norm1): LayerNorm((1024,), eps=1e-06, elementwise_affine=True)\n",
      "          (attn): Attention(\n",
      "            (qkv): Linear(in_features=1024, out_features=3072, bias=True)\n",
      "            (attn_drop): Dropout(p=0.0, inplace=False)\n",
      "            (proj): Linear(in_features=1024, out_features=1024, bias=True)\n",
      "            (proj_drop): Dropout(p=0.0, inplace=False)\n",
      "          )\n",
      "          (ls1): Identity()\n",
      "          (drop_path1): Identity()\n",
      "          (norm2): LayerNorm((1024,), eps=1e-06, elementwise_affine=True)\n",
      "          (mlp): Mlp(\n",
      "            (fc1): Linear(in_features=1024, out_features=4096, bias=True)\n",
      "            (act): GELU(approximate='none')\n",
      "            (drop1): Dropout(p=0.0, inplace=False)\n",
      "            (fc2): Linear(in_features=4096, out_features=1024, bias=True)\n",
      "            (drop2): Dropout(p=0.0, inplace=False)\n",
      "          )\n",
      "          (ls2): Identity()\n",
      "          (drop_path2): Identity()\n",
      "        )\n",
      "        (20): Block(\n",
      "          (norm1): LayerNorm((1024,), eps=1e-06, elementwise_affine=True)\n",
      "          (attn): Attention(\n",
      "            (qkv): Linear(in_features=1024, out_features=3072, bias=True)\n",
      "            (attn_drop): Dropout(p=0.0, inplace=False)\n",
      "            (proj): Linear(in_features=1024, out_features=1024, bias=True)\n",
      "            (proj_drop): Dropout(p=0.0, inplace=False)\n",
      "          )\n",
      "          (ls1): Identity()\n",
      "          (drop_path1): Identity()\n",
      "          (norm2): LayerNorm((1024,), eps=1e-06, elementwise_affine=True)\n",
      "          (mlp): Mlp(\n",
      "            (fc1): Linear(in_features=1024, out_features=4096, bias=True)\n",
      "            (act): GELU(approximate='none')\n",
      "            (drop1): Dropout(p=0.0, inplace=False)\n",
      "            (fc2): Linear(in_features=4096, out_features=1024, bias=True)\n",
      "            (drop2): Dropout(p=0.0, inplace=False)\n",
      "          )\n",
      "          (ls2): Identity()\n",
      "          (drop_path2): Identity()\n",
      "        )\n",
      "        (21): Block(\n",
      "          (norm1): LayerNorm((1024,), eps=1e-06, elementwise_affine=True)\n",
      "          (attn): Attention(\n",
      "            (qkv): Linear(in_features=1024, out_features=3072, bias=True)\n",
      "            (attn_drop): Dropout(p=0.0, inplace=False)\n",
      "            (proj): Linear(in_features=1024, out_features=1024, bias=True)\n",
      "            (proj_drop): Dropout(p=0.0, inplace=False)\n",
      "          )\n",
      "          (ls1): Identity()\n",
      "          (drop_path1): Identity()\n",
      "          (norm2): LayerNorm((1024,), eps=1e-06, elementwise_affine=True)\n",
      "          (mlp): Mlp(\n",
      "            (fc1): Linear(in_features=1024, out_features=4096, bias=True)\n",
      "            (act): GELU(approximate='none')\n",
      "            (drop1): Dropout(p=0.0, inplace=False)\n",
      "            (fc2): Linear(in_features=4096, out_features=1024, bias=True)\n",
      "            (drop2): Dropout(p=0.0, inplace=False)\n",
      "          )\n",
      "          (ls2): Identity()\n",
      "          (drop_path2): Identity()\n",
      "        )\n",
      "        (22): Block(\n",
      "          (norm1): LayerNorm((1024,), eps=1e-06, elementwise_affine=True)\n",
      "          (attn): Attention(\n",
      "            (qkv): Linear(in_features=1024, out_features=3072, bias=True)\n",
      "            (attn_drop): Dropout(p=0.0, inplace=False)\n",
      "            (proj): Linear(in_features=1024, out_features=1024, bias=True)\n",
      "            (proj_drop): Dropout(p=0.0, inplace=False)\n",
      "          )\n",
      "          (ls1): Identity()\n",
      "          (drop_path1): Identity()\n",
      "          (norm2): LayerNorm((1024,), eps=1e-06, elementwise_affine=True)\n",
      "          (mlp): Mlp(\n",
      "            (fc1): Linear(in_features=1024, out_features=4096, bias=True)\n",
      "            (act): GELU(approximate='none')\n",
      "            (drop1): Dropout(p=0.0, inplace=False)\n",
      "            (fc2): Linear(in_features=4096, out_features=1024, bias=True)\n",
      "            (drop2): Dropout(p=0.0, inplace=False)\n",
      "          )\n",
      "          (ls2): Identity()\n",
      "          (drop_path2): Identity()\n",
      "        )\n",
      "        (23): Block(\n",
      "          (norm1): LayerNorm((1024,), eps=1e-06, elementwise_affine=True)\n",
      "          (attn): Attention(\n",
      "            (qkv): Linear(in_features=1024, out_features=3072, bias=True)\n",
      "            (attn_drop): Dropout(p=0.0, inplace=False)\n",
      "            (proj): Linear(in_features=1024, out_features=1024, bias=True)\n",
      "            (proj_drop): Dropout(p=0.0, inplace=False)\n",
      "          )\n",
      "          (ls1): Identity()\n",
      "          (drop_path1): Identity()\n",
      "          (norm2): LayerNorm((1024,), eps=1e-06, elementwise_affine=True)\n",
      "          (mlp): Mlp(\n",
      "            (fc1): Linear(in_features=1024, out_features=4096, bias=True)\n",
      "            (act): GELU(approximate='none')\n",
      "            (drop1): Dropout(p=0.0, inplace=False)\n",
      "            (fc2): Linear(in_features=4096, out_features=1024, bias=True)\n",
      "            (drop2): Dropout(p=0.0, inplace=False)\n",
      "          )\n",
      "          (ls2): Identity()\n",
      "          (drop_path2): Identity()\n",
      "        )\n",
      "      )\n",
      "      (norm): LayerNorm((1024,), eps=1e-06, elementwise_affine=True)\n",
      "      (fc_norm): Identity()\n",
      "    )\n",
      "  )\n",
      "  (dynamics): CTRNN(\n",
      "    (ctrnn): CTRNN(\n",
      "      (fc_x2ah): Linear(in_features=1024, out_features=1024, bias=True)\n",
      "      (fc_h2ah): Linear(in_features=1024, out_features=1024, bias=False)\n",
      "      (fc_h2y): Linear(in_features=1024, out_features=1024, bias=True)\n",
      "    )\n",
      "    (regressor): Linear(in_features=1024, out_features=1024, bias=True)\n",
      "  )\n",
      ")\n",
      "======= Model layers =======\n",
      "['encoder', 'dynamics']\n"
     ]
    }
   ],
   "source": [
    "name = \"pfVC1_CTRNN_k700\"\n",
    "model, model_layers = load_pretrained_model(name)\n",
    "print(\"======= Model architecture =======\\n\", model)\n",
    "print(f\"======= Model layers =======\\n{model_layers}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c9676b00",
   "metadata": {},
   "source": [
    "## 1c. VC-1 + LSTM, dynamics module pretrained on Physion"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "20a44d37",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Loaded parameters from /om/weka/yanglab/anayebi/mental-sim/trained_models/VC-1+LSTM_physion.pt\n",
      "======= Model architecture =======\n",
      " FrozenPretrainedEncoder(\n",
      "  (encoder): VC1_pretrained(\n",
      "    (vc1): VisionTransformer(\n",
      "      (patch_embed): PatchEmbed(\n",
      "        (proj): Conv2d(3, 1024, kernel_size=(16, 16), stride=(16, 16))\n",
      "        (norm): Identity()\n",
      "      )\n",
      "      (pos_drop): Dropout(p=0.0, inplace=False)\n",
      "      (norm_pre): Identity()\n",
      "      (blocks): Sequential(\n",
      "        (0): Block(\n",
      "          (norm1): LayerNorm((1024,), eps=1e-06, elementwise_affine=True)\n",
      "          (attn): Attention(\n",
      "            (qkv): Linear(in_features=1024, out_features=3072, bias=True)\n",
      "            (attn_drop): Dropout(p=0.0, inplace=False)\n",
      "            (proj): Linear(in_features=1024, out_features=1024, bias=True)\n",
      "            (proj_drop): Dropout(p=0.0, inplace=False)\n",
      "          )\n",
      "          (ls1): Identity()\n",
      "          (drop_path1): Identity()\n",
      "          (norm2): LayerNorm((1024,), eps=1e-06, elementwise_affine=True)\n",
      "          (mlp): Mlp(\n",
      "            (fc1): Linear(in_features=1024, out_features=4096, bias=True)\n",
      "            (act): GELU(approximate='none')\n",
      "            (drop1): Dropout(p=0.0, inplace=False)\n",
      "            (fc2): Linear(in_features=4096, out_features=1024, bias=True)\n",
      "            (drop2): Dropout(p=0.0, inplace=False)\n",
      "          )\n",
      "          (ls2): Identity()\n",
      "          (drop_path2): Identity()\n",
      "        )\n",
      "        (1): Block(\n",
      "          (norm1): LayerNorm((1024,), eps=1e-06, elementwise_affine=True)\n",
      "          (attn): Attention(\n",
      "            (qkv): Linear(in_features=1024, out_features=3072, bias=True)\n",
      "            (attn_drop): Dropout(p=0.0, inplace=False)\n",
      "            (proj): Linear(in_features=1024, out_features=1024, bias=True)\n",
      "            (proj_drop): Dropout(p=0.0, inplace=False)\n",
      "          )\n",
      "          (ls1): Identity()\n",
      "          (drop_path1): Identity()\n",
      "          (norm2): LayerNorm((1024,), eps=1e-06, elementwise_affine=True)\n",
      "          (mlp): Mlp(\n",
      "            (fc1): Linear(in_features=1024, out_features=4096, bias=True)\n",
      "            (act): GELU(approximate='none')\n",
      "            (drop1): Dropout(p=0.0, inplace=False)\n",
      "            (fc2): Linear(in_features=4096, out_features=1024, bias=True)\n",
      "            (drop2): Dropout(p=0.0, inplace=False)\n",
      "          )\n",
      "          (ls2): Identity()\n",
      "          (drop_path2): Identity()\n",
      "        )\n",
      "        (2): Block(\n",
      "          (norm1): LayerNorm((1024,), eps=1e-06, elementwise_affine=True)\n",
      "          (attn): Attention(\n",
      "            (qkv): Linear(in_features=1024, out_features=3072, bias=True)\n",
      "            (attn_drop): Dropout(p=0.0, inplace=False)\n",
      "            (proj): Linear(in_features=1024, out_features=1024, bias=True)\n",
      "            (proj_drop): Dropout(p=0.0, inplace=False)\n",
      "          )\n",
      "          (ls1): Identity()\n",
      "          (drop_path1): Identity()\n",
      "          (norm2): LayerNorm((1024,), eps=1e-06, elementwise_affine=True)\n",
      "          (mlp): Mlp(\n",
      "            (fc1): Linear(in_features=1024, out_features=4096, bias=True)\n",
      "            (act): GELU(approximate='none')\n",
      "            (drop1): Dropout(p=0.0, inplace=False)\n",
      "            (fc2): Linear(in_features=4096, out_features=1024, bias=True)\n",
      "            (drop2): Dropout(p=0.0, inplace=False)\n",
      "          )\n",
      "          (ls2): Identity()\n",
      "          (drop_path2): Identity()\n",
      "        )\n",
      "        (3): Block(\n",
      "          (norm1): LayerNorm((1024,), eps=1e-06, elementwise_affine=True)\n",
      "          (attn): Attention(\n",
      "            (qkv): Linear(in_features=1024, out_features=3072, bias=True)\n",
      "            (attn_drop): Dropout(p=0.0, inplace=False)\n",
      "            (proj): Linear(in_features=1024, out_features=1024, bias=True)\n",
      "            (proj_drop): Dropout(p=0.0, inplace=False)\n",
      "          )\n",
      "          (ls1): Identity()\n",
      "          (drop_path1): Identity()\n",
      "          (norm2): LayerNorm((1024,), eps=1e-06, elementwise_affine=True)\n",
      "          (mlp): Mlp(\n",
      "            (fc1): Linear(in_features=1024, out_features=4096, bias=True)\n",
      "            (act): GELU(approximate='none')\n",
      "            (drop1): Dropout(p=0.0, inplace=False)\n",
      "            (fc2): Linear(in_features=4096, out_features=1024, bias=True)\n",
      "            (drop2): Dropout(p=0.0, inplace=False)\n",
      "          )\n",
      "          (ls2): Identity()\n",
      "          (drop_path2): Identity()\n",
      "        )\n",
      "        (4): Block(\n",
      "          (norm1): LayerNorm((1024,), eps=1e-06, elementwise_affine=True)\n",
      "          (attn): Attention(\n",
      "            (qkv): Linear(in_features=1024, out_features=3072, bias=True)\n",
      "            (attn_drop): Dropout(p=0.0, inplace=False)\n",
      "            (proj): Linear(in_features=1024, out_features=1024, bias=True)\n",
      "            (proj_drop): Dropout(p=0.0, inplace=False)\n",
      "          )\n",
      "          (ls1): Identity()\n",
      "          (drop_path1): Identity()\n",
      "          (norm2): LayerNorm((1024,), eps=1e-06, elementwise_affine=True)\n",
      "          (mlp): Mlp(\n",
      "            (fc1): Linear(in_features=1024, out_features=4096, bias=True)\n",
      "            (act): GELU(approximate='none')\n",
      "            (drop1): Dropout(p=0.0, inplace=False)\n",
      "            (fc2): Linear(in_features=4096, out_features=1024, bias=True)\n",
      "            (drop2): Dropout(p=0.0, inplace=False)\n",
      "          )\n",
      "          (ls2): Identity()\n",
      "          (drop_path2): Identity()\n",
      "        )\n",
      "        (5): Block(\n",
      "          (norm1): LayerNorm((1024,), eps=1e-06, elementwise_affine=True)\n",
      "          (attn): Attention(\n",
      "            (qkv): Linear(in_features=1024, out_features=3072, bias=True)\n",
      "            (attn_drop): Dropout(p=0.0, inplace=False)\n",
      "            (proj): Linear(in_features=1024, out_features=1024, bias=True)\n",
      "            (proj_drop): Dropout(p=0.0, inplace=False)\n",
      "          )\n",
      "          (ls1): Identity()\n",
      "          (drop_path1): Identity()\n",
      "          (norm2): LayerNorm((1024,), eps=1e-06, elementwise_affine=True)\n",
      "          (mlp): Mlp(\n",
      "            (fc1): Linear(in_features=1024, out_features=4096, bias=True)\n",
      "            (act): GELU(approximate='none')\n",
      "            (drop1): Dropout(p=0.0, inplace=False)\n",
      "            (fc2): Linear(in_features=4096, out_features=1024, bias=True)\n",
      "            (drop2): Dropout(p=0.0, inplace=False)\n",
      "          )\n",
      "          (ls2): Identity()\n",
      "          (drop_path2): Identity()\n",
      "        )\n",
      "        (6): Block(\n",
      "          (norm1): LayerNorm((1024,), eps=1e-06, elementwise_affine=True)\n",
      "          (attn): Attention(\n",
      "            (qkv): Linear(in_features=1024, out_features=3072, bias=True)\n",
      "            (attn_drop): Dropout(p=0.0, inplace=False)\n",
      "            (proj): Linear(in_features=1024, out_features=1024, bias=True)\n",
      "            (proj_drop): Dropout(p=0.0, inplace=False)\n",
      "          )\n",
      "          (ls1): Identity()\n",
      "          (drop_path1): Identity()\n",
      "          (norm2): LayerNorm((1024,), eps=1e-06, elementwise_affine=True)\n",
      "          (mlp): Mlp(\n",
      "            (fc1): Linear(in_features=1024, out_features=4096, bias=True)\n",
      "            (act): GELU(approximate='none')\n",
      "            (drop1): Dropout(p=0.0, inplace=False)\n",
      "            (fc2): Linear(in_features=4096, out_features=1024, bias=True)\n",
      "            (drop2): Dropout(p=0.0, inplace=False)\n",
      "          )\n",
      "          (ls2): Identity()\n",
      "          (drop_path2): Identity()\n",
      "        )\n",
      "        (7): Block(\n",
      "          (norm1): LayerNorm((1024,), eps=1e-06, elementwise_affine=True)\n",
      "          (attn): Attention(\n",
      "            (qkv): Linear(in_features=1024, out_features=3072, bias=True)\n",
      "            (attn_drop): Dropout(p=0.0, inplace=False)\n",
      "            (proj): Linear(in_features=1024, out_features=1024, bias=True)\n",
      "            (proj_drop): Dropout(p=0.0, inplace=False)\n",
      "          )\n",
      "          (ls1): Identity()\n",
      "          (drop_path1): Identity()\n",
      "          (norm2): LayerNorm((1024,), eps=1e-06, elementwise_affine=True)\n",
      "          (mlp): Mlp(\n",
      "            (fc1): Linear(in_features=1024, out_features=4096, bias=True)\n",
      "            (act): GELU(approximate='none')\n",
      "            (drop1): Dropout(p=0.0, inplace=False)\n",
      "            (fc2): Linear(in_features=4096, out_features=1024, bias=True)\n",
      "            (drop2): Dropout(p=0.0, inplace=False)\n",
      "          )\n",
      "          (ls2): Identity()\n",
      "          (drop_path2): Identity()\n",
      "        )\n",
      "        (8): Block(\n",
      "          (norm1): LayerNorm((1024,), eps=1e-06, elementwise_affine=True)\n",
      "          (attn): Attention(\n",
      "            (qkv): Linear(in_features=1024, out_features=3072, bias=True)\n",
      "            (attn_drop): Dropout(p=0.0, inplace=False)\n",
      "            (proj): Linear(in_features=1024, out_features=1024, bias=True)\n",
      "            (proj_drop): Dropout(p=0.0, inplace=False)\n",
      "          )\n",
      "          (ls1): Identity()\n",
      "          (drop_path1): Identity()\n",
      "          (norm2): LayerNorm((1024,), eps=1e-06, elementwise_affine=True)\n",
      "          (mlp): Mlp(\n",
      "            (fc1): Linear(in_features=1024, out_features=4096, bias=True)\n",
      "            (act): GELU(approximate='none')\n",
      "            (drop1): Dropout(p=0.0, inplace=False)\n",
      "            (fc2): Linear(in_features=4096, out_features=1024, bias=True)\n",
      "            (drop2): Dropout(p=0.0, inplace=False)\n",
      "          )\n",
      "          (ls2): Identity()\n",
      "          (drop_path2): Identity()\n",
      "        )\n",
      "        (9): Block(\n",
      "          (norm1): LayerNorm((1024,), eps=1e-06, elementwise_affine=True)\n",
      "          (attn): Attention(\n",
      "            (qkv): Linear(in_features=1024, out_features=3072, bias=True)\n",
      "            (attn_drop): Dropout(p=0.0, inplace=False)\n",
      "            (proj): Linear(in_features=1024, out_features=1024, bias=True)\n",
      "            (proj_drop): Dropout(p=0.0, inplace=False)\n",
      "          )\n",
      "          (ls1): Identity()\n",
      "          (drop_path1): Identity()\n",
      "          (norm2): LayerNorm((1024,), eps=1e-06, elementwise_affine=True)\n",
      "          (mlp): Mlp(\n",
      "            (fc1): Linear(in_features=1024, out_features=4096, bias=True)\n",
      "            (act): GELU(approximate='none')\n",
      "            (drop1): Dropout(p=0.0, inplace=False)\n",
      "            (fc2): Linear(in_features=4096, out_features=1024, bias=True)\n",
      "            (drop2): Dropout(p=0.0, inplace=False)\n",
      "          )\n",
      "          (ls2): Identity()\n",
      "          (drop_path2): Identity()\n",
      "        )\n",
      "        (10): Block(\n",
      "          (norm1): LayerNorm((1024,), eps=1e-06, elementwise_affine=True)\n",
      "          (attn): Attention(\n",
      "            (qkv): Linear(in_features=1024, out_features=3072, bias=True)\n",
      "            (attn_drop): Dropout(p=0.0, inplace=False)\n",
      "            (proj): Linear(in_features=1024, out_features=1024, bias=True)\n",
      "            (proj_drop): Dropout(p=0.0, inplace=False)\n",
      "          )\n",
      "          (ls1): Identity()\n",
      "          (drop_path1): Identity()\n",
      "          (norm2): LayerNorm((1024,), eps=1e-06, elementwise_affine=True)\n",
      "          (mlp): Mlp(\n",
      "            (fc1): Linear(in_features=1024, out_features=4096, bias=True)\n",
      "            (act): GELU(approximate='none')\n",
      "            (drop1): Dropout(p=0.0, inplace=False)\n",
      "            (fc2): Linear(in_features=4096, out_features=1024, bias=True)\n",
      "            (drop2): Dropout(p=0.0, inplace=False)\n",
      "          )\n",
      "          (ls2): Identity()\n",
      "          (drop_path2): Identity()\n",
      "        )\n",
      "        (11): Block(\n",
      "          (norm1): LayerNorm((1024,), eps=1e-06, elementwise_affine=True)\n",
      "          (attn): Attention(\n",
      "            (qkv): Linear(in_features=1024, out_features=3072, bias=True)\n",
      "            (attn_drop): Dropout(p=0.0, inplace=False)\n",
      "            (proj): Linear(in_features=1024, out_features=1024, bias=True)\n",
      "            (proj_drop): Dropout(p=0.0, inplace=False)\n",
      "          )\n",
      "          (ls1): Identity()\n",
      "          (drop_path1): Identity()\n",
      "          (norm2): LayerNorm((1024,), eps=1e-06, elementwise_affine=True)\n",
      "          (mlp): Mlp(\n",
      "            (fc1): Linear(in_features=1024, out_features=4096, bias=True)\n",
      "            (act): GELU(approximate='none')\n",
      "            (drop1): Dropout(p=0.0, inplace=False)\n",
      "            (fc2): Linear(in_features=4096, out_features=1024, bias=True)\n",
      "            (drop2): Dropout(p=0.0, inplace=False)\n",
      "          )\n",
      "          (ls2): Identity()\n",
      "          (drop_path2): Identity()\n",
      "        )\n",
      "        (12): Block(\n",
      "          (norm1): LayerNorm((1024,), eps=1e-06, elementwise_affine=True)\n",
      "          (attn): Attention(\n",
      "            (qkv): Linear(in_features=1024, out_features=3072, bias=True)\n",
      "            (attn_drop): Dropout(p=0.0, inplace=False)\n",
      "            (proj): Linear(in_features=1024, out_features=1024, bias=True)\n",
      "            (proj_drop): Dropout(p=0.0, inplace=False)\n",
      "          )\n",
      "          (ls1): Identity()\n",
      "          (drop_path1): Identity()\n",
      "          (norm2): LayerNorm((1024,), eps=1e-06, elementwise_affine=True)\n",
      "          (mlp): Mlp(\n",
      "            (fc1): Linear(in_features=1024, out_features=4096, bias=True)\n",
      "            (act): GELU(approximate='none')\n",
      "            (drop1): Dropout(p=0.0, inplace=False)\n",
      "            (fc2): Linear(in_features=4096, out_features=1024, bias=True)\n",
      "            (drop2): Dropout(p=0.0, inplace=False)\n",
      "          )\n",
      "          (ls2): Identity()\n",
      "          (drop_path2): Identity()\n",
      "        )\n",
      "        (13): Block(\n",
      "          (norm1): LayerNorm((1024,), eps=1e-06, elementwise_affine=True)\n",
      "          (attn): Attention(\n",
      "            (qkv): Linear(in_features=1024, out_features=3072, bias=True)\n",
      "            (attn_drop): Dropout(p=0.0, inplace=False)\n",
      "            (proj): Linear(in_features=1024, out_features=1024, bias=True)\n",
      "            (proj_drop): Dropout(p=0.0, inplace=False)\n",
      "          )\n",
      "          (ls1): Identity()\n",
      "          (drop_path1): Identity()\n",
      "          (norm2): LayerNorm((1024,), eps=1e-06, elementwise_affine=True)\n",
      "          (mlp): Mlp(\n",
      "            (fc1): Linear(in_features=1024, out_features=4096, bias=True)\n",
      "            (act): GELU(approximate='none')\n",
      "            (drop1): Dropout(p=0.0, inplace=False)\n",
      "            (fc2): Linear(in_features=4096, out_features=1024, bias=True)\n",
      "            (drop2): Dropout(p=0.0, inplace=False)\n",
      "          )\n",
      "          (ls2): Identity()\n",
      "          (drop_path2): Identity()\n",
      "        )\n",
      "        (14): Block(\n",
      "          (norm1): LayerNorm((1024,), eps=1e-06, elementwise_affine=True)\n",
      "          (attn): Attention(\n",
      "            (qkv): Linear(in_features=1024, out_features=3072, bias=True)\n",
      "            (attn_drop): Dropout(p=0.0, inplace=False)\n",
      "            (proj): Linear(in_features=1024, out_features=1024, bias=True)\n",
      "            (proj_drop): Dropout(p=0.0, inplace=False)\n",
      "          )\n",
      "          (ls1): Identity()\n",
      "          (drop_path1): Identity()\n",
      "          (norm2): LayerNorm((1024,), eps=1e-06, elementwise_affine=True)\n",
      "          (mlp): Mlp(\n",
      "            (fc1): Linear(in_features=1024, out_features=4096, bias=True)\n",
      "            (act): GELU(approximate='none')\n",
      "            (drop1): Dropout(p=0.0, inplace=False)\n",
      "            (fc2): Linear(in_features=4096, out_features=1024, bias=True)\n",
      "            (drop2): Dropout(p=0.0, inplace=False)\n",
      "          )\n",
      "          (ls2): Identity()\n",
      "          (drop_path2): Identity()\n",
      "        )\n",
      "        (15): Block(\n",
      "          (norm1): LayerNorm((1024,), eps=1e-06, elementwise_affine=True)\n",
      "          (attn): Attention(\n",
      "            (qkv): Linear(in_features=1024, out_features=3072, bias=True)\n",
      "            (attn_drop): Dropout(p=0.0, inplace=False)\n",
      "            (proj): Linear(in_features=1024, out_features=1024, bias=True)\n",
      "            (proj_drop): Dropout(p=0.0, inplace=False)\n",
      "          )\n",
      "          (ls1): Identity()\n",
      "          (drop_path1): Identity()\n",
      "          (norm2): LayerNorm((1024,), eps=1e-06, elementwise_affine=True)\n",
      "          (mlp): Mlp(\n",
      "            (fc1): Linear(in_features=1024, out_features=4096, bias=True)\n",
      "            (act): GELU(approximate='none')\n",
      "            (drop1): Dropout(p=0.0, inplace=False)\n",
      "            (fc2): Linear(in_features=4096, out_features=1024, bias=True)\n",
      "            (drop2): Dropout(p=0.0, inplace=False)\n",
      "          )\n",
      "          (ls2): Identity()\n",
      "          (drop_path2): Identity()\n",
      "        )\n",
      "        (16): Block(\n",
      "          (norm1): LayerNorm((1024,), eps=1e-06, elementwise_affine=True)\n",
      "          (attn): Attention(\n",
      "            (qkv): Linear(in_features=1024, out_features=3072, bias=True)\n",
      "            (attn_drop): Dropout(p=0.0, inplace=False)\n",
      "            (proj): Linear(in_features=1024, out_features=1024, bias=True)\n",
      "            (proj_drop): Dropout(p=0.0, inplace=False)\n",
      "          )\n",
      "          (ls1): Identity()\n",
      "          (drop_path1): Identity()\n",
      "          (norm2): LayerNorm((1024,), eps=1e-06, elementwise_affine=True)\n",
      "          (mlp): Mlp(\n",
      "            (fc1): Linear(in_features=1024, out_features=4096, bias=True)\n",
      "            (act): GELU(approximate='none')\n",
      "            (drop1): Dropout(p=0.0, inplace=False)\n",
      "            (fc2): Linear(in_features=4096, out_features=1024, bias=True)\n",
      "            (drop2): Dropout(p=0.0, inplace=False)\n",
      "          )\n",
      "          (ls2): Identity()\n",
      "          (drop_path2): Identity()\n",
      "        )\n",
      "        (17): Block(\n",
      "          (norm1): LayerNorm((1024,), eps=1e-06, elementwise_affine=True)\n",
      "          (attn): Attention(\n",
      "            (qkv): Linear(in_features=1024, out_features=3072, bias=True)\n",
      "            (attn_drop): Dropout(p=0.0, inplace=False)\n",
      "            (proj): Linear(in_features=1024, out_features=1024, bias=True)\n",
      "            (proj_drop): Dropout(p=0.0, inplace=False)\n",
      "          )\n",
      "          (ls1): Identity()\n",
      "          (drop_path1): Identity()\n",
      "          (norm2): LayerNorm((1024,), eps=1e-06, elementwise_affine=True)\n",
      "          (mlp): Mlp(\n",
      "            (fc1): Linear(in_features=1024, out_features=4096, bias=True)\n",
      "            (act): GELU(approximate='none')\n",
      "            (drop1): Dropout(p=0.0, inplace=False)\n",
      "            (fc2): Linear(in_features=4096, out_features=1024, bias=True)\n",
      "            (drop2): Dropout(p=0.0, inplace=False)\n",
      "          )\n",
      "          (ls2): Identity()\n",
      "          (drop_path2): Identity()\n",
      "        )\n",
      "        (18): Block(\n",
      "          (norm1): LayerNorm((1024,), eps=1e-06, elementwise_affine=True)\n",
      "          (attn): Attention(\n",
      "            (qkv): Linear(in_features=1024, out_features=3072, bias=True)\n",
      "            (attn_drop): Dropout(p=0.0, inplace=False)\n",
      "            (proj): Linear(in_features=1024, out_features=1024, bias=True)\n",
      "            (proj_drop): Dropout(p=0.0, inplace=False)\n",
      "          )\n",
      "          (ls1): Identity()\n",
      "          (drop_path1): Identity()\n",
      "          (norm2): LayerNorm((1024,), eps=1e-06, elementwise_affine=True)\n",
      "          (mlp): Mlp(\n",
      "            (fc1): Linear(in_features=1024, out_features=4096, bias=True)\n",
      "            (act): GELU(approximate='none')\n",
      "            (drop1): Dropout(p=0.0, inplace=False)\n",
      "            (fc2): Linear(in_features=4096, out_features=1024, bias=True)\n",
      "            (drop2): Dropout(p=0.0, inplace=False)\n",
      "          )\n",
      "          (ls2): Identity()\n",
      "          (drop_path2): Identity()\n",
      "        )\n",
      "        (19): Block(\n",
      "          (norm1): LayerNorm((1024,), eps=1e-06, elementwise_affine=True)\n",
      "          (attn): Attention(\n",
      "            (qkv): Linear(in_features=1024, out_features=3072, bias=True)\n",
      "            (attn_drop): Dropout(p=0.0, inplace=False)\n",
      "            (proj): Linear(in_features=1024, out_features=1024, bias=True)\n",
      "            (proj_drop): Dropout(p=0.0, inplace=False)\n",
      "          )\n",
      "          (ls1): Identity()\n",
      "          (drop_path1): Identity()\n",
      "          (norm2): LayerNorm((1024,), eps=1e-06, elementwise_affine=True)\n",
      "          (mlp): Mlp(\n",
      "            (fc1): Linear(in_features=1024, out_features=4096, bias=True)\n",
      "            (act): GELU(approximate='none')\n",
      "            (drop1): Dropout(p=0.0, inplace=False)\n",
      "            (fc2): Linear(in_features=4096, out_features=1024, bias=True)\n",
      "            (drop2): Dropout(p=0.0, inplace=False)\n",
      "          )\n",
      "          (ls2): Identity()\n",
      "          (drop_path2): Identity()\n",
      "        )\n",
      "        (20): Block(\n",
      "          (norm1): LayerNorm((1024,), eps=1e-06, elementwise_affine=True)\n",
      "          (attn): Attention(\n",
      "            (qkv): Linear(in_features=1024, out_features=3072, bias=True)\n",
      "            (attn_drop): Dropout(p=0.0, inplace=False)\n",
      "            (proj): Linear(in_features=1024, out_features=1024, bias=True)\n",
      "            (proj_drop): Dropout(p=0.0, inplace=False)\n",
      "          )\n",
      "          (ls1): Identity()\n",
      "          (drop_path1): Identity()\n",
      "          (norm2): LayerNorm((1024,), eps=1e-06, elementwise_affine=True)\n",
      "          (mlp): Mlp(\n",
      "            (fc1): Linear(in_features=1024, out_features=4096, bias=True)\n",
      "            (act): GELU(approximate='none')\n",
      "            (drop1): Dropout(p=0.0, inplace=False)\n",
      "            (fc2): Linear(in_features=4096, out_features=1024, bias=True)\n",
      "            (drop2): Dropout(p=0.0, inplace=False)\n",
      "          )\n",
      "          (ls2): Identity()\n",
      "          (drop_path2): Identity()\n",
      "        )\n",
      "        (21): Block(\n",
      "          (norm1): LayerNorm((1024,), eps=1e-06, elementwise_affine=True)\n",
      "          (attn): Attention(\n",
      "            (qkv): Linear(in_features=1024, out_features=3072, bias=True)\n",
      "            (attn_drop): Dropout(p=0.0, inplace=False)\n",
      "            (proj): Linear(in_features=1024, out_features=1024, bias=True)\n",
      "            (proj_drop): Dropout(p=0.0, inplace=False)\n",
      "          )\n",
      "          (ls1): Identity()\n",
      "          (drop_path1): Identity()\n",
      "          (norm2): LayerNorm((1024,), eps=1e-06, elementwise_affine=True)\n",
      "          (mlp): Mlp(\n",
      "            (fc1): Linear(in_features=1024, out_features=4096, bias=True)\n",
      "            (act): GELU(approximate='none')\n",
      "            (drop1): Dropout(p=0.0, inplace=False)\n",
      "            (fc2): Linear(in_features=4096, out_features=1024, bias=True)\n",
      "            (drop2): Dropout(p=0.0, inplace=False)\n",
      "          )\n",
      "          (ls2): Identity()\n",
      "          (drop_path2): Identity()\n",
      "        )\n",
      "        (22): Block(\n",
      "          (norm1): LayerNorm((1024,), eps=1e-06, elementwise_affine=True)\n",
      "          (attn): Attention(\n",
      "            (qkv): Linear(in_features=1024, out_features=3072, bias=True)\n",
      "            (attn_drop): Dropout(p=0.0, inplace=False)\n",
      "            (proj): Linear(in_features=1024, out_features=1024, bias=True)\n",
      "            (proj_drop): Dropout(p=0.0, inplace=False)\n",
      "          )\n",
      "          (ls1): Identity()\n",
      "          (drop_path1): Identity()\n",
      "          (norm2): LayerNorm((1024,), eps=1e-06, elementwise_affine=True)\n",
      "          (mlp): Mlp(\n",
      "            (fc1): Linear(in_features=1024, out_features=4096, bias=True)\n",
      "            (act): GELU(approximate='none')\n",
      "            (drop1): Dropout(p=0.0, inplace=False)\n",
      "            (fc2): Linear(in_features=4096, out_features=1024, bias=True)\n",
      "            (drop2): Dropout(p=0.0, inplace=False)\n",
      "          )\n",
      "          (ls2): Identity()\n",
      "          (drop_path2): Identity()\n",
      "        )\n",
      "        (23): Block(\n",
      "          (norm1): LayerNorm((1024,), eps=1e-06, elementwise_affine=True)\n",
      "          (attn): Attention(\n",
      "            (qkv): Linear(in_features=1024, out_features=3072, bias=True)\n",
      "            (attn_drop): Dropout(p=0.0, inplace=False)\n",
      "            (proj): Linear(in_features=1024, out_features=1024, bias=True)\n",
      "            (proj_drop): Dropout(p=0.0, inplace=False)\n",
      "          )\n",
      "          (ls1): Identity()\n",
      "          (drop_path1): Identity()\n",
      "          (norm2): LayerNorm((1024,), eps=1e-06, elementwise_affine=True)\n",
      "          (mlp): Mlp(\n",
      "            (fc1): Linear(in_features=1024, out_features=4096, bias=True)\n",
      "            (act): GELU(approximate='none')\n",
      "            (drop1): Dropout(p=0.0, inplace=False)\n",
      "            (fc2): Linear(in_features=4096, out_features=1024, bias=True)\n",
      "            (drop2): Dropout(p=0.0, inplace=False)\n",
      "          )\n",
      "          (ls2): Identity()\n",
      "          (drop_path2): Identity()\n",
      "        )\n",
      "      )\n",
      "      (norm): LayerNorm((1024,), eps=1e-06, elementwise_affine=True)\n",
      "      (fc_norm): Identity()\n",
      "    )\n",
      "  )\n",
      "  (dynamics): LSTM(\n",
      "    (lstm): LSTM(1024, 1024)\n",
      "    (regressor): Linear(in_features=1024, out_features=1024, bias=True)\n",
      "  )\n",
      ")\n",
      "======= Model layers =======\n",
      "['encoder', 'dynamics']\n"
     ]
    }
   ],
   "source": [
    "name = \"pfVC1_LSTM_physion\"\n",
    "model, model_layers = load_pretrained_model(name)\n",
    "print(\"======= Model architecture =======\\n\", model)\n",
    "print(f\"======= Model layers =======\\n{model_layers}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1d6e29d1",
   "metadata": {},
   "source": [
    "## 1d. VC-1 + LSTM, dynamics module pretrained on Kinetics-700"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "eacb33a4",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Loaded parameters from /om/weka/yanglab/anayebi/mental-sim/trained_models/VC-1+LSTM_k700.pt\n",
      "======= Model architecture =======\n",
      " FrozenPretrainedEncoder(\n",
      "  (encoder): VC1_pretrained(\n",
      "    (vc1): VisionTransformer(\n",
      "      (patch_embed): PatchEmbed(\n",
      "        (proj): Conv2d(3, 1024, kernel_size=(16, 16), stride=(16, 16))\n",
      "        (norm): Identity()\n",
      "      )\n",
      "      (pos_drop): Dropout(p=0.0, inplace=False)\n",
      "      (norm_pre): Identity()\n",
      "      (blocks): Sequential(\n",
      "        (0): Block(\n",
      "          (norm1): LayerNorm((1024,), eps=1e-06, elementwise_affine=True)\n",
      "          (attn): Attention(\n",
      "            (qkv): Linear(in_features=1024, out_features=3072, bias=True)\n",
      "            (attn_drop): Dropout(p=0.0, inplace=False)\n",
      "            (proj): Linear(in_features=1024, out_features=1024, bias=True)\n",
      "            (proj_drop): Dropout(p=0.0, inplace=False)\n",
      "          )\n",
      "          (ls1): Identity()\n",
      "          (drop_path1): Identity()\n",
      "          (norm2): LayerNorm((1024,), eps=1e-06, elementwise_affine=True)\n",
      "          (mlp): Mlp(\n",
      "            (fc1): Linear(in_features=1024, out_features=4096, bias=True)\n",
      "            (act): GELU(approximate='none')\n",
      "            (drop1): Dropout(p=0.0, inplace=False)\n",
      "            (fc2): Linear(in_features=4096, out_features=1024, bias=True)\n",
      "            (drop2): Dropout(p=0.0, inplace=False)\n",
      "          )\n",
      "          (ls2): Identity()\n",
      "          (drop_path2): Identity()\n",
      "        )\n",
      "        (1): Block(\n",
      "          (norm1): LayerNorm((1024,), eps=1e-06, elementwise_affine=True)\n",
      "          (attn): Attention(\n",
      "            (qkv): Linear(in_features=1024, out_features=3072, bias=True)\n",
      "            (attn_drop): Dropout(p=0.0, inplace=False)\n",
      "            (proj): Linear(in_features=1024, out_features=1024, bias=True)\n",
      "            (proj_drop): Dropout(p=0.0, inplace=False)\n",
      "          )\n",
      "          (ls1): Identity()\n",
      "          (drop_path1): Identity()\n",
      "          (norm2): LayerNorm((1024,), eps=1e-06, elementwise_affine=True)\n",
      "          (mlp): Mlp(\n",
      "            (fc1): Linear(in_features=1024, out_features=4096, bias=True)\n",
      "            (act): GELU(approximate='none')\n",
      "            (drop1): Dropout(p=0.0, inplace=False)\n",
      "            (fc2): Linear(in_features=4096, out_features=1024, bias=True)\n",
      "            (drop2): Dropout(p=0.0, inplace=False)\n",
      "          )\n",
      "          (ls2): Identity()\n",
      "          (drop_path2): Identity()\n",
      "        )\n",
      "        (2): Block(\n",
      "          (norm1): LayerNorm((1024,), eps=1e-06, elementwise_affine=True)\n",
      "          (attn): Attention(\n",
      "            (qkv): Linear(in_features=1024, out_features=3072, bias=True)\n",
      "            (attn_drop): Dropout(p=0.0, inplace=False)\n",
      "            (proj): Linear(in_features=1024, out_features=1024, bias=True)\n",
      "            (proj_drop): Dropout(p=0.0, inplace=False)\n",
      "          )\n",
      "          (ls1): Identity()\n",
      "          (drop_path1): Identity()\n",
      "          (norm2): LayerNorm((1024,), eps=1e-06, elementwise_affine=True)\n",
      "          (mlp): Mlp(\n",
      "            (fc1): Linear(in_features=1024, out_features=4096, bias=True)\n",
      "            (act): GELU(approximate='none')\n",
      "            (drop1): Dropout(p=0.0, inplace=False)\n",
      "            (fc2): Linear(in_features=4096, out_features=1024, bias=True)\n",
      "            (drop2): Dropout(p=0.0, inplace=False)\n",
      "          )\n",
      "          (ls2): Identity()\n",
      "          (drop_path2): Identity()\n",
      "        )\n",
      "        (3): Block(\n",
      "          (norm1): LayerNorm((1024,), eps=1e-06, elementwise_affine=True)\n",
      "          (attn): Attention(\n",
      "            (qkv): Linear(in_features=1024, out_features=3072, bias=True)\n",
      "            (attn_drop): Dropout(p=0.0, inplace=False)\n",
      "            (proj): Linear(in_features=1024, out_features=1024, bias=True)\n",
      "            (proj_drop): Dropout(p=0.0, inplace=False)\n",
      "          )\n",
      "          (ls1): Identity()\n",
      "          (drop_path1): Identity()\n",
      "          (norm2): LayerNorm((1024,), eps=1e-06, elementwise_affine=True)\n",
      "          (mlp): Mlp(\n",
      "            (fc1): Linear(in_features=1024, out_features=4096, bias=True)\n",
      "            (act): GELU(approximate='none')\n",
      "            (drop1): Dropout(p=0.0, inplace=False)\n",
      "            (fc2): Linear(in_features=4096, out_features=1024, bias=True)\n",
      "            (drop2): Dropout(p=0.0, inplace=False)\n",
      "          )\n",
      "          (ls2): Identity()\n",
      "          (drop_path2): Identity()\n",
      "        )\n",
      "        (4): Block(\n",
      "          (norm1): LayerNorm((1024,), eps=1e-06, elementwise_affine=True)\n",
      "          (attn): Attention(\n",
      "            (qkv): Linear(in_features=1024, out_features=3072, bias=True)\n",
      "            (attn_drop): Dropout(p=0.0, inplace=False)\n",
      "            (proj): Linear(in_features=1024, out_features=1024, bias=True)\n",
      "            (proj_drop): Dropout(p=0.0, inplace=False)\n",
      "          )\n",
      "          (ls1): Identity()\n",
      "          (drop_path1): Identity()\n",
      "          (norm2): LayerNorm((1024,), eps=1e-06, elementwise_affine=True)\n",
      "          (mlp): Mlp(\n",
      "            (fc1): Linear(in_features=1024, out_features=4096, bias=True)\n",
      "            (act): GELU(approximate='none')\n",
      "            (drop1): Dropout(p=0.0, inplace=False)\n",
      "            (fc2): Linear(in_features=4096, out_features=1024, bias=True)\n",
      "            (drop2): Dropout(p=0.0, inplace=False)\n",
      "          )\n",
      "          (ls2): Identity()\n",
      "          (drop_path2): Identity()\n",
      "        )\n",
      "        (5): Block(\n",
      "          (norm1): LayerNorm((1024,), eps=1e-06, elementwise_affine=True)\n",
      "          (attn): Attention(\n",
      "            (qkv): Linear(in_features=1024, out_features=3072, bias=True)\n",
      "            (attn_drop): Dropout(p=0.0, inplace=False)\n",
      "            (proj): Linear(in_features=1024, out_features=1024, bias=True)\n",
      "            (proj_drop): Dropout(p=0.0, inplace=False)\n",
      "          )\n",
      "          (ls1): Identity()\n",
      "          (drop_path1): Identity()\n",
      "          (norm2): LayerNorm((1024,), eps=1e-06, elementwise_affine=True)\n",
      "          (mlp): Mlp(\n",
      "            (fc1): Linear(in_features=1024, out_features=4096, bias=True)\n",
      "            (act): GELU(approximate='none')\n",
      "            (drop1): Dropout(p=0.0, inplace=False)\n",
      "            (fc2): Linear(in_features=4096, out_features=1024, bias=True)\n",
      "            (drop2): Dropout(p=0.0, inplace=False)\n",
      "          )\n",
      "          (ls2): Identity()\n",
      "          (drop_path2): Identity()\n",
      "        )\n",
      "        (6): Block(\n",
      "          (norm1): LayerNorm((1024,), eps=1e-06, elementwise_affine=True)\n",
      "          (attn): Attention(\n",
      "            (qkv): Linear(in_features=1024, out_features=3072, bias=True)\n",
      "            (attn_drop): Dropout(p=0.0, inplace=False)\n",
      "            (proj): Linear(in_features=1024, out_features=1024, bias=True)\n",
      "            (proj_drop): Dropout(p=0.0, inplace=False)\n",
      "          )\n",
      "          (ls1): Identity()\n",
      "          (drop_path1): Identity()\n",
      "          (norm2): LayerNorm((1024,), eps=1e-06, elementwise_affine=True)\n",
      "          (mlp): Mlp(\n",
      "            (fc1): Linear(in_features=1024, out_features=4096, bias=True)\n",
      "            (act): GELU(approximate='none')\n",
      "            (drop1): Dropout(p=0.0, inplace=False)\n",
      "            (fc2): Linear(in_features=4096, out_features=1024, bias=True)\n",
      "            (drop2): Dropout(p=0.0, inplace=False)\n",
      "          )\n",
      "          (ls2): Identity()\n",
      "          (drop_path2): Identity()\n",
      "        )\n",
      "        (7): Block(\n",
      "          (norm1): LayerNorm((1024,), eps=1e-06, elementwise_affine=True)\n",
      "          (attn): Attention(\n",
      "            (qkv): Linear(in_features=1024, out_features=3072, bias=True)\n",
      "            (attn_drop): Dropout(p=0.0, inplace=False)\n",
      "            (proj): Linear(in_features=1024, out_features=1024, bias=True)\n",
      "            (proj_drop): Dropout(p=0.0, inplace=False)\n",
      "          )\n",
      "          (ls1): Identity()\n",
      "          (drop_path1): Identity()\n",
      "          (norm2): LayerNorm((1024,), eps=1e-06, elementwise_affine=True)\n",
      "          (mlp): Mlp(\n",
      "            (fc1): Linear(in_features=1024, out_features=4096, bias=True)\n",
      "            (act): GELU(approximate='none')\n",
      "            (drop1): Dropout(p=0.0, inplace=False)\n",
      "            (fc2): Linear(in_features=4096, out_features=1024, bias=True)\n",
      "            (drop2): Dropout(p=0.0, inplace=False)\n",
      "          )\n",
      "          (ls2): Identity()\n",
      "          (drop_path2): Identity()\n",
      "        )\n",
      "        (8): Block(\n",
      "          (norm1): LayerNorm((1024,), eps=1e-06, elementwise_affine=True)\n",
      "          (attn): Attention(\n",
      "            (qkv): Linear(in_features=1024, out_features=3072, bias=True)\n",
      "            (attn_drop): Dropout(p=0.0, inplace=False)\n",
      "            (proj): Linear(in_features=1024, out_features=1024, bias=True)\n",
      "            (proj_drop): Dropout(p=0.0, inplace=False)\n",
      "          )\n",
      "          (ls1): Identity()\n",
      "          (drop_path1): Identity()\n",
      "          (norm2): LayerNorm((1024,), eps=1e-06, elementwise_affine=True)\n",
      "          (mlp): Mlp(\n",
      "            (fc1): Linear(in_features=1024, out_features=4096, bias=True)\n",
      "            (act): GELU(approximate='none')\n",
      "            (drop1): Dropout(p=0.0, inplace=False)\n",
      "            (fc2): Linear(in_features=4096, out_features=1024, bias=True)\n",
      "            (drop2): Dropout(p=0.0, inplace=False)\n",
      "          )\n",
      "          (ls2): Identity()\n",
      "          (drop_path2): Identity()\n",
      "        )\n",
      "        (9): Block(\n",
      "          (norm1): LayerNorm((1024,), eps=1e-06, elementwise_affine=True)\n",
      "          (attn): Attention(\n",
      "            (qkv): Linear(in_features=1024, out_features=3072, bias=True)\n",
      "            (attn_drop): Dropout(p=0.0, inplace=False)\n",
      "            (proj): Linear(in_features=1024, out_features=1024, bias=True)\n",
      "            (proj_drop): Dropout(p=0.0, inplace=False)\n",
      "          )\n",
      "          (ls1): Identity()\n",
      "          (drop_path1): Identity()\n",
      "          (norm2): LayerNorm((1024,), eps=1e-06, elementwise_affine=True)\n",
      "          (mlp): Mlp(\n",
      "            (fc1): Linear(in_features=1024, out_features=4096, bias=True)\n",
      "            (act): GELU(approximate='none')\n",
      "            (drop1): Dropout(p=0.0, inplace=False)\n",
      "            (fc2): Linear(in_features=4096, out_features=1024, bias=True)\n",
      "            (drop2): Dropout(p=0.0, inplace=False)\n",
      "          )\n",
      "          (ls2): Identity()\n",
      "          (drop_path2): Identity()\n",
      "        )\n",
      "        (10): Block(\n",
      "          (norm1): LayerNorm((1024,), eps=1e-06, elementwise_affine=True)\n",
      "          (attn): Attention(\n",
      "            (qkv): Linear(in_features=1024, out_features=3072, bias=True)\n",
      "            (attn_drop): Dropout(p=0.0, inplace=False)\n",
      "            (proj): Linear(in_features=1024, out_features=1024, bias=True)\n",
      "            (proj_drop): Dropout(p=0.0, inplace=False)\n",
      "          )\n",
      "          (ls1): Identity()\n",
      "          (drop_path1): Identity()\n",
      "          (norm2): LayerNorm((1024,), eps=1e-06, elementwise_affine=True)\n",
      "          (mlp): Mlp(\n",
      "            (fc1): Linear(in_features=1024, out_features=4096, bias=True)\n",
      "            (act): GELU(approximate='none')\n",
      "            (drop1): Dropout(p=0.0, inplace=False)\n",
      "            (fc2): Linear(in_features=4096, out_features=1024, bias=True)\n",
      "            (drop2): Dropout(p=0.0, inplace=False)\n",
      "          )\n",
      "          (ls2): Identity()\n",
      "          (drop_path2): Identity()\n",
      "        )\n",
      "        (11): Block(\n",
      "          (norm1): LayerNorm((1024,), eps=1e-06, elementwise_affine=True)\n",
      "          (attn): Attention(\n",
      "            (qkv): Linear(in_features=1024, out_features=3072, bias=True)\n",
      "            (attn_drop): Dropout(p=0.0, inplace=False)\n",
      "            (proj): Linear(in_features=1024, out_features=1024, bias=True)\n",
      "            (proj_drop): Dropout(p=0.0, inplace=False)\n",
      "          )\n",
      "          (ls1): Identity()\n",
      "          (drop_path1): Identity()\n",
      "          (norm2): LayerNorm((1024,), eps=1e-06, elementwise_affine=True)\n",
      "          (mlp): Mlp(\n",
      "            (fc1): Linear(in_features=1024, out_features=4096, bias=True)\n",
      "            (act): GELU(approximate='none')\n",
      "            (drop1): Dropout(p=0.0, inplace=False)\n",
      "            (fc2): Linear(in_features=4096, out_features=1024, bias=True)\n",
      "            (drop2): Dropout(p=0.0, inplace=False)\n",
      "          )\n",
      "          (ls2): Identity()\n",
      "          (drop_path2): Identity()\n",
      "        )\n",
      "        (12): Block(\n",
      "          (norm1): LayerNorm((1024,), eps=1e-06, elementwise_affine=True)\n",
      "          (attn): Attention(\n",
      "            (qkv): Linear(in_features=1024, out_features=3072, bias=True)\n",
      "            (attn_drop): Dropout(p=0.0, inplace=False)\n",
      "            (proj): Linear(in_features=1024, out_features=1024, bias=True)\n",
      "            (proj_drop): Dropout(p=0.0, inplace=False)\n",
      "          )\n",
      "          (ls1): Identity()\n",
      "          (drop_path1): Identity()\n",
      "          (norm2): LayerNorm((1024,), eps=1e-06, elementwise_affine=True)\n",
      "          (mlp): Mlp(\n",
      "            (fc1): Linear(in_features=1024, out_features=4096, bias=True)\n",
      "            (act): GELU(approximate='none')\n",
      "            (drop1): Dropout(p=0.0, inplace=False)\n",
      "            (fc2): Linear(in_features=4096, out_features=1024, bias=True)\n",
      "            (drop2): Dropout(p=0.0, inplace=False)\n",
      "          )\n",
      "          (ls2): Identity()\n",
      "          (drop_path2): Identity()\n",
      "        )\n",
      "        (13): Block(\n",
      "          (norm1): LayerNorm((1024,), eps=1e-06, elementwise_affine=True)\n",
      "          (attn): Attention(\n",
      "            (qkv): Linear(in_features=1024, out_features=3072, bias=True)\n",
      "            (attn_drop): Dropout(p=0.0, inplace=False)\n",
      "            (proj): Linear(in_features=1024, out_features=1024, bias=True)\n",
      "            (proj_drop): Dropout(p=0.0, inplace=False)\n",
      "          )\n",
      "          (ls1): Identity()\n",
      "          (drop_path1): Identity()\n",
      "          (norm2): LayerNorm((1024,), eps=1e-06, elementwise_affine=True)\n",
      "          (mlp): Mlp(\n",
      "            (fc1): Linear(in_features=1024, out_features=4096, bias=True)\n",
      "            (act): GELU(approximate='none')\n",
      "            (drop1): Dropout(p=0.0, inplace=False)\n",
      "            (fc2): Linear(in_features=4096, out_features=1024, bias=True)\n",
      "            (drop2): Dropout(p=0.0, inplace=False)\n",
      "          )\n",
      "          (ls2): Identity()\n",
      "          (drop_path2): Identity()\n",
      "        )\n",
      "        (14): Block(\n",
      "          (norm1): LayerNorm((1024,), eps=1e-06, elementwise_affine=True)\n",
      "          (attn): Attention(\n",
      "            (qkv): Linear(in_features=1024, out_features=3072, bias=True)\n",
      "            (attn_drop): Dropout(p=0.0, inplace=False)\n",
      "            (proj): Linear(in_features=1024, out_features=1024, bias=True)\n",
      "            (proj_drop): Dropout(p=0.0, inplace=False)\n",
      "          )\n",
      "          (ls1): Identity()\n",
      "          (drop_path1): Identity()\n",
      "          (norm2): LayerNorm((1024,), eps=1e-06, elementwise_affine=True)\n",
      "          (mlp): Mlp(\n",
      "            (fc1): Linear(in_features=1024, out_features=4096, bias=True)\n",
      "            (act): GELU(approximate='none')\n",
      "            (drop1): Dropout(p=0.0, inplace=False)\n",
      "            (fc2): Linear(in_features=4096, out_features=1024, bias=True)\n",
      "            (drop2): Dropout(p=0.0, inplace=False)\n",
      "          )\n",
      "          (ls2): Identity()\n",
      "          (drop_path2): Identity()\n",
      "        )\n",
      "        (15): Block(\n",
      "          (norm1): LayerNorm((1024,), eps=1e-06, elementwise_affine=True)\n",
      "          (attn): Attention(\n",
      "            (qkv): Linear(in_features=1024, out_features=3072, bias=True)\n",
      "            (attn_drop): Dropout(p=0.0, inplace=False)\n",
      "            (proj): Linear(in_features=1024, out_features=1024, bias=True)\n",
      "            (proj_drop): Dropout(p=0.0, inplace=False)\n",
      "          )\n",
      "          (ls1): Identity()\n",
      "          (drop_path1): Identity()\n",
      "          (norm2): LayerNorm((1024,), eps=1e-06, elementwise_affine=True)\n",
      "          (mlp): Mlp(\n",
      "            (fc1): Linear(in_features=1024, out_features=4096, bias=True)\n",
      "            (act): GELU(approximate='none')\n",
      "            (drop1): Dropout(p=0.0, inplace=False)\n",
      "            (fc2): Linear(in_features=4096, out_features=1024, bias=True)\n",
      "            (drop2): Dropout(p=0.0, inplace=False)\n",
      "          )\n",
      "          (ls2): Identity()\n",
      "          (drop_path2): Identity()\n",
      "        )\n",
      "        (16): Block(\n",
      "          (norm1): LayerNorm((1024,), eps=1e-06, elementwise_affine=True)\n",
      "          (attn): Attention(\n",
      "            (qkv): Linear(in_features=1024, out_features=3072, bias=True)\n",
      "            (attn_drop): Dropout(p=0.0, inplace=False)\n",
      "            (proj): Linear(in_features=1024, out_features=1024, bias=True)\n",
      "            (proj_drop): Dropout(p=0.0, inplace=False)\n",
      "          )\n",
      "          (ls1): Identity()\n",
      "          (drop_path1): Identity()\n",
      "          (norm2): LayerNorm((1024,), eps=1e-06, elementwise_affine=True)\n",
      "          (mlp): Mlp(\n",
      "            (fc1): Linear(in_features=1024, out_features=4096, bias=True)\n",
      "            (act): GELU(approximate='none')\n",
      "            (drop1): Dropout(p=0.0, inplace=False)\n",
      "            (fc2): Linear(in_features=4096, out_features=1024, bias=True)\n",
      "            (drop2): Dropout(p=0.0, inplace=False)\n",
      "          )\n",
      "          (ls2): Identity()\n",
      "          (drop_path2): Identity()\n",
      "        )\n",
      "        (17): Block(\n",
      "          (norm1): LayerNorm((1024,), eps=1e-06, elementwise_affine=True)\n",
      "          (attn): Attention(\n",
      "            (qkv): Linear(in_features=1024, out_features=3072, bias=True)\n",
      "            (attn_drop): Dropout(p=0.0, inplace=False)\n",
      "            (proj): Linear(in_features=1024, out_features=1024, bias=True)\n",
      "            (proj_drop): Dropout(p=0.0, inplace=False)\n",
      "          )\n",
      "          (ls1): Identity()\n",
      "          (drop_path1): Identity()\n",
      "          (norm2): LayerNorm((1024,), eps=1e-06, elementwise_affine=True)\n",
      "          (mlp): Mlp(\n",
      "            (fc1): Linear(in_features=1024, out_features=4096, bias=True)\n",
      "            (act): GELU(approximate='none')\n",
      "            (drop1): Dropout(p=0.0, inplace=False)\n",
      "            (fc2): Linear(in_features=4096, out_features=1024, bias=True)\n",
      "            (drop2): Dropout(p=0.0, inplace=False)\n",
      "          )\n",
      "          (ls2): Identity()\n",
      "          (drop_path2): Identity()\n",
      "        )\n",
      "        (18): Block(\n",
      "          (norm1): LayerNorm((1024,), eps=1e-06, elementwise_affine=True)\n",
      "          (attn): Attention(\n",
      "            (qkv): Linear(in_features=1024, out_features=3072, bias=True)\n",
      "            (attn_drop): Dropout(p=0.0, inplace=False)\n",
      "            (proj): Linear(in_features=1024, out_features=1024, bias=True)\n",
      "            (proj_drop): Dropout(p=0.0, inplace=False)\n",
      "          )\n",
      "          (ls1): Identity()\n",
      "          (drop_path1): Identity()\n",
      "          (norm2): LayerNorm((1024,), eps=1e-06, elementwise_affine=True)\n",
      "          (mlp): Mlp(\n",
      "            (fc1): Linear(in_features=1024, out_features=4096, bias=True)\n",
      "            (act): GELU(approximate='none')\n",
      "            (drop1): Dropout(p=0.0, inplace=False)\n",
      "            (fc2): Linear(in_features=4096, out_features=1024, bias=True)\n",
      "            (drop2): Dropout(p=0.0, inplace=False)\n",
      "          )\n",
      "          (ls2): Identity()\n",
      "          (drop_path2): Identity()\n",
      "        )\n",
      "        (19): Block(\n",
      "          (norm1): LayerNorm((1024,), eps=1e-06, elementwise_affine=True)\n",
      "          (attn): Attention(\n",
      "            (qkv): Linear(in_features=1024, out_features=3072, bias=True)\n",
      "            (attn_drop): Dropout(p=0.0, inplace=False)\n",
      "            (proj): Linear(in_features=1024, out_features=1024, bias=True)\n",
      "            (proj_drop): Dropout(p=0.0, inplace=False)\n",
      "          )\n",
      "          (ls1): Identity()\n",
      "          (drop_path1): Identity()\n",
      "          (norm2): LayerNorm((1024,), eps=1e-06, elementwise_affine=True)\n",
      "          (mlp): Mlp(\n",
      "            (fc1): Linear(in_features=1024, out_features=4096, bias=True)\n",
      "            (act): GELU(approximate='none')\n",
      "            (drop1): Dropout(p=0.0, inplace=False)\n",
      "            (fc2): Linear(in_features=4096, out_features=1024, bias=True)\n",
      "            (drop2): Dropout(p=0.0, inplace=False)\n",
      "          )\n",
      "          (ls2): Identity()\n",
      "          (drop_path2): Identity()\n",
      "        )\n",
      "        (20): Block(\n",
      "          (norm1): LayerNorm((1024,), eps=1e-06, elementwise_affine=True)\n",
      "          (attn): Attention(\n",
      "            (qkv): Linear(in_features=1024, out_features=3072, bias=True)\n",
      "            (attn_drop): Dropout(p=0.0, inplace=False)\n",
      "            (proj): Linear(in_features=1024, out_features=1024, bias=True)\n",
      "            (proj_drop): Dropout(p=0.0, inplace=False)\n",
      "          )\n",
      "          (ls1): Identity()\n",
      "          (drop_path1): Identity()\n",
      "          (norm2): LayerNorm((1024,), eps=1e-06, elementwise_affine=True)\n",
      "          (mlp): Mlp(\n",
      "            (fc1): Linear(in_features=1024, out_features=4096, bias=True)\n",
      "            (act): GELU(approximate='none')\n",
      "            (drop1): Dropout(p=0.0, inplace=False)\n",
      "            (fc2): Linear(in_features=4096, out_features=1024, bias=True)\n",
      "            (drop2): Dropout(p=0.0, inplace=False)\n",
      "          )\n",
      "          (ls2): Identity()\n",
      "          (drop_path2): Identity()\n",
      "        )\n",
      "        (21): Block(\n",
      "          (norm1): LayerNorm((1024,), eps=1e-06, elementwise_affine=True)\n",
      "          (attn): Attention(\n",
      "            (qkv): Linear(in_features=1024, out_features=3072, bias=True)\n",
      "            (attn_drop): Dropout(p=0.0, inplace=False)\n",
      "            (proj): Linear(in_features=1024, out_features=1024, bias=True)\n",
      "            (proj_drop): Dropout(p=0.0, inplace=False)\n",
      "          )\n",
      "          (ls1): Identity()\n",
      "          (drop_path1): Identity()\n",
      "          (norm2): LayerNorm((1024,), eps=1e-06, elementwise_affine=True)\n",
      "          (mlp): Mlp(\n",
      "            (fc1): Linear(in_features=1024, out_features=4096, bias=True)\n",
      "            (act): GELU(approximate='none')\n",
      "            (drop1): Dropout(p=0.0, inplace=False)\n",
      "            (fc2): Linear(in_features=4096, out_features=1024, bias=True)\n",
      "            (drop2): Dropout(p=0.0, inplace=False)\n",
      "          )\n",
      "          (ls2): Identity()\n",
      "          (drop_path2): Identity()\n",
      "        )\n",
      "        (22): Block(\n",
      "          (norm1): LayerNorm((1024,), eps=1e-06, elementwise_affine=True)\n",
      "          (attn): Attention(\n",
      "            (qkv): Linear(in_features=1024, out_features=3072, bias=True)\n",
      "            (attn_drop): Dropout(p=0.0, inplace=False)\n",
      "            (proj): Linear(in_features=1024, out_features=1024, bias=True)\n",
      "            (proj_drop): Dropout(p=0.0, inplace=False)\n",
      "          )\n",
      "          (ls1): Identity()\n",
      "          (drop_path1): Identity()\n",
      "          (norm2): LayerNorm((1024,), eps=1e-06, elementwise_affine=True)\n",
      "          (mlp): Mlp(\n",
      "            (fc1): Linear(in_features=1024, out_features=4096, bias=True)\n",
      "            (act): GELU(approximate='none')\n",
      "            (drop1): Dropout(p=0.0, inplace=False)\n",
      "            (fc2): Linear(in_features=4096, out_features=1024, bias=True)\n",
      "            (drop2): Dropout(p=0.0, inplace=False)\n",
      "          )\n",
      "          (ls2): Identity()\n",
      "          (drop_path2): Identity()\n",
      "        )\n",
      "        (23): Block(\n",
      "          (norm1): LayerNorm((1024,), eps=1e-06, elementwise_affine=True)\n",
      "          (attn): Attention(\n",
      "            (qkv): Linear(in_features=1024, out_features=3072, bias=True)\n",
      "            (attn_drop): Dropout(p=0.0, inplace=False)\n",
      "            (proj): Linear(in_features=1024, out_features=1024, bias=True)\n",
      "            (proj_drop): Dropout(p=0.0, inplace=False)\n",
      "          )\n",
      "          (ls1): Identity()\n",
      "          (drop_path1): Identity()\n",
      "          (norm2): LayerNorm((1024,), eps=1e-06, elementwise_affine=True)\n",
      "          (mlp): Mlp(\n",
      "            (fc1): Linear(in_features=1024, out_features=4096, bias=True)\n",
      "            (act): GELU(approximate='none')\n",
      "            (drop1): Dropout(p=0.0, inplace=False)\n",
      "            (fc2): Linear(in_features=4096, out_features=1024, bias=True)\n",
      "            (drop2): Dropout(p=0.0, inplace=False)\n",
      "          )\n",
      "          (ls2): Identity()\n",
      "          (drop_path2): Identity()\n",
      "        )\n",
      "      )\n",
      "      (norm): LayerNorm((1024,), eps=1e-06, elementwise_affine=True)\n",
      "      (fc_norm): Identity()\n",
      "    )\n",
      "  )\n",
      "  (dynamics): LSTM(\n",
      "    (lstm): LSTM(1024, 1024)\n",
      "    (regressor): Linear(in_features=1024, out_features=1024, bias=True)\n",
      "  )\n",
      ")\n",
      "======= Model layers =======\n",
      "['encoder', 'dynamics']\n"
     ]
    }
   ],
   "source": [
    "name = \"pfVC1_LSTM_k700\"\n",
    "model, model_layers = load_pretrained_model(name)\n",
    "print(\"======= Model architecture =======\\n\", model)\n",
    "print(f\"======= Model layers =======\\n{model_layers}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "20e9358d",
   "metadata": {},
   "source": [
    "# 2. R3M + Dynamics"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ec16153a",
   "metadata": {},
   "source": [
    "## 2a. R3M + CTRNN, dynamics module pretrained on Physion"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "43d0593d",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/om2/user/anayebi/miniconda/envs/pong_env/lib/python3.9/site-packages/torch/nn/_reduction.py:42: UserWarning: size_average and reduce args will be deprecated, please use reduction='none' instead.\n",
      "  warnings.warn(warning.format(ret))\n",
      "/om2/user/anayebi/miniconda/envs/pong_env/lib/python3.9/site-packages/torchvision/models/_utils.py:208: UserWarning: The parameter 'pretrained' is deprecated since 0.13 and may be removed in the future, please use 'weights' instead.\n",
      "  warnings.warn(\n",
      "/om2/user/anayebi/miniconda/envs/pong_env/lib/python3.9/site-packages/torchvision/models/_utils.py:223: UserWarning: Arguments other than a weight enum or `None` for 'weights' are deprecated since 0.13 and may be removed in the future. The current behavior is equivalent to passing `weights=None`.\n",
      "  warnings.warn(msg)\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Loaded parameters from /om/weka/yanglab/anayebi/mental-sim/trained_models/R3M+CTRNN_physion.pt\n",
      "======= Model architecture =======\n",
      " FrozenPretrainedEncoder(\n",
      "  (encoder): R3M_pretrained(\n",
      "    (r3m): DataParallel(\n",
      "      (module): R3M(\n",
      "        (cs): CosineSimilarity()\n",
      "        (bce): BCELoss()\n",
      "        (sigm): Sigmoid()\n",
      "        (convnet): ResNet(\n",
      "          (conv1): Conv2d(3, 64, kernel_size=(7, 7), stride=(2, 2), padding=(3, 3), bias=False)\n",
      "          (bn1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "          (relu): ReLU(inplace=True)\n",
      "          (maxpool): MaxPool2d(kernel_size=3, stride=2, padding=1, dilation=1, ceil_mode=False)\n",
      "          (layer1): Sequential(\n",
      "            (0): Bottleneck(\n",
      "              (conv1): Conv2d(64, 64, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "              (bn1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "              (conv2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
      "              (bn2): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "              (conv3): Conv2d(64, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "              (bn3): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "              (relu): ReLU(inplace=True)\n",
      "              (downsample): Sequential(\n",
      "                (0): Conv2d(64, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "                (1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "              )\n",
      "            )\n",
      "            (1): Bottleneck(\n",
      "              (conv1): Conv2d(256, 64, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "              (bn1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "              (conv2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
      "              (bn2): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "              (conv3): Conv2d(64, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "              (bn3): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "              (relu): ReLU(inplace=True)\n",
      "            )\n",
      "            (2): Bottleneck(\n",
      "              (conv1): Conv2d(256, 64, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "              (bn1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "              (conv2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
      "              (bn2): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "              (conv3): Conv2d(64, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "              (bn3): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "              (relu): ReLU(inplace=True)\n",
      "            )\n",
      "          )\n",
      "          (layer2): Sequential(\n",
      "            (0): Bottleneck(\n",
      "              (conv1): Conv2d(256, 128, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "              (bn1): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "              (conv2): Conv2d(128, 128, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)\n",
      "              (bn2): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "              (conv3): Conv2d(128, 512, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "              (bn3): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "              (relu): ReLU(inplace=True)\n",
      "              (downsample): Sequential(\n",
      "                (0): Conv2d(256, 512, kernel_size=(1, 1), stride=(2, 2), bias=False)\n",
      "                (1): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "              )\n",
      "            )\n",
      "            (1): Bottleneck(\n",
      "              (conv1): Conv2d(512, 128, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "              (bn1): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "              (conv2): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
      "              (bn2): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "              (conv3): Conv2d(128, 512, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "              (bn3): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "              (relu): ReLU(inplace=True)\n",
      "            )\n",
      "            (2): Bottleneck(\n",
      "              (conv1): Conv2d(512, 128, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "              (bn1): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "              (conv2): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
      "              (bn2): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "              (conv3): Conv2d(128, 512, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "              (bn3): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "              (relu): ReLU(inplace=True)\n",
      "            )\n",
      "            (3): Bottleneck(\n",
      "              (conv1): Conv2d(512, 128, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "              (bn1): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "              (conv2): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
      "              (bn2): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "              (conv3): Conv2d(128, 512, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "              (bn3): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "              (relu): ReLU(inplace=True)\n",
      "            )\n",
      "          )\n",
      "          (layer3): Sequential(\n",
      "            (0): Bottleneck(\n",
      "              (conv1): Conv2d(512, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "              (bn1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "              (conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)\n",
      "              (bn2): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "              (conv3): Conv2d(256, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "              (bn3): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "              (relu): ReLU(inplace=True)\n",
      "              (downsample): Sequential(\n",
      "                (0): Conv2d(512, 1024, kernel_size=(1, 1), stride=(2, 2), bias=False)\n",
      "                (1): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "              )\n",
      "            )\n",
      "            (1): Bottleneck(\n",
      "              (conv1): Conv2d(1024, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "              (bn1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "              (conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
      "              (bn2): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "              (conv3): Conv2d(256, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "              (bn3): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "              (relu): ReLU(inplace=True)\n",
      "            )\n",
      "            (2): Bottleneck(\n",
      "              (conv1): Conv2d(1024, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "              (bn1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "              (conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
      "              (bn2): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "              (conv3): Conv2d(256, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "              (bn3): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "              (relu): ReLU(inplace=True)\n",
      "            )\n",
      "            (3): Bottleneck(\n",
      "              (conv1): Conv2d(1024, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "              (bn1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "              (conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
      "              (bn2): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "              (conv3): Conv2d(256, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "              (bn3): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "              (relu): ReLU(inplace=True)\n",
      "            )\n",
      "            (4): Bottleneck(\n",
      "              (conv1): Conv2d(1024, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "              (bn1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "              (conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
      "              (bn2): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "              (conv3): Conv2d(256, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "              (bn3): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "              (relu): ReLU(inplace=True)\n",
      "            )\n",
      "            (5): Bottleneck(\n",
      "              (conv1): Conv2d(1024, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "              (bn1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "              (conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
      "              (bn2): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "              (conv3): Conv2d(256, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "              (bn3): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "              (relu): ReLU(inplace=True)\n",
      "            )\n",
      "          )\n",
      "          (layer4): Sequential(\n",
      "            (0): Bottleneck(\n",
      "              (conv1): Conv2d(1024, 512, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "              (bn1): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "              (conv2): Conv2d(512, 512, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)\n",
      "              (bn2): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "              (conv3): Conv2d(512, 2048, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "              (bn3): BatchNorm2d(2048, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "              (relu): ReLU(inplace=True)\n",
      "              (downsample): Sequential(\n",
      "                (0): Conv2d(1024, 2048, kernel_size=(1, 1), stride=(2, 2), bias=False)\n",
      "                (1): BatchNorm2d(2048, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "              )\n",
      "            )\n",
      "            (1): Bottleneck(\n",
      "              (conv1): Conv2d(2048, 512, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "              (bn1): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "              (conv2): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
      "              (bn2): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "              (conv3): Conv2d(512, 2048, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "              (bn3): BatchNorm2d(2048, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "              (relu): ReLU(inplace=True)\n",
      "            )\n",
      "            (2): Bottleneck(\n",
      "              (conv1): Conv2d(2048, 512, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "              (bn1): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "              (conv2): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
      "              (bn2): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "              (conv3): Conv2d(512, 2048, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "              (bn3): BatchNorm2d(2048, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "              (relu): ReLU(inplace=True)\n",
      "            )\n",
      "          )\n",
      "          (avgpool): AdaptiveAvgPool2d(output_size=(1, 1))\n",
      "          (fc): Identity()\n",
      "        )\n",
      "        (normlayer): Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])\n",
      "      )\n",
      "    )\n",
      "  )\n",
      "  (dynamics): CTRNN(\n",
      "    (ctrnn): CTRNN(\n",
      "      (fc_x2ah): Linear(in_features=2048, out_features=1024, bias=True)\n",
      "      (fc_h2ah): Linear(in_features=1024, out_features=1024, bias=False)\n",
      "      (fc_h2y): Linear(in_features=1024, out_features=1024, bias=True)\n",
      "    )\n",
      "    (regressor): Linear(in_features=1024, out_features=2048, bias=True)\n",
      "  )\n",
      ")\n",
      "======= Model layers =======\n",
      "['encoder', 'dynamics']\n"
     ]
    }
   ],
   "source": [
    "name = \"pfR3M_CTRNN_physion\"\n",
    "model, model_layers = load_pretrained_model(name)\n",
    "print(\"======= Model architecture =======\\n\", model)\n",
    "print(f\"======= Model layers =======\\n{model_layers}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f604a190",
   "metadata": {},
   "source": [
    "## 2b. R3M + CTRNN, dynamics module pretrained on Kinetics-700"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "7eb19a12",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Loaded parameters from /om/weka/yanglab/anayebi/mental-sim/trained_models/R3M+CTRNN_k700.pt\n",
      "======= Model architecture =======\n",
      " FrozenPretrainedEncoder(\n",
      "  (encoder): R3M_pretrained(\n",
      "    (r3m): DataParallel(\n",
      "      (module): R3M(\n",
      "        (cs): CosineSimilarity()\n",
      "        (bce): BCELoss()\n",
      "        (sigm): Sigmoid()\n",
      "        (convnet): ResNet(\n",
      "          (conv1): Conv2d(3, 64, kernel_size=(7, 7), stride=(2, 2), padding=(3, 3), bias=False)\n",
      "          (bn1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "          (relu): ReLU(inplace=True)\n",
      "          (maxpool): MaxPool2d(kernel_size=3, stride=2, padding=1, dilation=1, ceil_mode=False)\n",
      "          (layer1): Sequential(\n",
      "            (0): Bottleneck(\n",
      "              (conv1): Conv2d(64, 64, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "              (bn1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "              (conv2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
      "              (bn2): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "              (conv3): Conv2d(64, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "              (bn3): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "              (relu): ReLU(inplace=True)\n",
      "              (downsample): Sequential(\n",
      "                (0): Conv2d(64, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "                (1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "              )\n",
      "            )\n",
      "            (1): Bottleneck(\n",
      "              (conv1): Conv2d(256, 64, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "              (bn1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "              (conv2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
      "              (bn2): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "              (conv3): Conv2d(64, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "              (bn3): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "              (relu): ReLU(inplace=True)\n",
      "            )\n",
      "            (2): Bottleneck(\n",
      "              (conv1): Conv2d(256, 64, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "              (bn1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "              (conv2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
      "              (bn2): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "              (conv3): Conv2d(64, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "              (bn3): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "              (relu): ReLU(inplace=True)\n",
      "            )\n",
      "          )\n",
      "          (layer2): Sequential(\n",
      "            (0): Bottleneck(\n",
      "              (conv1): Conv2d(256, 128, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "              (bn1): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "              (conv2): Conv2d(128, 128, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)\n",
      "              (bn2): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "              (conv3): Conv2d(128, 512, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "              (bn3): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "              (relu): ReLU(inplace=True)\n",
      "              (downsample): Sequential(\n",
      "                (0): Conv2d(256, 512, kernel_size=(1, 1), stride=(2, 2), bias=False)\n",
      "                (1): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "              )\n",
      "            )\n",
      "            (1): Bottleneck(\n",
      "              (conv1): Conv2d(512, 128, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "              (bn1): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "              (conv2): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
      "              (bn2): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "              (conv3): Conv2d(128, 512, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "              (bn3): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "              (relu): ReLU(inplace=True)\n",
      "            )\n",
      "            (2): Bottleneck(\n",
      "              (conv1): Conv2d(512, 128, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "              (bn1): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "              (conv2): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
      "              (bn2): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "              (conv3): Conv2d(128, 512, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "              (bn3): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "              (relu): ReLU(inplace=True)\n",
      "            )\n",
      "            (3): Bottleneck(\n",
      "              (conv1): Conv2d(512, 128, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "              (bn1): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "              (conv2): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
      "              (bn2): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "              (conv3): Conv2d(128, 512, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "              (bn3): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "              (relu): ReLU(inplace=True)\n",
      "            )\n",
      "          )\n",
      "          (layer3): Sequential(\n",
      "            (0): Bottleneck(\n",
      "              (conv1): Conv2d(512, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "              (bn1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "              (conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)\n",
      "              (bn2): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "              (conv3): Conv2d(256, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "              (bn3): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "              (relu): ReLU(inplace=True)\n",
      "              (downsample): Sequential(\n",
      "                (0): Conv2d(512, 1024, kernel_size=(1, 1), stride=(2, 2), bias=False)\n",
      "                (1): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "              )\n",
      "            )\n",
      "            (1): Bottleneck(\n",
      "              (conv1): Conv2d(1024, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "              (bn1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "              (conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
      "              (bn2): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "              (conv3): Conv2d(256, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "              (bn3): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "              (relu): ReLU(inplace=True)\n",
      "            )\n",
      "            (2): Bottleneck(\n",
      "              (conv1): Conv2d(1024, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "              (bn1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "              (conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
      "              (bn2): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "              (conv3): Conv2d(256, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "              (bn3): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "              (relu): ReLU(inplace=True)\n",
      "            )\n",
      "            (3): Bottleneck(\n",
      "              (conv1): Conv2d(1024, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "              (bn1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "              (conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
      "              (bn2): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "              (conv3): Conv2d(256, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "              (bn3): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "              (relu): ReLU(inplace=True)\n",
      "            )\n",
      "            (4): Bottleneck(\n",
      "              (conv1): Conv2d(1024, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "              (bn1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "              (conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
      "              (bn2): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "              (conv3): Conv2d(256, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "              (bn3): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "              (relu): ReLU(inplace=True)\n",
      "            )\n",
      "            (5): Bottleneck(\n",
      "              (conv1): Conv2d(1024, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "              (bn1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "              (conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
      "              (bn2): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "              (conv3): Conv2d(256, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "              (bn3): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "              (relu): ReLU(inplace=True)\n",
      "            )\n",
      "          )\n",
      "          (layer4): Sequential(\n",
      "            (0): Bottleneck(\n",
      "              (conv1): Conv2d(1024, 512, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "              (bn1): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "              (conv2): Conv2d(512, 512, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)\n",
      "              (bn2): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "              (conv3): Conv2d(512, 2048, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "              (bn3): BatchNorm2d(2048, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "              (relu): ReLU(inplace=True)\n",
      "              (downsample): Sequential(\n",
      "                (0): Conv2d(1024, 2048, kernel_size=(1, 1), stride=(2, 2), bias=False)\n",
      "                (1): BatchNorm2d(2048, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "              )\n",
      "            )\n",
      "            (1): Bottleneck(\n",
      "              (conv1): Conv2d(2048, 512, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "              (bn1): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "              (conv2): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
      "              (bn2): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "              (conv3): Conv2d(512, 2048, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "              (bn3): BatchNorm2d(2048, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "              (relu): ReLU(inplace=True)\n",
      "            )\n",
      "            (2): Bottleneck(\n",
      "              (conv1): Conv2d(2048, 512, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "              (bn1): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "              (conv2): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
      "              (bn2): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "              (conv3): Conv2d(512, 2048, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "              (bn3): BatchNorm2d(2048, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "              (relu): ReLU(inplace=True)\n",
      "            )\n",
      "          )\n",
      "          (avgpool): AdaptiveAvgPool2d(output_size=(1, 1))\n",
      "          (fc): Identity()\n",
      "        )\n",
      "        (normlayer): Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])\n",
      "      )\n",
      "    )\n",
      "  )\n",
      "  (dynamics): CTRNN(\n",
      "    (ctrnn): CTRNN(\n",
      "      (fc_x2ah): Linear(in_features=2048, out_features=1024, bias=True)\n",
      "      (fc_h2ah): Linear(in_features=1024, out_features=1024, bias=False)\n",
      "      (fc_h2y): Linear(in_features=1024, out_features=1024, bias=True)\n",
      "    )\n",
      "    (regressor): Linear(in_features=1024, out_features=2048, bias=True)\n",
      "  )\n",
      ")\n",
      "======= Model layers =======\n",
      "['encoder', 'dynamics']\n"
     ]
    }
   ],
   "source": [
    "name = \"pfR3M_CTRNN_k700\"\n",
    "model, model_layers = load_pretrained_model(name)\n",
    "print(\"======= Model architecture =======\\n\", model)\n",
    "print(f\"======= Model layers =======\\n{model_layers}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f093dd66",
   "metadata": {},
   "source": [
    "## 2c. R3M + LSTM, dynamics module pretrained on Physion"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "5db8760b",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Loaded parameters from /om/weka/yanglab/anayebi/mental-sim/trained_models/R3M+LSTM_physion.pt\n",
      "======= Model architecture =======\n",
      " FrozenPretrainedEncoder(\n",
      "  (encoder): R3M_pretrained(\n",
      "    (r3m): DataParallel(\n",
      "      (module): R3M(\n",
      "        (cs): CosineSimilarity()\n",
      "        (bce): BCELoss()\n",
      "        (sigm): Sigmoid()\n",
      "        (convnet): ResNet(\n",
      "          (conv1): Conv2d(3, 64, kernel_size=(7, 7), stride=(2, 2), padding=(3, 3), bias=False)\n",
      "          (bn1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "          (relu): ReLU(inplace=True)\n",
      "          (maxpool): MaxPool2d(kernel_size=3, stride=2, padding=1, dilation=1, ceil_mode=False)\n",
      "          (layer1): Sequential(\n",
      "            (0): Bottleneck(\n",
      "              (conv1): Conv2d(64, 64, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "              (bn1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "              (conv2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
      "              (bn2): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "              (conv3): Conv2d(64, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "              (bn3): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "              (relu): ReLU(inplace=True)\n",
      "              (downsample): Sequential(\n",
      "                (0): Conv2d(64, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "                (1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "              )\n",
      "            )\n",
      "            (1): Bottleneck(\n",
      "              (conv1): Conv2d(256, 64, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "              (bn1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "              (conv2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
      "              (bn2): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "              (conv3): Conv2d(64, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "              (bn3): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "              (relu): ReLU(inplace=True)\n",
      "            )\n",
      "            (2): Bottleneck(\n",
      "              (conv1): Conv2d(256, 64, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "              (bn1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "              (conv2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
      "              (bn2): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "              (conv3): Conv2d(64, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "              (bn3): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "              (relu): ReLU(inplace=True)\n",
      "            )\n",
      "          )\n",
      "          (layer2): Sequential(\n",
      "            (0): Bottleneck(\n",
      "              (conv1): Conv2d(256, 128, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "              (bn1): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "              (conv2): Conv2d(128, 128, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)\n",
      "              (bn2): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "              (conv3): Conv2d(128, 512, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "              (bn3): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "              (relu): ReLU(inplace=True)\n",
      "              (downsample): Sequential(\n",
      "                (0): Conv2d(256, 512, kernel_size=(1, 1), stride=(2, 2), bias=False)\n",
      "                (1): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "              )\n",
      "            )\n",
      "            (1): Bottleneck(\n",
      "              (conv1): Conv2d(512, 128, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "              (bn1): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "              (conv2): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
      "              (bn2): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "              (conv3): Conv2d(128, 512, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "              (bn3): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "              (relu): ReLU(inplace=True)\n",
      "            )\n",
      "            (2): Bottleneck(\n",
      "              (conv1): Conv2d(512, 128, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "              (bn1): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "              (conv2): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
      "              (bn2): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "              (conv3): Conv2d(128, 512, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "              (bn3): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "              (relu): ReLU(inplace=True)\n",
      "            )\n",
      "            (3): Bottleneck(\n",
      "              (conv1): Conv2d(512, 128, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "              (bn1): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "              (conv2): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
      "              (bn2): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "              (conv3): Conv2d(128, 512, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "              (bn3): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "              (relu): ReLU(inplace=True)\n",
      "            )\n",
      "          )\n",
      "          (layer3): Sequential(\n",
      "            (0): Bottleneck(\n",
      "              (conv1): Conv2d(512, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "              (bn1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "              (conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)\n",
      "              (bn2): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "              (conv3): Conv2d(256, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "              (bn3): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "              (relu): ReLU(inplace=True)\n",
      "              (downsample): Sequential(\n",
      "                (0): Conv2d(512, 1024, kernel_size=(1, 1), stride=(2, 2), bias=False)\n",
      "                (1): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "              )\n",
      "            )\n",
      "            (1): Bottleneck(\n",
      "              (conv1): Conv2d(1024, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "              (bn1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "              (conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
      "              (bn2): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "              (conv3): Conv2d(256, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "              (bn3): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "              (relu): ReLU(inplace=True)\n",
      "            )\n",
      "            (2): Bottleneck(\n",
      "              (conv1): Conv2d(1024, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "              (bn1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "              (conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
      "              (bn2): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "              (conv3): Conv2d(256, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "              (bn3): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "              (relu): ReLU(inplace=True)\n",
      "            )\n",
      "            (3): Bottleneck(\n",
      "              (conv1): Conv2d(1024, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "              (bn1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "              (conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
      "              (bn2): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "              (conv3): Conv2d(256, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "              (bn3): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "              (relu): ReLU(inplace=True)\n",
      "            )\n",
      "            (4): Bottleneck(\n",
      "              (conv1): Conv2d(1024, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "              (bn1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "              (conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
      "              (bn2): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "              (conv3): Conv2d(256, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "              (bn3): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "              (relu): ReLU(inplace=True)\n",
      "            )\n",
      "            (5): Bottleneck(\n",
      "              (conv1): Conv2d(1024, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "              (bn1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "              (conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
      "              (bn2): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "              (conv3): Conv2d(256, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "              (bn3): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "              (relu): ReLU(inplace=True)\n",
      "            )\n",
      "          )\n",
      "          (layer4): Sequential(\n",
      "            (0): Bottleneck(\n",
      "              (conv1): Conv2d(1024, 512, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "              (bn1): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "              (conv2): Conv2d(512, 512, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)\n",
      "              (bn2): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "              (conv3): Conv2d(512, 2048, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "              (bn3): BatchNorm2d(2048, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "              (relu): ReLU(inplace=True)\n",
      "              (downsample): Sequential(\n",
      "                (0): Conv2d(1024, 2048, kernel_size=(1, 1), stride=(2, 2), bias=False)\n",
      "                (1): BatchNorm2d(2048, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "              )\n",
      "            )\n",
      "            (1): Bottleneck(\n",
      "              (conv1): Conv2d(2048, 512, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "              (bn1): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "              (conv2): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
      "              (bn2): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "              (conv3): Conv2d(512, 2048, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "              (bn3): BatchNorm2d(2048, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "              (relu): ReLU(inplace=True)\n",
      "            )\n",
      "            (2): Bottleneck(\n",
      "              (conv1): Conv2d(2048, 512, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "              (bn1): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "              (conv2): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
      "              (bn2): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "              (conv3): Conv2d(512, 2048, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "              (bn3): BatchNorm2d(2048, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "              (relu): ReLU(inplace=True)\n",
      "            )\n",
      "          )\n",
      "          (avgpool): AdaptiveAvgPool2d(output_size=(1, 1))\n",
      "          (fc): Identity()\n",
      "        )\n",
      "        (normlayer): Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])\n",
      "      )\n",
      "    )\n",
      "  )\n",
      "  (dynamics): LSTM(\n",
      "    (lstm): LSTM(2048, 1024)\n",
      "    (regressor): Linear(in_features=1024, out_features=2048, bias=True)\n",
      "  )\n",
      ")\n",
      "======= Model layers =======\n",
      "['encoder', 'dynamics']\n"
     ]
    }
   ],
   "source": [
    "name = \"pfR3M_LSTM_physion\"\n",
    "model, model_layers = load_pretrained_model(name)\n",
    "print(\"======= Model architecture =======\\n\", model)\n",
    "print(f\"======= Model layers =======\\n{model_layers}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "88895192",
   "metadata": {},
   "source": [
    "## 2d. R3M + LSTM, dynamics module pretrained on Kinetics-700"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "eba3eb3d",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Loaded parameters from /om/weka/yanglab/anayebi/mental-sim/trained_models/R3M+LSTM_k700.pt\n",
      "======= Model architecture =======\n",
      " FrozenPretrainedEncoder(\n",
      "  (encoder): R3M_pretrained(\n",
      "    (r3m): DataParallel(\n",
      "      (module): R3M(\n",
      "        (cs): CosineSimilarity()\n",
      "        (bce): BCELoss()\n",
      "        (sigm): Sigmoid()\n",
      "        (convnet): ResNet(\n",
      "          (conv1): Conv2d(3, 64, kernel_size=(7, 7), stride=(2, 2), padding=(3, 3), bias=False)\n",
      "          (bn1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "          (relu): ReLU(inplace=True)\n",
      "          (maxpool): MaxPool2d(kernel_size=3, stride=2, padding=1, dilation=1, ceil_mode=False)\n",
      "          (layer1): Sequential(\n",
      "            (0): Bottleneck(\n",
      "              (conv1): Conv2d(64, 64, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "              (bn1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "              (conv2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
      "              (bn2): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "              (conv3): Conv2d(64, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "              (bn3): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "              (relu): ReLU(inplace=True)\n",
      "              (downsample): Sequential(\n",
      "                (0): Conv2d(64, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "                (1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "              )\n",
      "            )\n",
      "            (1): Bottleneck(\n",
      "              (conv1): Conv2d(256, 64, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "              (bn1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "              (conv2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
      "              (bn2): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "              (conv3): Conv2d(64, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "              (bn3): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "              (relu): ReLU(inplace=True)\n",
      "            )\n",
      "            (2): Bottleneck(\n",
      "              (conv1): Conv2d(256, 64, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "              (bn1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "              (conv2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
      "              (bn2): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "              (conv3): Conv2d(64, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "              (bn3): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "              (relu): ReLU(inplace=True)\n",
      "            )\n",
      "          )\n",
      "          (layer2): Sequential(\n",
      "            (0): Bottleneck(\n",
      "              (conv1): Conv2d(256, 128, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "              (bn1): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "              (conv2): Conv2d(128, 128, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)\n",
      "              (bn2): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "              (conv3): Conv2d(128, 512, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "              (bn3): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "              (relu): ReLU(inplace=True)\n",
      "              (downsample): Sequential(\n",
      "                (0): Conv2d(256, 512, kernel_size=(1, 1), stride=(2, 2), bias=False)\n",
      "                (1): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "              )\n",
      "            )\n",
      "            (1): Bottleneck(\n",
      "              (conv1): Conv2d(512, 128, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "              (bn1): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "              (conv2): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
      "              (bn2): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "              (conv3): Conv2d(128, 512, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "              (bn3): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "              (relu): ReLU(inplace=True)\n",
      "            )\n",
      "            (2): Bottleneck(\n",
      "              (conv1): Conv2d(512, 128, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "              (bn1): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "              (conv2): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
      "              (bn2): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "              (conv3): Conv2d(128, 512, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "              (bn3): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "              (relu): ReLU(inplace=True)\n",
      "            )\n",
      "            (3): Bottleneck(\n",
      "              (conv1): Conv2d(512, 128, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "              (bn1): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "              (conv2): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
      "              (bn2): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "              (conv3): Conv2d(128, 512, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "              (bn3): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "              (relu): ReLU(inplace=True)\n",
      "            )\n",
      "          )\n",
      "          (layer3): Sequential(\n",
      "            (0): Bottleneck(\n",
      "              (conv1): Conv2d(512, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "              (bn1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "              (conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)\n",
      "              (bn2): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "              (conv3): Conv2d(256, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "              (bn3): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "              (relu): ReLU(inplace=True)\n",
      "              (downsample): Sequential(\n",
      "                (0): Conv2d(512, 1024, kernel_size=(1, 1), stride=(2, 2), bias=False)\n",
      "                (1): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "              )\n",
      "            )\n",
      "            (1): Bottleneck(\n",
      "              (conv1): Conv2d(1024, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "              (bn1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "              (conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
      "              (bn2): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "              (conv3): Conv2d(256, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "              (bn3): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "              (relu): ReLU(inplace=True)\n",
      "            )\n",
      "            (2): Bottleneck(\n",
      "              (conv1): Conv2d(1024, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "              (bn1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "              (conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
      "              (bn2): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "              (conv3): Conv2d(256, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "              (bn3): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "              (relu): ReLU(inplace=True)\n",
      "            )\n",
      "            (3): Bottleneck(\n",
      "              (conv1): Conv2d(1024, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "              (bn1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "              (conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
      "              (bn2): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "              (conv3): Conv2d(256, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "              (bn3): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "              (relu): ReLU(inplace=True)\n",
      "            )\n",
      "            (4): Bottleneck(\n",
      "              (conv1): Conv2d(1024, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "              (bn1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "              (conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
      "              (bn2): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "              (conv3): Conv2d(256, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "              (bn3): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "              (relu): ReLU(inplace=True)\n",
      "            )\n",
      "            (5): Bottleneck(\n",
      "              (conv1): Conv2d(1024, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "              (bn1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "              (conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
      "              (bn2): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "              (conv3): Conv2d(256, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "              (bn3): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "              (relu): ReLU(inplace=True)\n",
      "            )\n",
      "          )\n",
      "          (layer4): Sequential(\n",
      "            (0): Bottleneck(\n",
      "              (conv1): Conv2d(1024, 512, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "              (bn1): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "              (conv2): Conv2d(512, 512, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)\n",
      "              (bn2): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "              (conv3): Conv2d(512, 2048, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "              (bn3): BatchNorm2d(2048, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "              (relu): ReLU(inplace=True)\n",
      "              (downsample): Sequential(\n",
      "                (0): Conv2d(1024, 2048, kernel_size=(1, 1), stride=(2, 2), bias=False)\n",
      "                (1): BatchNorm2d(2048, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "              )\n",
      "            )\n",
      "            (1): Bottleneck(\n",
      "              (conv1): Conv2d(2048, 512, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "              (bn1): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "              (conv2): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
      "              (bn2): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "              (conv3): Conv2d(512, 2048, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "              (bn3): BatchNorm2d(2048, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "              (relu): ReLU(inplace=True)\n",
      "            )\n",
      "            (2): Bottleneck(\n",
      "              (conv1): Conv2d(2048, 512, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "              (bn1): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "              (conv2): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
      "              (bn2): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "              (conv3): Conv2d(512, 2048, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "              (bn3): BatchNorm2d(2048, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "              (relu): ReLU(inplace=True)\n",
      "            )\n",
      "          )\n",
      "          (avgpool): AdaptiveAvgPool2d(output_size=(1, 1))\n",
      "          (fc): Identity()\n",
      "        )\n",
      "        (normlayer): Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])\n",
      "      )\n",
      "    )\n",
      "  )\n",
      "  (dynamics): LSTM(\n",
      "    (lstm): LSTM(2048, 1024)\n",
      "    (regressor): Linear(in_features=1024, out_features=2048, bias=True)\n",
      "  )\n",
      ")\n",
      "======= Model layers =======\n",
      "['encoder', 'dynamics']\n"
     ]
    }
   ],
   "source": [
    "name = \"pfR3M_LSTM_k700\"\n",
    "model, model_layers = load_pretrained_model(name)\n",
    "print(\"======= Model architecture =======\\n\", model)\n",
    "print(f\"======= Model layers =======\\n{model_layers}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b2d34cd7",
   "metadata": {},
   "source": [
    "# 3. FitVid end-to-end pretrained on Physion with RandAugment (64x64 px images)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "id": "fe96eb8d",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Loaded parameters from /om/weka/yanglab/anayebi/mental-sim/trained_models/FitVid_physion_64x64.pt\n",
      "======= Model architecture =======\n",
      " FitVid(\n",
      "  (encoder): NvaeEncoder(\n",
      "    (blocks): ModuleList(\n",
      "      (0): ModuleList(\n",
      "        (0): EncoderBlock(\n",
      "          (conv1): Conv2d(3, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
      "          (conv2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=same, bias=False)\n",
      "          (shortcut): Conv2d(3, 64, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "          (normalize_input): BatchNorm2d(3, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "          (normalize_hidden): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "          (normalize_residual): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "          (se): SEBlock(\n",
      "            (squeeze): Conv2d(64, 4, kernel_size=(1, 1), stride=(1, 1))\n",
      "            (expand): Conv2d(4, 64, kernel_size=(1, 1), stride=(1, 1))\n",
      "          )\n",
      "          (act): SiLU()\n",
      "        )\n",
      "        (1): EncoderBlock(\n",
      "          (conv1): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
      "          (conv2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=same, bias=False)\n",
      "          (shortcut): Conv2d(64, 64, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "          (normalize_input): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "          (normalize_hidden): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "          (normalize_residual): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "          (se): SEBlock(\n",
      "            (squeeze): Conv2d(64, 4, kernel_size=(1, 1), stride=(1, 1))\n",
      "            (expand): Conv2d(4, 64, kernel_size=(1, 1), stride=(1, 1))\n",
      "          )\n",
      "          (act): SiLU()\n",
      "        )\n",
      "      )\n",
      "      (1): ModuleList(\n",
      "        (0): EncoderBlock(\n",
      "          (conv1): Conv2d(64, 128, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)\n",
      "          (conv2): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=same, bias=False)\n",
      "          (shortcut): Conv2d(64, 128, kernel_size=(1, 1), stride=(2, 2), bias=False)\n",
      "          (normalize_input): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "          (normalize_hidden): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "          (normalize_residual): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "          (se): SEBlock(\n",
      "            (squeeze): Conv2d(128, 8, kernel_size=(1, 1), stride=(1, 1))\n",
      "            (expand): Conv2d(8, 128, kernel_size=(1, 1), stride=(1, 1))\n",
      "          )\n",
      "          (act): SiLU()\n",
      "        )\n",
      "        (1): EncoderBlock(\n",
      "          (conv1): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
      "          (conv2): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=same, bias=False)\n",
      "          (shortcut): Conv2d(128, 128, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "          (normalize_input): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "          (normalize_hidden): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "          (normalize_residual): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "          (se): SEBlock(\n",
      "            (squeeze): Conv2d(128, 8, kernel_size=(1, 1), stride=(1, 1))\n",
      "            (expand): Conv2d(8, 128, kernel_size=(1, 1), stride=(1, 1))\n",
      "          )\n",
      "          (act): SiLU()\n",
      "        )\n",
      "      )\n",
      "      (2): ModuleList(\n",
      "        (0): EncoderBlock(\n",
      "          (conv1): Conv2d(128, 256, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)\n",
      "          (conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=same, bias=False)\n",
      "          (shortcut): Conv2d(128, 256, kernel_size=(1, 1), stride=(2, 2), bias=False)\n",
      "          (normalize_input): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "          (normalize_hidden): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "          (normalize_residual): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "          (se): SEBlock(\n",
      "            (squeeze): Conv2d(256, 16, kernel_size=(1, 1), stride=(1, 1))\n",
      "            (expand): Conv2d(16, 256, kernel_size=(1, 1), stride=(1, 1))\n",
      "          )\n",
      "          (act): SiLU()\n",
      "        )\n",
      "        (1): EncoderBlock(\n",
      "          (conv1): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
      "          (conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=same, bias=False)\n",
      "          (shortcut): Conv2d(256, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "          (normalize_input): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "          (normalize_hidden): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "          (normalize_residual): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "          (se): SEBlock(\n",
      "            (squeeze): Conv2d(256, 16, kernel_size=(1, 1), stride=(1, 1))\n",
      "            (expand): Conv2d(16, 256, kernel_size=(1, 1), stride=(1, 1))\n",
      "          )\n",
      "          (act): SiLU()\n",
      "        )\n",
      "      )\n",
      "      (3): ModuleList(\n",
      "        (0): EncoderBlock(\n",
      "          (conv1): Conv2d(256, 512, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)\n",
      "          (conv2): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=same, bias=False)\n",
      "          (shortcut): Conv2d(256, 512, kernel_size=(1, 1), stride=(2, 2), bias=False)\n",
      "          (normalize_input): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "          (normalize_hidden): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "          (normalize_residual): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "          (se): SEBlock(\n",
      "            (squeeze): Conv2d(512, 32, kernel_size=(1, 1), stride=(1, 1))\n",
      "            (expand): Conv2d(32, 512, kernel_size=(1, 1), stride=(1, 1))\n",
      "          )\n",
      "          (act): SiLU()\n",
      "        )\n",
      "        (1): EncoderBlock(\n",
      "          (conv1): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
      "          (conv2): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=same, bias=False)\n",
      "          (shortcut): Conv2d(512, 512, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "          (normalize_input): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "          (normalize_hidden): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "          (normalize_residual): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "          (se): SEBlock(\n",
      "            (squeeze): Conv2d(512, 32, kernel_size=(1, 1), stride=(1, 1))\n",
      "            (expand): Conv2d(32, 512, kernel_size=(1, 1), stride=(1, 1))\n",
      "          )\n",
      "          (act): SiLU()\n",
      "        )\n",
      "      )\n",
      "    )\n",
      "    (linear): Linear(in_features=512, out_features=128, bias=True)\n",
      "  )\n",
      "  (decoder): NvaeDecoder(\n",
      "    (linear): Linear(in_features=128, out_features=32768, bias=True)\n",
      "    (blocks): ModuleList(\n",
      "      (0): ModuleList(\n",
      "        (0-1): 2 x DecoderBlock(\n",
      "          (act): SiLU()\n",
      "          (conv1): Conv2d(512, 2048, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "          (conv2): Conv2d(2048, 2048, kernel_size=(5, 5), stride=(1, 1), padding=same, bias=False)\n",
      "          (conv3): Conv2d(2048, 512, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "          (shortcut): Conv2d(512, 512, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "          (normalize_input): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "          (normalize_hidden1): BatchNorm2d(2048, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "          (normalize_hidden2): BatchNorm2d(2048, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "          (normalize_output): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "          (normalize_residual): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "          (se): SEBlock(\n",
      "            (squeeze): Conv2d(512, 32, kernel_size=(1, 1), stride=(1, 1))\n",
      "            (expand): Conv2d(32, 512, kernel_size=(1, 1), stride=(1, 1))\n",
      "          )\n",
      "        )\n",
      "      )\n",
      "      (1): ModuleList(\n",
      "        (0): DecoderBlock(\n",
      "          (act): SiLU()\n",
      "          (conv1): Conv2d(512, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "          (conv2): Conv2d(1024, 1024, kernel_size=(5, 5), stride=(1, 1), padding=same, bias=False)\n",
      "          (conv3): Conv2d(1024, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "          (shortcut): Conv2d(512, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "          (normalize_input): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "          (normalize_hidden1): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "          (normalize_hidden2): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "          (normalize_output): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "          (normalize_residual): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "          (se): SEBlock(\n",
      "            (squeeze): Conv2d(256, 16, kernel_size=(1, 1), stride=(1, 1))\n",
      "            (expand): Conv2d(16, 256, kernel_size=(1, 1), stride=(1, 1))\n",
      "          )\n",
      "        )\n",
      "        (1): DecoderBlock(\n",
      "          (act): SiLU()\n",
      "          (conv1): Conv2d(256, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "          (conv2): Conv2d(1024, 1024, kernel_size=(5, 5), stride=(1, 1), padding=same, bias=False)\n",
      "          (conv3): Conv2d(1024, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "          (shortcut): Conv2d(256, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "          (normalize_input): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "          (normalize_hidden1): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "          (normalize_hidden2): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "          (normalize_output): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "          (normalize_residual): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "          (se): SEBlock(\n",
      "            (squeeze): Conv2d(256, 16, kernel_size=(1, 1), stride=(1, 1))\n",
      "            (expand): Conv2d(16, 256, kernel_size=(1, 1), stride=(1, 1))\n",
      "          )\n",
      "        )\n",
      "      )\n",
      "      (2): ModuleList(\n",
      "        (0): DecoderBlock(\n",
      "          (act): SiLU()\n",
      "          (conv1): Conv2d(256, 512, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "          (conv2): Conv2d(512, 512, kernel_size=(5, 5), stride=(1, 1), padding=same, bias=False)\n",
      "          (conv3): Conv2d(512, 128, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "          (shortcut): Conv2d(256, 128, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "          (normalize_input): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "          (normalize_hidden1): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "          (normalize_hidden2): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "          (normalize_output): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "          (normalize_residual): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "          (se): SEBlock(\n",
      "            (squeeze): Conv2d(128, 8, kernel_size=(1, 1), stride=(1, 1))\n",
      "            (expand): Conv2d(8, 128, kernel_size=(1, 1), stride=(1, 1))\n",
      "          )\n",
      "        )\n",
      "        (1): DecoderBlock(\n",
      "          (act): SiLU()\n",
      "          (conv1): Conv2d(128, 512, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "          (conv2): Conv2d(512, 512, kernel_size=(5, 5), stride=(1, 1), padding=same, bias=False)\n",
      "          (conv3): Conv2d(512, 128, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "          (shortcut): Conv2d(128, 128, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "          (normalize_input): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "          (normalize_hidden1): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "          (normalize_hidden2): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "          (normalize_output): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "          (normalize_residual): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "          (se): SEBlock(\n",
      "            (squeeze): Conv2d(128, 8, kernel_size=(1, 1), stride=(1, 1))\n",
      "            (expand): Conv2d(8, 128, kernel_size=(1, 1), stride=(1, 1))\n",
      "          )\n",
      "        )\n",
      "      )\n",
      "      (3): ModuleList(\n",
      "        (0): DecoderBlock(\n",
      "          (act): SiLU()\n",
      "          (conv1): Conv2d(128, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "          (conv2): Conv2d(256, 256, kernel_size=(5, 5), stride=(1, 1), padding=same, bias=False)\n",
      "          (conv3): Conv2d(256, 64, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "          (shortcut): Conv2d(128, 64, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "          (normalize_input): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "          (normalize_hidden1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "          (normalize_hidden2): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "          (normalize_output): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "          (normalize_residual): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "          (se): SEBlock(\n",
      "            (squeeze): Conv2d(64, 4, kernel_size=(1, 1), stride=(1, 1))\n",
      "            (expand): Conv2d(4, 64, kernel_size=(1, 1), stride=(1, 1))\n",
      "          )\n",
      "        )\n",
      "        (1): DecoderBlock(\n",
      "          (act): SiLU()\n",
      "          (conv1): Conv2d(64, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "          (conv2): Conv2d(256, 256, kernel_size=(5, 5), stride=(1, 1), padding=same, bias=False)\n",
      "          (conv3): Conv2d(256, 64, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "          (shortcut): Conv2d(64, 64, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "          (normalize_input): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "          (normalize_hidden1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "          (normalize_hidden2): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "          (normalize_output): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "          (normalize_residual): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "          (se): SEBlock(\n",
      "            (squeeze): Conv2d(64, 4, kernel_size=(1, 1), stride=(1, 1))\n",
      "            (expand): Conv2d(4, 64, kernel_size=(1, 1), stride=(1, 1))\n",
      "          )\n",
      "        )\n",
      "      )\n",
      "    )\n",
      "    (output_conv): Conv2d(64, 3, kernel_size=(3, 3), stride=(1, 1), padding=same, bias=False)\n",
      "  )\n",
      "  (frame_predictor): MultiGaussianLSTM(\n",
      "    (embed): Linear(in_features=138, out_features=256, bias=True)\n",
      "    (mean): Linear(in_features=256, out_features=128, bias=True)\n",
      "    (logvar): Linear(in_features=256, out_features=128, bias=True)\n",
      "    (layers): ModuleList(\n",
      "      (0-1): 2 x LSTMCell(256, 256)\n",
      "    )\n",
      "  )\n",
      "  (posterior): MultiGaussianLSTM(\n",
      "    (embed): Linear(in_features=128, out_features=256, bias=True)\n",
      "    (mean): Linear(in_features=256, out_features=10, bias=True)\n",
      "    (logvar): Linear(in_features=256, out_features=10, bias=True)\n",
      "    (layers): ModuleList(\n",
      "      (0): LSTMCell(256, 256)\n",
      "    )\n",
      "  )\n",
      "  (prior): MultiGaussianLSTM(\n",
      "    (embed): Linear(in_features=128, out_features=256, bias=True)\n",
      "    (mean): Linear(in_features=256, out_features=10, bias=True)\n",
      "    (logvar): Linear(in_features=256, out_features=10, bias=True)\n",
      "    (layers): ModuleList(\n",
      "      (0): LSTMCell(256, 256)\n",
      "    )\n",
      "  )\n",
      ")\n",
      "======= Model layers =======\n",
      "['encoder', 'dynamics', 'decoder']\n"
     ]
    }
   ],
   "source": [
    "name = \"fitvid_ctxt7_physion_aug_frames_64x64\"\n",
    "model, model_layers = load_pretrained_model(name)\n",
    "print(\"======= Model architecture =======\\n\", model)\n",
    "print(f\"======= Model layers =======\\n{model_layers}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "066e32b0",
   "metadata": {},
   "source": [
    "# 4. SVG end-to-end pretrained on Physion with 128x128 px images"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "id": "fcec767a",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Loaded parameters from /om/weka/yanglab/anayebi/mental-sim/trained_models/SVG_physion_128x128.pt\n",
      "======= Model architecture =======\n",
      " SVG(\n",
      "  (encoder): encoder(\n",
      "    (c1): Sequential(\n",
      "      (0): vgg_layer(\n",
      "        (main): Sequential(\n",
      "          (0): Conv2d(3, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "          (1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "          (2): LeakyReLU(negative_slope=0.2, inplace=True)\n",
      "        )\n",
      "      )\n",
      "      (1): vgg_layer(\n",
      "        (main): Sequential(\n",
      "          (0): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "          (1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "          (2): LeakyReLU(negative_slope=0.2, inplace=True)\n",
      "        )\n",
      "      )\n",
      "    )\n",
      "    (c2): Sequential(\n",
      "      (0): vgg_layer(\n",
      "        (main): Sequential(\n",
      "          (0): Conv2d(64, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "          (1): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "          (2): LeakyReLU(negative_slope=0.2, inplace=True)\n",
      "        )\n",
      "      )\n",
      "      (1): vgg_layer(\n",
      "        (main): Sequential(\n",
      "          (0): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "          (1): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "          (2): LeakyReLU(negative_slope=0.2, inplace=True)\n",
      "        )\n",
      "      )\n",
      "    )\n",
      "    (c3): Sequential(\n",
      "      (0): vgg_layer(\n",
      "        (main): Sequential(\n",
      "          (0): Conv2d(128, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "          (1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "          (2): LeakyReLU(negative_slope=0.2, inplace=True)\n",
      "        )\n",
      "      )\n",
      "      (1): vgg_layer(\n",
      "        (main): Sequential(\n",
      "          (0): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "          (1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "          (2): LeakyReLU(negative_slope=0.2, inplace=True)\n",
      "        )\n",
      "      )\n",
      "      (2): vgg_layer(\n",
      "        (main): Sequential(\n",
      "          (0): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "          (1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "          (2): LeakyReLU(negative_slope=0.2, inplace=True)\n",
      "        )\n",
      "      )\n",
      "    )\n",
      "    (c4): Sequential(\n",
      "      (0): vgg_layer(\n",
      "        (main): Sequential(\n",
      "          (0): Conv2d(256, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "          (1): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "          (2): LeakyReLU(negative_slope=0.2, inplace=True)\n",
      "        )\n",
      "      )\n",
      "      (1): vgg_layer(\n",
      "        (main): Sequential(\n",
      "          (0): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "          (1): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "          (2): LeakyReLU(negative_slope=0.2, inplace=True)\n",
      "        )\n",
      "      )\n",
      "      (2): vgg_layer(\n",
      "        (main): Sequential(\n",
      "          (0): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "          (1): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "          (2): LeakyReLU(negative_slope=0.2, inplace=True)\n",
      "        )\n",
      "      )\n",
      "    )\n",
      "    (c5): Sequential(\n",
      "      (0): vgg_layer(\n",
      "        (main): Sequential(\n",
      "          (0): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "          (1): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "          (2): LeakyReLU(negative_slope=0.2, inplace=True)\n",
      "        )\n",
      "      )\n",
      "      (1): vgg_layer(\n",
      "        (main): Sequential(\n",
      "          (0): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "          (1): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "          (2): LeakyReLU(negative_slope=0.2, inplace=True)\n",
      "        )\n",
      "      )\n",
      "      (2): vgg_layer(\n",
      "        (main): Sequential(\n",
      "          (0): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "          (1): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "          (2): LeakyReLU(negative_slope=0.2, inplace=True)\n",
      "        )\n",
      "      )\n",
      "    )\n",
      "    (c6): Sequential(\n",
      "      (0): Conv2d(512, 128, kernel_size=(4, 4), stride=(1, 1))\n",
      "      (1): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      (2): Tanh()\n",
      "    )\n",
      "    (mp): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)\n",
      "  )\n",
      "  (decoder): decoder(\n",
      "    (upc1): Sequential(\n",
      "      (0): ConvTranspose2d(128, 512, kernel_size=(4, 4), stride=(1, 1))\n",
      "      (1): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      (2): LeakyReLU(negative_slope=0.2, inplace=True)\n",
      "    )\n",
      "    (upc2): Sequential(\n",
      "      (0): vgg_layer(\n",
      "        (main): Sequential(\n",
      "          (0): Conv2d(1024, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "          (1): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "          (2): LeakyReLU(negative_slope=0.2, inplace=True)\n",
      "        )\n",
      "      )\n",
      "      (1): vgg_layer(\n",
      "        (main): Sequential(\n",
      "          (0): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "          (1): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "          (2): LeakyReLU(negative_slope=0.2, inplace=True)\n",
      "        )\n",
      "      )\n",
      "      (2): vgg_layer(\n",
      "        (main): Sequential(\n",
      "          (0): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "          (1): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "          (2): LeakyReLU(negative_slope=0.2, inplace=True)\n",
      "        )\n",
      "      )\n",
      "    )\n",
      "    (upc3): Sequential(\n",
      "      (0): vgg_layer(\n",
      "        (main): Sequential(\n",
      "          (0): Conv2d(1024, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "          (1): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "          (2): LeakyReLU(negative_slope=0.2, inplace=True)\n",
      "        )\n",
      "      )\n",
      "      (1): vgg_layer(\n",
      "        (main): Sequential(\n",
      "          (0): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "          (1): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "          (2): LeakyReLU(negative_slope=0.2, inplace=True)\n",
      "        )\n",
      "      )\n",
      "      (2): vgg_layer(\n",
      "        (main): Sequential(\n",
      "          (0): Conv2d(512, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "          (1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "          (2): LeakyReLU(negative_slope=0.2, inplace=True)\n",
      "        )\n",
      "      )\n",
      "    )\n",
      "    (upc4): Sequential(\n",
      "      (0): vgg_layer(\n",
      "        (main): Sequential(\n",
      "          (0): Conv2d(512, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "          (1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "          (2): LeakyReLU(negative_slope=0.2, inplace=True)\n",
      "        )\n",
      "      )\n",
      "      (1): vgg_layer(\n",
      "        (main): Sequential(\n",
      "          (0): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "          (1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "          (2): LeakyReLU(negative_slope=0.2, inplace=True)\n",
      "        )\n",
      "      )\n",
      "      (2): vgg_layer(\n",
      "        (main): Sequential(\n",
      "          (0): Conv2d(256, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "          (1): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "          (2): LeakyReLU(negative_slope=0.2, inplace=True)\n",
      "        )\n",
      "      )\n",
      "    )\n",
      "    (upc5): Sequential(\n",
      "      (0): vgg_layer(\n",
      "        (main): Sequential(\n",
      "          (0): Conv2d(256, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "          (1): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "          (2): LeakyReLU(negative_slope=0.2, inplace=True)\n",
      "        )\n",
      "      )\n",
      "      (1): vgg_layer(\n",
      "        (main): Sequential(\n",
      "          (0): Conv2d(128, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "          (1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "          (2): LeakyReLU(negative_slope=0.2, inplace=True)\n",
      "        )\n",
      "      )\n",
      "    )\n",
      "    (upc6): Sequential(\n",
      "      (0): vgg_layer(\n",
      "        (main): Sequential(\n",
      "          (0): Conv2d(128, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "          (1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "          (2): LeakyReLU(negative_slope=0.2, inplace=True)\n",
      "        )\n",
      "      )\n",
      "      (1): ConvTranspose2d(64, 3, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "      (2): Sigmoid()\n",
      "    )\n",
      "    (up): UpsamplingNearest2d(scale_factor=2.0, mode='nearest')\n",
      "  )\n",
      "  (frame_predictor): lstm(\n",
      "    (embed): Linear(in_features=138, out_features=256, bias=True)\n",
      "    (lstm): ModuleList(\n",
      "      (0-1): 2 x LSTMCell(256, 256)\n",
      "    )\n",
      "    (output): Sequential(\n",
      "      (0): Linear(in_features=256, out_features=128, bias=True)\n",
      "      (1): Tanh()\n",
      "    )\n",
      "  )\n",
      "  (posterior): gaussian_lstm(\n",
      "    (embed): Linear(in_features=128, out_features=256, bias=True)\n",
      "    (lstm): ModuleList(\n",
      "      (0): LSTMCell(256, 256)\n",
      "    )\n",
      "    (mu_net): Linear(in_features=256, out_features=10, bias=True)\n",
      "    (logvar_net): Linear(in_features=256, out_features=10, bias=True)\n",
      "  )\n",
      "  (prior): gaussian_lstm(\n",
      "    (embed): Linear(in_features=128, out_features=256, bias=True)\n",
      "    (lstm): ModuleList(\n",
      "      (0): LSTMCell(256, 256)\n",
      "    )\n",
      "    (mu_net): Linear(in_features=256, out_features=10, bias=True)\n",
      "    (logvar_net): Linear(in_features=256, out_features=10, bias=True)\n",
      "  )\n",
      ")\n",
      "======= Model layers =======\n",
      "['encoder', 'dynamics', 'decoder']\n"
     ]
    }
   ],
   "source": [
    "name = \"svg_physion_128x128\"\n",
    "model, model_layers = load_pretrained_model(name)\n",
    "print(\"======= Model architecture =======\\n\", model)\n",
    "print(f\"======= Model layers =======\\n{model_layers}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "89270c36",
   "metadata": {},
   "source": [
    "# 5. C-SWM with Large Encoder end-to-end pretrained on Physion"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "id": "69d8bcc7",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Loaded parameters from /om/weka/yanglab/anayebi/mental-sim/trained_models/CSWM_large_physion.pt\n",
      "======= Model architecture =======\n",
      " PassiveVideoCSWM(\n",
      "  (model): ContrastiveSWM(\n",
      "    (obj_extractor): EncoderCNNLarge(\n",
      "      (cnn1): Conv2d(21, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "      (act1): ReLU()\n",
      "      (ln1): BatchNorm2d(32, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      (cnn2): Conv2d(32, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "      (act2): ReLU()\n",
      "      (ln2): BatchNorm2d(32, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      (cnn3): Conv2d(32, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "      (act3): ReLU()\n",
      "      (ln3): BatchNorm2d(32, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      (cnn4): Conv2d(32, 10, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "      (act4): Sigmoid()\n",
      "    )\n",
      "    (obj_encoder): EncoderMLP(\n",
      "      (fc1): Linear(in_features=50176, out_features=512, bias=True)\n",
      "      (fc2): Linear(in_features=512, out_features=512, bias=True)\n",
      "      (fc3): Linear(in_features=512, out_features=128, bias=True)\n",
      "      (ln): LayerNorm((512,), eps=1e-05, elementwise_affine=True)\n",
      "      (act1): ReLU()\n",
      "      (act2): ReLU()\n",
      "    )\n",
      "    (transition_model): TransitionGNN(\n",
      "      (edge_mlp): Sequential(\n",
      "        (0): Linear(in_features=256, out_features=512, bias=True)\n",
      "        (1): ReLU()\n",
      "        (2): Linear(in_features=512, out_features=512, bias=True)\n",
      "        (3): LayerNorm((512,), eps=1e-05, elementwise_affine=True)\n",
      "        (4): ReLU()\n",
      "        (5): Linear(in_features=512, out_features=512, bias=True)\n",
      "      )\n",
      "      (node_mlp): Sequential(\n",
      "        (0): Linear(in_features=640, out_features=512, bias=True)\n",
      "        (1): ReLU()\n",
      "        (2): Linear(in_features=512, out_features=512, bias=True)\n",
      "        (3): LayerNorm((512,), eps=1e-05, elementwise_affine=True)\n",
      "        (4): ReLU()\n",
      "        (5): Linear(in_features=512, out_features=128, bias=True)\n",
      "      )\n",
      "    )\n",
      "  )\n",
      ")\n",
      "======= Model layers =======\n",
      "['encoder', 'dynamics']\n"
     ]
    }
   ],
   "source": [
    "name = \"large_CSWM_physion\"\n",
    "model, model_layers = load_pretrained_model(name)\n",
    "print(\"======= Model architecture =======\\n\", model)\n",
    "print(f\"======= Model layers =======\\n{model_layers}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a9550d88",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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.9.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}