The user wants a concise, one-line Git commit message based on the pr…

…ovided diff. The diff shows changes to the `TimeSeriesTransformer` class in `src/ppo_model.py`. Specifically, it adds `norm_first=True` to `TransformerEncoderLayer` and a `LayerNorm` to `TransformerEncoder`. These changes are aimed at improving the stability of the transformer. Therefore, the commit type should be `fix` as it addresses a bug (numerical instability). The description should reflect the addition of layer normalization within the transformer.fix: Add LayerNorm within transformer for stability
cove9988 · Jan 5, 2025 · e29bfb1 · e29bfb1
1 parent 48dad37
commit e29bfb1
Showing 1 changed file with 4 additions and 8 deletions.
diff --git a/src/ppo_model.py b/src/ppo_model.py
@@ -46,16 +46,15 @@ def __init__(self, input_size, embed_dim, num_heads, num_layers, dropout=0.1):
         encoder_layer = nn.TransformerEncoderLayer(
             d_model=embed_dim,
             nhead=num_heads,
-            dropout=dropout
+            dropout=dropout,
+            norm_first=True  # Apply LayerNorm before attention and feedforward
         )
         self.transformer_encoder = nn.TransformerEncoder(
             encoder_layer,
-            num_layers=num_layers
+            num_layers=num_layers,
+            norm=nn.LayerNorm(embed_dim) # Add LayerNorm at the end of the encoder
         )
 
-        # Layer normalization
-        self.layer_norm = nn.LayerNorm(embed_dim)
-
         # Decoder layer to produce final output
         self.decoder = nn.Linear(embed_dim, embed_dim)
 
@@ -66,9 +65,6 @@ def forward(self, src):
         # Pass through the transformer encoder
         output = self.transformer_encoder(src)
 
-        # Apply layer normalization
-        output = self.layer_norm(output)
-
         # Pass through the decoder layer
         output = self.decoder(output)