diff --git a/config/trainer.yaml b/config/trainer.yaml index 81f5e88..f8022a0 100644 --- a/config/trainer.yaml +++ b/config/trainer.yaml @@ -54,7 +54,7 @@ training: should: True learning_rate: 0.0001 tokenizer: - batch_num_samples: 32 + batch_num_samples: 128 grad_acc_steps: 1 max_grad_norm: 10.0 start_after_epochs: 5 @@ -71,7 +71,7 @@ training: grad_acc_steps: 1 max_grad_norm: 10.0 start_after_epochs: 50 - steps_per_epoch: 200 + steps_per_epoch: 20 imagine_horizon: ${common.sequence_length} burn_in: 20 gamma: 0.995 diff --git a/config/world_model/default.yaml b/config/world_model/default.yaml index cb9fcd3..16769b7 100644 --- a/config/world_model/default.yaml +++ b/config/world_model/default.yaml @@ -1,6 +1,6 @@ _target_: models.TransformerConfig tokens_per_block: 17 -max_blocks: 20 +max_blocks: 10 attention: 'causal' num_layers: 10 num_heads: 4 diff --git a/img/2023-11-16-13-01-11.png b/img/2023-11-16-13-01-11.png new file mode 100644 index 0000000..00f3a57 Binary files /dev/null and b/img/2023-11-16-13-01-11.png differ diff --git a/research_journal.md b/research_journal.md index bfd295d..de33a90 100644 --- a/research_journal.md +++ b/research_journal.md @@ -163,3 +163,39 @@ hm maybe it's just the face it has to backprop throguh the whole LLM :( damn... well running eval on the transformer brought it down from 100sec to 60, but it's still huge. But then why is the model training fast? It makes not sense + +# 2023-11-16 12:54:48 + +Why is agent so slow? Lets find out +- look at diagram +- look at train_agent + - to tokenizer.compute_loss is just tokenizer + - world_model.compute_loss user tokenizer with no grad + - actor_critic? takes an hour!! + - imagine (with grad?) + - x20 = horizon + - self(obs) + - WorldModelEnv.step this has no grad! + - transformer + - tokenizer with no grad + - compute_lambda_returns with no grad + + +So changes: +- the world model step always had no grad! +- I just made the lstm smaller and the horizon smaller +- from 1h to 3m. Reasonable. + + +Experiment: +- try no grad on the model? ok it now takes 20 minutes to train... still slow +with a smaller lstm and only 10 steos ut tajes 8 mins, + +![](img/2023-11-16-13-01-11.png) + +Experience collection (train_dataset): 100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 200/200 [00:05<00:00, 36.11it/s] +Training tokenizer: 100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 200/200 [00:55<00:00, 3.61it/s] +Training world_model: 100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 200/200 [01:12<00:00, 2.76it/s] +Training tokenizer: 55sec +Training world_model 72 sec +train actor_critic 3min. It looks like it scales with lstm size! diff --git a/src/envs/world_model_env.py b/src/envs/world_model_env.py index 6188d2c..73e3852 100644 --- a/src/envs/world_model_env.py +++ b/src/envs/world_model_env.py @@ -65,13 +65,12 @@ class WorldModelEnv: token = action.clone().detach() if isinstance(action, torch.Tensor) else torch.tensor(action, dtype=torch.long) token = token.reshape(-1, 1).to(self.device) # (B, 1) + for k in range(num_passes): # assumption that there is only one action token. outputs_wm = self.world_model(token, past_keys_values=self.keys_values_wm) output_sequence.append(outputs_wm.output_sequence) - - # if outputs_wm.logits_rewards.shape[1] > 0: if k == 0: reward = Categorical(logits=outputs_wm.logits_rewards).sample().float().cpu().numpy().reshape(-1) - 1 # (B,) done = Categorical(logits=outputs_wm.logits_ends).sample().cpu().numpy().astype(bool).reshape(-1) # (B,) diff --git a/src/models/actor_critic.py b/src/models/actor_critic.py index 69dc3de..0fff2b6 100644 --- a/src/models/actor_critic.py +++ b/src/models/actor_critic.py @@ -36,6 +36,7 @@ class ImagineOutput: class ActorCritic(nn.Module): def __init__(self, act_vocab_size, use_original_obs: bool = False) -> None: super().__init__() + shrink = 4 self.use_original_obs = use_original_obs self.conv1 = nn.Conv2d(3, 32, 3, stride=1, padding=1) self.maxp1 = nn.MaxPool2d(2, 2) @@ -43,15 +44,15 @@ class ActorCritic(nn.Module): self.maxp2 = nn.MaxPool2d(2, 2) self.conv3 = nn.Conv2d(32, 64, 3, stride=1, padding=1) self.maxp3 = nn.MaxPool2d(2, 2) - self.conv4 = nn.Conv2d(64, 64, 3, stride=1, padding=1) + self.conv4 = nn.Conv2d(64, 64//shrink, 3, stride=1, padding=1) self.maxp4 = nn.MaxPool2d(2, 2) - self.lstm_dim = 512 - self.lstm = nn.LSTMCell(1024, self.lstm_dim) + self.lstm_dim = 64 + self.lstm = nn.LSTMCell(1024//shrink, self.lstm_dim) self.hx, self.cx = None, None - self.critic_linear = nn.Linear(512, 1) - self.actor_linear = nn.Linear(512, act_vocab_size) + self.critic_linear = nn.Linear(self.lstm_dim, 1) + self.actor_linear = nn.Linear(self.lstm_dim, act_vocab_size) def __repr__(self) -> str: return "actor_critic" @@ -85,7 +86,7 @@ class ActorCritic(nn.Module): x = F.relu(self.maxp2(self.conv2(x))) x = F.relu(self.maxp3(self.conv3(x))) x = F.relu(self.maxp4(self.conv4(x))) - x = torch.flatten(x, start_dim=1) + x = torch.flatten(x, start_dim=1) # [b=32, 64//shrink, 4, 4] if mask_padding is None: self.hx, self.cx = self.lstm(x, (self.hx, self.cx)) diff --git a/src/trainer.py b/src/trainer.py index ec94572..8b53aa5 100644 --- a/src/trainer.py +++ b/src/trainer.py @@ -169,6 +169,7 @@ class Trainer: if max_grad_norm is not None: torch.nn.utils.clip_grad_norm_(component.parameters(), max_grad_norm) + optimizer.step() metrics = {f'{str(component)}/train/total_loss': loss_total_epoch, **intermediate_losses}