mirror of
https://github.com/wassname/iris_bigvae.git
synced 2026-09-09 11:24:31 +08:00
748 lines
33 KiB
Plaintext
748 lines
33 KiB
Plaintext
{
|
|
"cells": [
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": 1,
|
|
"metadata": {},
|
|
"outputs": [
|
|
{
|
|
"name": "stderr",
|
|
"output_type": "stream",
|
|
"text": [
|
|
"\n"
|
|
]
|
|
}
|
|
],
|
|
"source": [
|
|
"# autoreload import your package\n",
|
|
"%load_ext autoreload\n",
|
|
"%autoreload 2\n",
|
|
"\n",
|
|
"import gym\n",
|
|
"import numpy as np\n",
|
|
"import matplotlib.pyplot as plt\n",
|
|
"%matplotlib inline\n",
|
|
"plt.style.use('ggplot')\n"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {},
|
|
"source": [
|
|
"## Debug model components\n",
|
|
"\n",
|
|
"### Using trainer? :poop:\n",
|
|
"\n",
|
|
"Hyrda is really annoying\n",
|
|
"\n",
|
|
"\n"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": 2,
|
|
"metadata": {},
|
|
"outputs": [
|
|
{
|
|
"name": "stderr",
|
|
"output_type": "stream",
|
|
"text": [
|
|
"/media/wassname/SGIronWolf/projects5/worldmodels/iris_bigvae/.venv/lib/python3.9/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
|
|
" from .autonotebook import tqdm as notebook_tqdm\n",
|
|
"Failed to detect the name of this notebook, you can set it manually with the WANDB_NOTEBOOK_NAME environment variable to enable code saving.\n"
|
|
]
|
|
},
|
|
{
|
|
"name": "stdout",
|
|
"output_type": "stream",
|
|
"text": [
|
|
"{'wandb': {'mode': 'disabled', 'project': 'iris', 'entity': None, 'name': None, 'group': None, 'tags': None, 'notes': None}, 'initialization': {'path_to_checkpoint': None, 'load_tokenizer': False, 'load_world_model': False, 'load_actor_critic': False}, 'common': {'epochs': 600, 'device': 'cuda:0', 'do_checkpoint': False, 'seed': 0, 'sequence_length': '${world_model.max_blocks}', 'resume': True}, 'collection': {'train': {'num_envs': 1, 'stop_after_epochs': 500, 'num_episodes_to_save': 10, 'config': {'epsilon': 0.01, 'should_sample': True, 'temperature': 1.0, 'num_steps': 200, 'burn_in': '${training.actor_critic.burn_in}'}}, 'test': {'num_envs': 8, 'num_episodes_to_save': '${collection.train.num_episodes_to_save}', 'config': {'epsilon': 0.0, 'should_sample': True, 'temperature': 0.5, 'num_episodes': 16, 'burn_in': '${training.actor_critic.burn_in}'}}}, 'training': {'should': True, 'learning_rate': 0.0001, 'tokenizer': {'batch_num_samples': 128, 'grad_acc_steps': 1, 'max_grad_norm': 10.0, 'start_after_epochs': 1, 'steps_per_epoch': 10}, 'world_model': {'batch_num_samples': 4, 'grad_acc_steps': 1, 'max_grad_norm': 10.0, 'weight_decay': 0.01, 'start_after_epochs': 1, 'steps_per_epoch': 10}, 'actor_critic': {'batch_num_samples': 4, 'grad_acc_steps': 1, 'max_grad_norm': 10.0, 'start_after_epochs': 1, 'steps_per_epoch': 10, 'imagine_horizon': '${common.sequence_length}', 'burn_in': 20, 'gamma': 0.995, 'lambda_': 0.95, 'entropy_weight': 0.001}}, 'evaluation': {'should': True, 'every': 5, 'tokenizer': {'batch_num_samples': '${training.tokenizer.batch_num_samples}', 'start_after_epochs': '${training.tokenizer.start_after_epochs}', 'save_reconstructions': True}, 'world_model': {'batch_num_samples': '${training.world_model.batch_num_samples}', 'start_after_epochs': '${training.world_model.start_after_epochs}'}, 'actor_critic': {'num_episodes_to_save': '${training.actor_critic.batch_num_samples}', 'horizon': '${training.actor_critic.imagine_horizon}', 'start_after_epochs': '${training.actor_critic.start_after_epochs}'}}, 'tokenizer': {'_target_': 'src.models.tokenizer.Tokenizer', 'vocab_size': 2048, 'embed_dim': 2048, 'encoder': {'_target_': 'src.models.tokenizer.Encoder', 'config': {'_target_': 'src.models.tokenizer.EncoderDecoderConfig', 'resolution': 64, 'in_channels': 3, 'z_channels': 2048, 'ch': 64, 'ch_mult': [1, 1, 1, 1, 1], 'num_res_blocks': 2, 'attn_resolutions': [8, 16], 'out_ch': 3, 'dropout': 0.0}}, 'decoder': {'_target_': 'src.models.tokenizer.Decoder', 'config': '${..encoder.config}'}}, 'world_model': {'_target_': 'src.models.TransformerConfig', 'max_blocks': 10, 'num_layers': 1, 'num_heads': 1, 'embed_dim': 2048, 'dropout': 0.1, 'model_name': 'PY007/TinyLlama-1.1B-intermediate-step-715k-1.5T', 'rank': 32, 'tokens_per_block': 17}, 'actor_critic': {'use_original_obs': False, 'lstm_dim': 512}, 'env': {'train': {'_target_': 'src.envs.make_env', 'id': 'CrafterReward-v1', 'size': 64, 'max_episode_steps': 20000, 'noop_max': 30, 'frame_skip': 4, 'done_on_life_loss': True, 'clip_reward': False}, 'test': {'_target_': '${..train._target_}', 'id': '${..train.id}', 'size': '${..train.size}', 'max_episode_steps': 108000, 'noop_max': 1, 'frame_skip': '${..train.frame_skip}', 'done_on_life_loss': False, 'clip_reward': False}, 'keymap': 'atari/${.train.id}'}, 'datasets': {'train': {'_target_': 'src.dataset.EpisodesDatasetRamMonitoring', 'max_ram_usage': '30G', 'name': 'train_dataset'}, 'test': {'_target_': 'src.dataset.EpisodesDataset', 'max_num_episodes': None, 'name': 'test_dataset'}}}\n",
|
|
"Tokenizer : shape of latent is (2048, 4, 4).\n"
|
|
]
|
|
},
|
|
{
|
|
"name": "stderr",
|
|
"output_type": "stream",
|
|
"text": [
|
|
"/media/wassname/SGIronWolf/projects5/worldmodels/iris_bigvae/.venv/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",
|
|
"/media/wassname/SGIronWolf/projects5/worldmodels/iris_bigvae/.venv/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=VGG16_Weights.IMAGENET1K_V1`. You can also use `weights=VGG16_Weights.DEFAULT` to get the most up-to-date weights.\n",
|
|
" warnings.warn(msg)\n",
|
|
"Using pad_token, but it is not set yet.\n"
|
|
]
|
|
},
|
|
{
|
|
"name": "stdout",
|
|
"output_type": "stream",
|
|
"text": [
|
|
"trainable params: 50,462,720 || all params: 1,150,511,104 || trainable%: 4.386113252149889\n",
|
|
"None\n",
|
|
"32314243 parameters in agent.tokenizer\n",
|
|
"752979973 parameters in agent.world_model\n",
|
|
"3224626 parameters in agent.actor_critic\n"
|
|
]
|
|
}
|
|
],
|
|
"source": [
|
|
"\n",
|
|
"import os\n",
|
|
"os.environ['WANDB_MODE'] = 'disabled'\n",
|
|
"\n",
|
|
"import hydra\n",
|
|
"from hydra import initialize, initialize_config_module, initialize_config_dir, compose\n",
|
|
"from omegaconf import OmegaConf\n",
|
|
"\n",
|
|
"from pathlib import Path\n",
|
|
"from datetime import datetime\n",
|
|
"\n",
|
|
"from src.trainer import Trainer\n",
|
|
"\n",
|
|
"\n",
|
|
"class Trainer2(Trainer):\n",
|
|
" \n",
|
|
" def load_checkpoint(self, *args, **kwargs):\n",
|
|
" pass\n",
|
|
"\n",
|
|
"\n",
|
|
"\n",
|
|
"ts = datetime.now().strftime(\"%Y-%m-%d/%H-%M-%S\")\n",
|
|
"run_dir = Path(f\"..outputs/{ts}\").absolute()\n",
|
|
"run_dir.mkdir(parents=True, exist_ok=True)\n",
|
|
"abs_config_dir=os.path.abspath(\"../config\")\n",
|
|
"os.chdir(run_dir)\n",
|
|
"# with initialize_config_dir(version_base=None, config_dir=abs_config_dir):\n",
|
|
"with initialize(version_base=None, config_path=\"../config\"):\n",
|
|
" cfg = compose(config_name='trainer', overrides=[\n",
|
|
" f'hydra.run.dir={run_dir}',\n",
|
|
" # f\"initialization.path_to_checkpoint={str(path_to_checkpoint.absolute())}\",\n",
|
|
" 'wandb.mode=disabled',\n",
|
|
" \"env.train.id=CrafterReward-v1\",\n",
|
|
" \"training.tokenizer.start_after_epochs=1\",\n",
|
|
" \"training.world_model.start_after_epochs=1\",\n",
|
|
" \"training.actor_critic.start_after_epochs=1\",\n",
|
|
" \"training.tokenizer.steps_per_epoch=10\",\n",
|
|
" \"training.world_model.steps_per_epoch=10\",\n",
|
|
" \"training.actor_critic.steps_per_epoch=10\",\n",
|
|
" \"common.do_checkpoint=False\",\n",
|
|
" \"common.resume=True\",\n",
|
|
" \"training.world_model.batch_num_samples=4\",\n",
|
|
" \"training.actor_critic.batch_num_samples=4\",\n",
|
|
" ])\n",
|
|
" print(cfg)\n",
|
|
"\n",
|
|
" with run_dir:\n",
|
|
" Path('media/episodes/train').mkdir(parents=True, exist_ok=True)\n",
|
|
" trainer = Trainer2(cfg)\n",
|
|
" trainer\n"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": []
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {},
|
|
"source": [
|
|
"\n",
|
|
"## Trainer train_agent\n"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": 3,
|
|
"metadata": {},
|
|
"outputs": [
|
|
{
|
|
"name": "stdout",
|
|
"output_type": "stream",
|
|
"text": [
|
|
"Experience collection (train_dataset): 100%|██████████| 200/200 [00:03<00:00, 58.00it/s]\n"
|
|
]
|
|
},
|
|
{
|
|
"data": {
|
|
"text/plain": [
|
|
"[{'train_dataset/episode_length': 183,\n",
|
|
" 'train_dataset/episode_return': tensor(0.1000),\n",
|
|
" 'train_dataset/episode_num': 0,\n",
|
|
" 'train_dataset/action_histogram': <wandb.sdk.data_types.histogram.Histogram at 0x7f97f7373eb0>},\n",
|
|
" {'train_dataset/#episodes': 2,\n",
|
|
" 'train_dataset/#steps': 200,\n",
|
|
" 'train_dataset/return': 0.100000024}]"
|
|
]
|
|
},
|
|
"execution_count": 3,
|
|
"metadata": {},
|
|
"output_type": "execute_result"
|
|
}
|
|
],
|
|
"source": [
|
|
"self=trainer\n",
|
|
"epoch = 52\n",
|
|
"\n",
|
|
"# get out first exp\n",
|
|
"self.train_collector.collect(self.agent, epoch, **self.cfg.collection.train.config)\n"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": 4,
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"self.agent.train()\n",
|
|
"self.agent.zero_grad()\n",
|
|
"\n",
|
|
"metrics_tokenizer, metrics_world_model, metrics_actor_critic = {}, {}, {}\n",
|
|
"\n",
|
|
"cfg_tokenizer = self.cfg.training.tokenizer\n",
|
|
"cfg_world_model = self.cfg.training.world_model\n",
|
|
"cfg_actor_critic = self.cfg.training.actor_critic\n",
|
|
"\n",
|
|
"# if epoch > cfg_tokenizer.start_after_epochs:\n",
|
|
"# metrics_tokenizer = self.train_component(self.agent.tokenizer, self.optimizer_tokenizer, sequence_length=1, sample_from_start=True, **cfg_tokenizer)\n",
|
|
"# self.agent.tokenizer.eval()\n",
|
|
"\n",
|
|
"# if epoch > cfg_world_model.start_after_epochs:\n",
|
|
"# metrics_world_model = self.train_component(self.agent.world_model, self.optimizer_world_model, sequence_length=self.cfg.common.sequence_length, sample_from_start=True, tokenizer=self.agent.tokenizer, **cfg_world_model)\n",
|
|
"# self.agent.world_model.eval()\n",
|
|
"\n",
|
|
"# if epoch > cfg_actor_critic.start_after_epochs:\n",
|
|
"# metrics_actor_critic = self.train_component(self.agent.actor_critic, self.optimizer_actor_critic, sequence_length=1 + self.cfg.training.actor_critic.burn_in, sample_from_start=False, tokenizer=self.agent.tokenizer, world_model=self.agent.world_model, **cfg_actor_critic)\n",
|
|
"# self.agent.actor_critic.eval()\n"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": 5,
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"import torch\n",
|
|
"from torchinfo import summary\n",
|
|
"import torch\n",
|
|
"from einops import rearrange\n"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {},
|
|
"source": [
|
|
"## Directly benchmark models"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": 6,
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"tokenizer = self.agent.tokenizer\n",
|
|
"world_model = self.agent.world_model\n",
|
|
"actor_critic = self.agent.actor_critic\n"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": 7,
|
|
"metadata": {},
|
|
"outputs": [
|
|
{
|
|
"data": {
|
|
"text/plain": [
|
|
"4"
|
|
]
|
|
},
|
|
"execution_count": 7,
|
|
"metadata": {},
|
|
"output_type": "execute_result"
|
|
}
|
|
],
|
|
"source": [
|
|
"batch_num_samples = cfg.training.world_model.batch_num_samples\n",
|
|
"sequence_length = cfg.common.sequence_length\n",
|
|
"sample_from_start = False\n",
|
|
"# train_dataset = instantiate(cfg.datasets.train)\n",
|
|
"batch_num_samples\n"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": 8,
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"batch = self.train_dataset.sample_batch(batch_num_samples, sequence_length, sample_from_start)\n",
|
|
"batch = {k: v.to(self.device) for k, v in batch.items()}\n"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": 9,
|
|
"metadata": {},
|
|
"outputs": [
|
|
{
|
|
"name": "stdout",
|
|
"output_type": "stream",
|
|
"text": [
|
|
"CPU times: user 190 ms, sys: 4.87 ms, total: 194 ms\n",
|
|
"Wall time: 195 ms\n"
|
|
]
|
|
},
|
|
{
|
|
"data": {
|
|
"text/plain": [
|
|
"<src.utils.LossWithIntermediateLosses at 0x7f97ed202a30>"
|
|
]
|
|
},
|
|
"execution_count": 9,
|
|
"metadata": {},
|
|
"output_type": "execute_result"
|
|
}
|
|
],
|
|
"source": [
|
|
"%%time\n",
|
|
"self.agent.world_model.compute_loss(batch, tokenizer=self.agent.tokenizer)\n"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": 10,
|
|
"metadata": {},
|
|
"outputs": [
|
|
{
|
|
"name": "stdout",
|
|
"output_type": "stream",
|
|
"text": [
|
|
"CPU times: user 9.35 s, sys: 17.3 ms, total: 9.37 s\n",
|
|
"Wall time: 9.37 s\n"
|
|
]
|
|
},
|
|
{
|
|
"data": {
|
|
"text/plain": [
|
|
"<src.utils.LossWithIntermediateLosses at 0x7f97f274beb0>"
|
|
]
|
|
},
|
|
"execution_count": 10,
|
|
"metadata": {},
|
|
"output_type": "execute_result"
|
|
}
|
|
],
|
|
"source": [
|
|
"%%time\n",
|
|
"# TODO: why is this so slow?\n",
|
|
"cfg_actor_critic = self.cfg.training.actor_critic\n",
|
|
"self.agent.actor_critic.compute_loss(batch, tokenizer=self.agent.tokenizer, world_model=self.agent.world_model, **cfg_actor_critic)\n"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": 11,
|
|
"metadata": {},
|
|
"outputs": [
|
|
{
|
|
"name": "stdout",
|
|
"output_type": "stream",
|
|
"text": [
|
|
"CPU times: user 9.11 s, sys: 15.5 ms, total: 9.13 s\n",
|
|
"Wall time: 9.13 s\n"
|
|
]
|
|
}
|
|
],
|
|
"source": [
|
|
"%%time\n",
|
|
"# is this the slow part... yes. damn\n",
|
|
"actor_critic.imagine(batch, tokenizer, world_model, horizon=10);\n"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": 12,
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"# # takes 0.1 s, fast\n",
|
|
"# wm_env = WorldModelEnv(tokenizer, world_model, device)\n",
|
|
"# wm_env\n"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": 13,
|
|
"metadata": {},
|
|
"outputs": [
|
|
{
|
|
"name": "stdout",
|
|
"output_type": "stream",
|
|
"text": [
|
|
"torch.Size([4, 3, 64, 64])\n",
|
|
"CPU times: user 105 ms, sys: 207 µs, total: 105 ms\n",
|
|
"Wall time: 105 ms\n"
|
|
]
|
|
}
|
|
],
|
|
"source": [
|
|
"%%time\n",
|
|
"# this takes 0.1 seconds and is run 10+ time. So 1 second. Hmm\n",
|
|
"from src.envs.world_model_env import WorldModelEnv, Categorical\n",
|
|
"initial_observations = batch['observations']\n",
|
|
"\n",
|
|
"# get the right obs\n",
|
|
"wm_env = WorldModelEnv(self.agent.tokenizer, self.agent.world_model, self.device)\n",
|
|
"obs = wm_env.reset_from_initial_observations(initial_observations[:, -1])\n",
|
|
"print(obs.shape)\n",
|
|
"\n",
|
|
"\n",
|
|
"# make sure hidden states are right\n",
|
|
"self.agent.actor_critic.reset(obs.shape[0])\n",
|
|
"\n"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": 14,
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"import gc\n",
|
|
"gc.collect()\n",
|
|
"torch.cuda.empty_cache()\n",
|
|
"# obs\n"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": 15,
|
|
"metadata": {},
|
|
"outputs": [
|
|
{
|
|
"name": "stdout",
|
|
"output_type": "stream",
|
|
"text": [
|
|
"CPU times: user 1.6 ms, sys: 309 µs, total: 1.9 ms\n",
|
|
"Wall time: 1.72 ms\n"
|
|
]
|
|
}
|
|
],
|
|
"source": [
|
|
"%%time\n",
|
|
"# 700us\n",
|
|
"# fast, executed 10+ times\n",
|
|
"outputs_ac = actor_critic(obs)\n"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": 16,
|
|
"metadata": {},
|
|
"outputs": [
|
|
{
|
|
"data": {
|
|
"text/plain": [
|
|
"torch.Size([4, 1, 17])"
|
|
]
|
|
},
|
|
"execution_count": 16,
|
|
"metadata": {},
|
|
"output_type": "execute_result"
|
|
}
|
|
],
|
|
"source": [
|
|
"outputs_ac.logits_actions.shape\n",
|
|
"# action_token.shape\n"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": 17,
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"# %%timeit\n",
|
|
"# slow! takes 1s, executed 10+ times this is the culprit, not the lstm. hmm\n",
|
|
"k=3\n",
|
|
"horizon = 6\n",
|
|
"action_token = Categorical(logits=outputs_ac.logits_actions).sample()\n",
|
|
"obs, reward, done, _ = wm_env.step(action_token, should_predict_next_obs=(k < horizon - 1))\n"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": 18,
|
|
"metadata": {},
|
|
"outputs": [
|
|
{
|
|
"name": "stdout",
|
|
"output_type": "stream",
|
|
"text": [
|
|
"66.5 ms ± 1.53 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)\n"
|
|
]
|
|
}
|
|
],
|
|
"source": [
|
|
"%%timeit\n",
|
|
"# 62ms\n",
|
|
"# this is the slow part again. no grad and eval don't hepl\n",
|
|
"outputs_wm = world_model(action_token, past_keys_values=wm_env.keys_values_wm)\n"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": 19,
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"\n",
|
|
"num_steps=1\n",
|
|
"prev_steps=0\n",
|
|
"sequences = world_model.embedder(action_token, num_steps, prev_steps) + world_model.pos_emb(prev_steps + torch.arange(num_steps, device=action_token.device))\n"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": 20,
|
|
"metadata": {},
|
|
"outputs": [
|
|
{
|
|
"ename": "AssertionError",
|
|
"evalue": "",
|
|
"output_type": "error",
|
|
"traceback": [
|
|
"\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
|
|
"\u001b[0;31mAssertionError\u001b[0m Traceback (most recent call last)",
|
|
"\u001b[1;32m/media/wassname/SGIronWolf/projects5/worldmodels/iris_bigvae/notebooks/01_debug_models.ipynb Cell 24\u001b[0m line \u001b[0;36m1\n\u001b[0;32m----> <a href='vscode-notebook-cell:/media/wassname/SGIronWolf/projects5/worldmodels/iris_bigvae/notebooks/01_debug_models.ipynb#X64sZmlsZQ%3D%3D?line=0'>1</a>\u001b[0m get_ipython()\u001b[39m.\u001b[39;49mrun_cell_magic(\u001b[39m'\u001b[39;49m\u001b[39mtimeit\u001b[39;49m\u001b[39m'\u001b[39;49m, \u001b[39m'\u001b[39;49m\u001b[39m'\u001b[39;49m, \u001b[39m\"\u001b[39;49m\u001b[39m# ofc it\u001b[39;49m\u001b[39m'\u001b[39;49m\u001b[39ms the transformer that\u001b[39;49m\u001b[39m'\u001b[39;49m\u001b[39ms slow. I guess we just call it was more than during training\u001b[39;49m\u001b[39m\\n\u001b[39;49;00m\u001b[39mpast_keys_values = wm_env.keys_values_wm\u001b[39;49m\u001b[39m\\n\u001b[39;49;00m\u001b[39mx = world_model.transformer(sequences, past_keys_values)\u001b[39;49m\u001b[39m\\n\u001b[39;49;00m\u001b[39m\"\u001b[39;49m)\n",
|
|
"File \u001b[0;32m/media/wassname/SGIronWolf/projects5/worldmodels/iris_bigvae/.venv/lib/python3.9/site-packages/IPython/core/interactiveshell.py:2515\u001b[0m, in \u001b[0;36mInteractiveShell.run_cell_magic\u001b[0;34m(self, magic_name, line, cell)\u001b[0m\n\u001b[1;32m 2513\u001b[0m \u001b[39mwith\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39mbuiltin_trap:\n\u001b[1;32m 2514\u001b[0m args \u001b[39m=\u001b[39m (magic_arg_s, cell)\n\u001b[0;32m-> 2515\u001b[0m result \u001b[39m=\u001b[39m fn(\u001b[39m*\u001b[39;49margs, \u001b[39m*\u001b[39;49m\u001b[39m*\u001b[39;49mkwargs)\n\u001b[1;32m 2517\u001b[0m \u001b[39m# The code below prevents the output from being displayed\u001b[39;00m\n\u001b[1;32m 2518\u001b[0m \u001b[39m# when using magics with decorator @output_can_be_silenced\u001b[39;00m\n\u001b[1;32m 2519\u001b[0m \u001b[39m# when the last Python token in the expression is a ';'.\u001b[39;00m\n\u001b[1;32m 2520\u001b[0m \u001b[39mif\u001b[39;00m \u001b[39mgetattr\u001b[39m(fn, magic\u001b[39m.\u001b[39mMAGIC_OUTPUT_CAN_BE_SILENCED, \u001b[39mFalse\u001b[39;00m):\n",
|
|
"File \u001b[0;32m/media/wassname/SGIronWolf/projects5/worldmodels/iris_bigvae/.venv/lib/python3.9/site-packages/IPython/core/magics/execution.py:1189\u001b[0m, in \u001b[0;36mExecutionMagics.timeit\u001b[0;34m(self, line, cell, local_ns)\u001b[0m\n\u001b[1;32m 1186\u001b[0m \u001b[39mif\u001b[39;00m time_number \u001b[39m>\u001b[39m\u001b[39m=\u001b[39m \u001b[39m0.2\u001b[39m:\n\u001b[1;32m 1187\u001b[0m \u001b[39mbreak\u001b[39;00m\n\u001b[0;32m-> 1189\u001b[0m all_runs \u001b[39m=\u001b[39m timer\u001b[39m.\u001b[39;49mrepeat(repeat, number)\n\u001b[1;32m 1190\u001b[0m best \u001b[39m=\u001b[39m \u001b[39mmin\u001b[39m(all_runs) \u001b[39m/\u001b[39m number\n\u001b[1;32m 1191\u001b[0m worst \u001b[39m=\u001b[39m \u001b[39mmax\u001b[39m(all_runs) \u001b[39m/\u001b[39m number\n",
|
|
"File \u001b[0;32m~/miniforge3/lib/python3.9/timeit.py:205\u001b[0m, in \u001b[0;36mTimer.repeat\u001b[0;34m(self, repeat, number)\u001b[0m\n\u001b[1;32m 203\u001b[0m r \u001b[39m=\u001b[39m []\n\u001b[1;32m 204\u001b[0m \u001b[39mfor\u001b[39;00m i \u001b[39min\u001b[39;00m \u001b[39mrange\u001b[39m(repeat):\n\u001b[0;32m--> 205\u001b[0m t \u001b[39m=\u001b[39m \u001b[39mself\u001b[39;49m\u001b[39m.\u001b[39;49mtimeit(number)\n\u001b[1;32m 206\u001b[0m r\u001b[39m.\u001b[39mappend(t)\n\u001b[1;32m 207\u001b[0m \u001b[39mreturn\u001b[39;00m r\n",
|
|
"File \u001b[0;32m/media/wassname/SGIronWolf/projects5/worldmodels/iris_bigvae/.venv/lib/python3.9/site-packages/IPython/core/magics/execution.py:173\u001b[0m, in \u001b[0;36mTimer.timeit\u001b[0;34m(self, number)\u001b[0m\n\u001b[1;32m 171\u001b[0m gc\u001b[39m.\u001b[39mdisable()\n\u001b[1;32m 172\u001b[0m \u001b[39mtry\u001b[39;00m:\n\u001b[0;32m--> 173\u001b[0m timing \u001b[39m=\u001b[39m \u001b[39mself\u001b[39;49m\u001b[39m.\u001b[39;49minner(it, \u001b[39mself\u001b[39;49m\u001b[39m.\u001b[39;49mtimer)\n\u001b[1;32m 174\u001b[0m \u001b[39mfinally\u001b[39;00m:\n\u001b[1;32m 175\u001b[0m \u001b[39mif\u001b[39;00m gcold:\n",
|
|
"File \u001b[0;32m<magic-timeit>:3\u001b[0m, in \u001b[0;36minner\u001b[0;34m(_it, _timer)\u001b[0m\n",
|
|
"File \u001b[0;32m/media/wassname/SGIronWolf/projects5/worldmodels/iris_bigvae/.venv/lib/python3.9/site-packages/torch/nn/modules/module.py:1518\u001b[0m, in \u001b[0;36mModule._wrapped_call_impl\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m 1516\u001b[0m \u001b[39mreturn\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_compiled_call_impl(\u001b[39m*\u001b[39margs, \u001b[39m*\u001b[39m\u001b[39m*\u001b[39mkwargs) \u001b[39m# type: ignore[misc]\u001b[39;00m\n\u001b[1;32m 1517\u001b[0m \u001b[39melse\u001b[39;00m:\n\u001b[0;32m-> 1518\u001b[0m \u001b[39mreturn\u001b[39;00m \u001b[39mself\u001b[39;49m\u001b[39m.\u001b[39;49m_call_impl(\u001b[39m*\u001b[39;49margs, \u001b[39m*\u001b[39;49m\u001b[39m*\u001b[39;49mkwargs)\n",
|
|
"File \u001b[0;32m/media/wassname/SGIronWolf/projects5/worldmodels/iris_bigvae/.venv/lib/python3.9/site-packages/torch/nn/modules/module.py:1527\u001b[0m, in \u001b[0;36mModule._call_impl\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m 1522\u001b[0m \u001b[39m# If we don't have any hooks, we want to skip the rest of the logic in\u001b[39;00m\n\u001b[1;32m 1523\u001b[0m \u001b[39m# this function, and just call forward.\u001b[39;00m\n\u001b[1;32m 1524\u001b[0m \u001b[39mif\u001b[39;00m \u001b[39mnot\u001b[39;00m (\u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_backward_hooks \u001b[39mor\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_backward_pre_hooks \u001b[39mor\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_forward_hooks \u001b[39mor\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_forward_pre_hooks\n\u001b[1;32m 1525\u001b[0m \u001b[39mor\u001b[39;00m _global_backward_pre_hooks \u001b[39mor\u001b[39;00m _global_backward_hooks\n\u001b[1;32m 1526\u001b[0m \u001b[39mor\u001b[39;00m _global_forward_hooks \u001b[39mor\u001b[39;00m _global_forward_pre_hooks):\n\u001b[0;32m-> 1527\u001b[0m \u001b[39mreturn\u001b[39;00m forward_call(\u001b[39m*\u001b[39;49margs, \u001b[39m*\u001b[39;49m\u001b[39m*\u001b[39;49mkwargs)\n\u001b[1;32m 1529\u001b[0m \u001b[39mtry\u001b[39;00m:\n\u001b[1;32m 1530\u001b[0m result \u001b[39m=\u001b[39m \u001b[39mNone\u001b[39;00m\n",
|
|
"File \u001b[0;32m/media/wassname/SGIronWolf/projects5/worldmodels/iris_bigvae/src/models/transformer.py:69\u001b[0m, in \u001b[0;36mTransformer.forward\u001b[0;34m(self, sequences, past_keys_values)\u001b[0m\n\u001b[1;32m 66\u001b[0m \u001b[39m# k_size = (x.shape[0], x.shape[1], x.shape[1], 1)\u001b[39;00m\n\u001b[1;32m 67\u001b[0m \u001b[39m# v_size = past_keys_values[0]._v_cache._cache.size()\u001b[39;00m\n\u001b[1;32m 68\u001b[0m v_size \u001b[39m=\u001b[39m (k_size[\u001b[39m0\u001b[39m], k_size[\u001b[39m1\u001b[39m], x\u001b[39m.\u001b[39mshape[\u001b[39m1\u001b[39m], k_size[\u001b[39m3\u001b[39m])\n\u001b[0;32m---> 69\u001b[0m past_keys_values[\u001b[39m0\u001b[39;49m]\u001b[39m.\u001b[39;49mupdate(torch\u001b[39m.\u001b[39;49mrand(v_size), torch\u001b[39m.\u001b[39;49mrand(v_size))\n\u001b[1;32m 70\u001b[0m \u001b[39mreturn\u001b[39;00m x\n",
|
|
"File \u001b[0;32m/media/wassname/SGIronWolf/projects5/worldmodels/iris_bigvae/src/models/kv_caching.py:59\u001b[0m, in \u001b[0;36mKVCache.update\u001b[0;34m(self, k, v)\u001b[0m\n\u001b[1;32m 58\u001b[0m \u001b[39mdef\u001b[39;00m \u001b[39mupdate\u001b[39m(\u001b[39mself\u001b[39m, k: torch\u001b[39m.\u001b[39mTensor, v: torch\u001b[39m.\u001b[39mTensor):\n\u001b[0;32m---> 59\u001b[0m \u001b[39mself\u001b[39;49m\u001b[39m.\u001b[39;49m_k_cache\u001b[39m.\u001b[39;49mupdate(k)\n\u001b[1;32m 60\u001b[0m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_v_cache\u001b[39m.\u001b[39mupdate(v)\n",
|
|
"File \u001b[0;32m/media/wassname/SGIronWolf/projects5/worldmodels/iris_bigvae/src/models/kv_caching.py:33\u001b[0m, in \u001b[0;36mCache.update\u001b[0;34m(self, x)\u001b[0m\n\u001b[1;32m 31\u001b[0m \u001b[39mdef\u001b[39;00m \u001b[39mupdate\u001b[39m(\u001b[39mself\u001b[39m, x: torch\u001b[39m.\u001b[39mTensor) \u001b[39m-\u001b[39m\u001b[39m>\u001b[39m \u001b[39mNone\u001b[39;00m:\n\u001b[1;32m 32\u001b[0m \u001b[39massert\u001b[39;00m (x\u001b[39m.\u001b[39mndim \u001b[39m==\u001b[39m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_cache\u001b[39m.\u001b[39mndim) \u001b[39mand\u001b[39;00m \u001b[39mall\u001b[39m([x\u001b[39m.\u001b[39msize(i) \u001b[39m==\u001b[39m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_cache\u001b[39m.\u001b[39msize(i) \u001b[39mfor\u001b[39;00m i \u001b[39min\u001b[39;00m (\u001b[39m0\u001b[39m, \u001b[39m1\u001b[39m, \u001b[39m3\u001b[39m)])\n\u001b[0;32m---> 33\u001b[0m \u001b[39massert\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_size \u001b[39m+\u001b[39m x\u001b[39m.\u001b[39msize(\u001b[39m2\u001b[39m) \u001b[39m<\u001b[39m\u001b[39m=\u001b[39m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_cache\u001b[39m.\u001b[39mshape[\u001b[39m2\u001b[39m]\n\u001b[1;32m 34\u001b[0m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_cache \u001b[39m=\u001b[39m AssignWithoutInplaceCheck\u001b[39m.\u001b[39mapply(\u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_cache, x, \u001b[39m2\u001b[39m, \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_size, \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_size \u001b[39m+\u001b[39m x\u001b[39m.\u001b[39msize(\u001b[39m2\u001b[39m))\n\u001b[1;32m 35\u001b[0m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_size \u001b[39m+\u001b[39m\u001b[39m=\u001b[39m x\u001b[39m.\u001b[39msize(\u001b[39m2\u001b[39m)\n",
|
|
"\u001b[0;31mAssertionError\u001b[0m: "
|
|
]
|
|
}
|
|
],
|
|
"source": [
|
|
"%%timeit\n",
|
|
"# ofc it's the transformer that's slow. I guess we just call it was more than during training\n",
|
|
"past_keys_values = wm_env.keys_values_wm\n",
|
|
"x = world_model.transformer(sequences, past_keys_values)\n"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"# past_keys_values = wm_env.keys_values_wm\n",
|
|
"# x = world_model.transformer(sequences, past_keys_values)\n"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"%%timeit\n",
|
|
"# ofc it's the transformer that's slow. I guess we just call it was more than during training\n",
|
|
"past_keys_values = wm_env.keys_values_wm\n",
|
|
"x = world_model.transformer(sequences, past_keys_values)\n"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"logits_observations = world_model.head_observations(x, num_steps=num_steps, prev_steps=prev_steps)\n",
|
|
"logits_rewards = world_model.head_rewards(x, num_steps=num_steps, prev_steps=prev_steps)\n",
|
|
"logits_ends = world_model.head_ends(x, num_steps=num_steps, prev_steps=prev_steps)\n"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {},
|
|
"source": [
|
|
"## Torchinfo model sizes\n"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"observations = self.agent.tokenizer.preprocess_input(rearrange(batch['observations'], 'b t c h w -> (b t) c h w'))\n",
|
|
"# z, z_quantized, reconstructions = self.agent.tokenizer(observations, should_preprocess=False, should_postprocess=False)\n",
|
|
"summary(self.agent.tokenizer, input_data=observations)\n"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"\n",
|
|
"\n",
|
|
"with torch.no_grad():\n",
|
|
" obs_tokens = self.agent.tokenizer.encode(batch['observations'], should_preprocess=True).tokens # (BL, K)\n",
|
|
"\n",
|
|
"act_tokens = rearrange(batch['actions'], 'b l -> b l 1')\n",
|
|
"tokens = rearrange(torch.cat((obs_tokens, act_tokens), dim=2), 'b l k1 -> b (l k1)') # \n",
|
|
"\n",
|
|
"summary(self.agent.world_model, input_data=tokens)\n"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"from src.envs.world_model_env import WorldModelEnv\n",
|
|
"initial_observations = batch['observations']\n",
|
|
"\n",
|
|
"# get the right obs\n",
|
|
"wm_env = WorldModelEnv(self.agent.tokenizer, self.agent.world_model, self.device)\n",
|
|
"obs = wm_env.reset_from_initial_observations(initial_observations[:, -1])\n",
|
|
"obs.shape\n",
|
|
"\n",
|
|
"\n",
|
|
"# make sure hidden states are right\n",
|
|
"self.agent.actor_critic.reset(obs.shape[0])\n"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"\n",
|
|
"\n",
|
|
"from torchinfo import summary\n",
|
|
"summary(self.agent.actor_critic, input_data=obs)\n"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {},
|
|
"source": [
|
|
"## Debug env\n"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"import minihack\n",
|
|
"env = gym.make(\"MiniHack-River-v0\", observation_keys=(\"pixel_crop\", \"pixel\", 'blstats', 'message'))\n",
|
|
"env.reset() # each reset generates a new environment instance\n",
|
|
"obs, reward, end, info = env.step(1) # move agent '@' north\n",
|
|
"print(obs['pixel_crop'].shape)\n",
|
|
"plt.imshow(obs['pixel_crop'])\n",
|
|
"plt.show()\n",
|
|
"\n",
|
|
"print(obs['pixel'].shape)\n",
|
|
"plt.imshow(obs['pixel'])\n",
|
|
"\n"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"# # plt.imshow(obs['glyphs_crop'])\n",
|
|
"# obs['glyphs_crop'].shape\n",
|
|
"# obs['blstats']\n"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"import minihack\n",
|
|
"env = gym.make(\"MiniHack-Room-5x5-v0\", observation_keys=(\"pixel_crop\", \"pixel\", 'blstats', 'message'))\n",
|
|
"env.reset() # each reset generates a new environment instance\n",
|
|
"obs, reward, end, info = env.step(1) # move agent '@' north\n",
|
|
"print(obs['pixel_crop'].shape)\n",
|
|
"plt.imshow(obs['pixel_crop'])\n",
|
|
"plt.show()\n",
|
|
"\n",
|
|
"print(obs['pixel'].shape)\n",
|
|
"plt.imshow(obs['pixel'])\n",
|
|
"\n"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": []
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"import minihack\n",
|
|
"import crafter\n",
|
|
"env = gym.make(\"CrafterReward-v1\")\n",
|
|
"env.reset() # each reset generates a new environment instance\n",
|
|
"obs, reward, end, info = env.step(1) # move agent '@' north\n",
|
|
"print(obs.shape)\n",
|
|
"plt.imshow(obs)\n",
|
|
"plt.show()\n"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": []
|
|
}
|
|
],
|
|
"metadata": {
|
|
"kernelspec": {
|
|
"display_name": ".venv",
|
|
"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.16"
|
|
}
|
|
},
|
|
"nbformat": 4,
|
|
"nbformat_minor": 2
|
|
}
|