Merge pull request #6 from pierrot-lc/transformer-only

configs/exp/hard.yaml

@@ -6,12 +6,10 @@ env:
   batch_size: 10000
 
 model:
-  n_channels: 16
-  embedding_dim: 12
-  n_heads: 1
-  backbone_cnn_layers: 6
-  backbone_transformer_layers: 3
-  decoder_layers: 3
+  embedding_dim: 64
+  n_heads: 2
+  backbone_transformer_layers: 6
+  decoder_layers: 2
   dropout: 0.05
 
 optimizer:
@@ -21,7 +19,8 @@ optimizer:
 
 scheduler:
   warmup_steps: 0
-  cosine_tmax: 0
+  cosine_t0: 0
+  cosine_tmult: 0
 
 loss:
   gamma: 0.99
@@ -33,7 +32,7 @@ loss:
 
 trainer:
   clip_value: 1.0
-  scramble_size: 0.20
+  scramble_size: 0.50
 
 iterations:
   rollouts: 100

configs/exp/normal.yaml

@@ -3,42 +3,41 @@ group: Normal
 
 env:
   path: ./instances/eternity_A.txt
-  batch_size: 5000
+  batch_size: 10000
 
 model:
-  n_channels: 16
-  embedding_dim: 12
-  n_heads: 1
-  backbone_cnn_layers: 5
-  backbone_transformer_layers: 1
+  embedding_dim: 64
+  n_heads: 4
+  backbone_transformer_layers: 6
   decoder_layers: 2
   dropout: 0.05
 
 optimizer:
   optimizer: adamw
-  learning_rate: 1.0e-3
+  learning_rate: 1.0e-4
   weight_decay: 1.0e-3
 
 scheduler:
   warmup_steps: 0
-  cosine_tmax: 0
+  cosine_t0: 0
+  cosine_tmult: 0
 
 loss:
   gamma: 0.99
   gae_lambda: 0.95
   ppo_clip_ac: 0.30
   ppo_clip_vf: 0.30
   value_weight: 1.0e-3
-  entropy_weight: 5.0e-4
+  entropy_weight: 1.0e-3
 
 trainer:
   clip_value: 1.0
-  scramble_size: 0.20
+  scramble_size: 0.50
 
 iterations:
   rollouts: 60
   epochs: -1
   batches: 60
-  batch_size: 2500
+  batch_size: 5000
 
 checkpoint:

configs/exp/trivial.yaml

@@ -6,22 +6,21 @@ env:
   batch_size: 5000
 
 model:
-  n_channels: 6
-  embedding_dim: 12
-  n_heads: 1
-  backbone_cnn_layers: 4
-  backbone_transformer_layers: 1
+  embedding_dim: 24
+  n_heads: 2
+  backbone_transformer_layers: 4
   decoder_layers: 1
   dropout: 0.05
 
 optimizer:
-  optimizer: lion
+  optimizer: adamw
   learning_rate: 1.0e-3
   weight_decay: 1.0e-4
 
 scheduler:
   warmup_steps: 0
-  cosine_tmax: 0
+  cosine_t0: 0
+  cosine_tmult: 0
 
 loss:
   gamma: 0.99
@@ -33,7 +32,7 @@ loss:
 
 trainer:
   clip_value: 1.0
-  scramble_size: 0.25
+  scramble_size: 0.50
 
 iterations:
   rollouts: 10

configs/exp/trivial_B.yaml

@@ -3,25 +3,24 @@ group: Trivial B
 
 env:
   path: ./instances/eternity_trivial_B.txt
-  batch_size: 5000
+  batch_size: 10000
 
 model:
-  n_channels: 9
-  embedding_dim: 16
+  embedding_dim: 24
   n_heads: 2
-  backbone_cnn_layers: 4
-  backbone_transformer_layers: 1
+  backbone_transformer_layers: 4
   decoder_layers: 1
   dropout: 0.05
 
 optimizer:
   learning_rate: 1.0e-3
-  optimizer: lion
+  optimizer: adamw
   weight_decay: 1.0e-4
 
 scheduler:
   warmup_steps: 0
-  cosine_tmax: 0
+  cosine_t0: 0
+  cosine_tmult: 0
 
 loss:
   gamma: 0.99
@@ -33,12 +32,12 @@ loss:
 
 trainer:
   clip_value: 1.0
-  scramble_size: 0.10
+  scramble_size: 0.50
 
 iterations:
   rollouts: 40
   epochs: 100
   batches: 40
-  batch_size: 2500
+  batch_size: 5000
 
 checkpoint:

main.py

@@ -47,13 +47,10 @@ def init_model(config: DictConfig, env: EternityEnv) -> Policy:
     """Initialize the model."""
     model = config.exp.model
     return Policy(
-        n_classes=env.n_classes,
         board_width=env.board_size,
         board_height=env.board_size,
-        n_channels=model.n_channels,
         embedding_dim=model.embedding_dim,
         n_heads=model.n_heads,
-        backbone_cnn_layers=model.backbone_cnn_layers,
         backbone_transformer_layers=model.backbone_transformer_layers,
         decoder_layers=model.decoder_layers,
         dropout=model.dropout,
@@ -107,10 +104,11 @@ def init_scheduler(
         )
         schedulers.append(warmup_scheduler)
 
-    if scheduler.cosine_tmax > 0:
-        lr_scheduler = optim.lr_scheduler.CosineAnnealingLR(
+    if scheduler.cosine_t0 > 0:
+        lr_scheduler = optim.lr_scheduler.CosineAnnealingWarmRestarts(
             optimizer=optimizer,
-            T_max=scheduler.cosine_tmax,
+            T_0=scheduler.cosine_t0,
+            T_mult=scheduler.cosine_tmult,
         )
         schedulers.append(lr_scheduler)
 
@@ -173,6 +171,7 @@ def reload_checkpoint(config: DictConfig, trainer: Trainer):
     trainer.model.load_state_dict(state_dict["model"])
     # HACK: The training seems to not be stable when loading the optimizer state.
     # trainer.optimizer.load_state_dict(state_dict["optimizer"])
+    trainer.scheduler.load_state_dict(state_dict["scheduler"])
     print(f"Checkpoint from {checkpoint_path} loaded.")
 
 

src/model/backbone.py

@@ -2,9 +2,10 @@
 import torch.nn as nn
 from einops import rearrange
 from einops.layers.torch import Rearrange
-from positional_encodings.torch_encodings import PositionalEncoding2D
+from positional_encodings.torch_encodings import PositionalEncoding2D, Summer
 
 from ..environment import N_SIDES
+from .class_encoding import ClassEncoding
 
 
 class Backbone(nn.Module):
@@ -24,47 +25,25 @@ class Backbone(nn.Module):
 
     def __init__(
         self,
-        n_classes: int,
-        n_channels: int,
         embedding_dim: int,
         n_heads: int,
-        cnn_layers: int,
         transformer_layers: int,
         dropout: float,
     ):
         super().__init__()
 
         self.embed_board = nn.Sequential(
-            # Embed the classes of each size of the tiles.
-            Rearrange("b t h w -> b h w t"),
-            nn.Embedding(n_classes, n_channels),
+            # Encode the classes.
+            ClassEncoding(embedding_dim),
             # Merge the classes of each tile into a single embedding.
-            Rearrange("b h w t e -> b (t e) h w"),
-            nn.Conv2d(N_SIDES * n_channels, n_channels, kernel_size=1, padding="same"),
-            nn.GELU(),
-            nn.GroupNorm(1, n_channels),
+            Rearrange("b t h w e -> b h w (t e)"),
+            nn.Linear(N_SIDES * embedding_dim, embedding_dim),
+            # Add the 2D positional encodings.
+            Summer(PositionalEncoding2D(embedding_dim)),
+            # To transformer layout.
+            Rearrange("b h w e -> (h w) b e"),
         )
-        self.linear = nn.Sequential(
-            nn.Linear(n_channels, embedding_dim),
-            nn.LayerNorm(embedding_dim),
-        )
-        self.positional_enc = PositionalEncoding2D(embedding_dim)
 
-        self.cnn_layers = nn.ModuleList(
-            [
-                nn.Sequential(
-                    nn.Conv2d(
-                        n_channels,
-                        n_channels,
-                        kernel_size=3,
-                        padding="same",
-                    ),
-                    nn.GELU(),
-                    nn.GroupNorm(1, n_channels),
-                )
-                for _ in range(cnn_layers)
-            ]
-        )
         self.transformer_layers = nn.TransformerEncoder(
             nn.TransformerEncoderLayer(
                 d_model=embedding_dim,
@@ -92,16 +71,6 @@ def forward(
             The embedded game state.
                 Shape of [board_height x board_width, batch_size, embedding_dim].
         """
-        batch_size, _, board_height, board_width = tiles.shape
-
-        tiles = self.embed_board(tiles)
-        for layer in self.cnn_layers:
-            tiles = layer(tiles) + tiles
-
-        tokens = rearrange(tiles, "b e h w -> b h w e")
-        tokens = self.linear(tokens)
-        tokens = self.positional_enc(tokens) + tokens
-        tokens = rearrange(tokens, "b h w e -> (h w) b e")
+        tokens = self.embed_board(tiles)
         tokens = self.transformer_layers(tokens)
-
         return tokens

src/model/class_encoding.py

@@ -42,13 +42,13 @@ def forward(self, board: torch.Tensor) -> torch.Tensor:
         ---
         Args:
             board: The timesteps of the game states.
-                Shape of [batch_size, board_height, board_width, N_SIDES].
+                Shape of [...].
         ---
         Returns:
             The board with encoded classes.
-                Shape of [batch_size, board_height, board_width, N_SIDES, embedding_dim].
+                Shape of [..., embedding_dim].
         """
-        assert board.max() <= self.embedding_dim, "Not enough orthogonal vectors!"
+        assert board.max() < self.embedding_dim, "Not enough orthogonal vectors!"
 
         # Encode the classes of each tile.
         return self.class_enc(board)

src/model/model.py

@@ -16,13 +16,10 @@
 class Policy(nn.Module):
     def __init__(
         self,
-        n_classes: int,
         board_width: int,
         board_height: int,
-        n_channels: int,
         embedding_dim: int,
         n_heads: int,
-        backbone_cnn_layers: int,
         backbone_transformer_layers: int,
         decoder_layers: int,
         dropout: float,
@@ -33,11 +30,8 @@ def __init__(
         self.embedding_dim = embedding_dim
 
         self.backbone = Backbone(
-            n_classes,
-            n_channels,
             embedding_dim,
             n_heads,
-            backbone_cnn_layers,
             backbone_transformer_layers,
             dropout,
         )

src/policy_gradient/trainer.py

@@ -207,10 +207,12 @@ def evaluate(self) -> dict[str, Any]:
     def save_model(self, filepath: Path | str):
         model_state = self.model.state_dict()
         optimizer_state = self.optimizer.state_dict()
+        scheduler_state = self.scheduler.state_dict()
         torch.save(
             {
                 "model": model_state,
                 "optimizer": optimizer_state,
+                "scheduler": scheduler_state,
             },
             filepath,
         )