wrap

Author: tcapelle

model.py

@@ -222,6 +222,11 @@ def decode(self, encoder_output: torch.Tensor, src_mask: torch.Tensor, tgt: torc
     def project(self, x):
         # (batch, seq_len, vocab_size)
         return self.projection_layer(x)
+
+    def forward(self, encoder_output, encoder_mask, decoder_input, decoder_mask):      
+        encoder_output = self.encode(encoder_input, encoder_mask) # (B, seq_len, d_model)
+        decoder_output = self.decode(encoder_output, encoder_mask, decoder_input, decoder_mask) # (B, seq_len, d_model)
+        return self.project(decoder_output) # (B, seq_len, vocab_size))
 
 def build_transformer(src_vocab_size: int, tgt_vocab_size: int, src_seq_len: int, tgt_seq_len: int, d_model: int=512, N: int=6, h: int=8, dropout: float=0.1, d_ff: int=2048) -> Transformer:
     # Create the embedding layers

train.py

@@ -261,11 +261,12 @@ def train_model(config: ModelConfig):
             encoder_mask = batch['encoder_mask'].to(device) # (B, 1, 1, seq_len)
             decoder_mask = batch['decoder_mask'].to(device) # (B, 1, seq_len, seq_len)
 
-            # Run the tensors through the encoder, decoder and the projection layer
-            encoder_output = model.module.encode(encoder_input, encoder_mask) # (B, seq_len, d_model)
-            decoder_output = model.module.decode(encoder_output, encoder_mask, decoder_input, decoder_mask) # (B, seq_len, d_model)
-            proj_output = model.module.project(decoder_output) # (B, seq_len, vocab_size)
-
+            # # Run the tensors through the encoder, decoder and the projection layer
+            # encoder_output = model.module.encode(encoder_input, encoder_mask) # (B, seq_len, d_model)
+            # decoder_output = model.module.decode(encoder_output, encoder_mask, decoder_input, decoder_mask) # (B, seq_len, d_model)
+            # proj_output = model.module.project(decoder_output) # (B, seq_len, vocab_size)
+            proj_output = model(encoder_output, encoder_mask, decoder_input, decoder_mask)
+            
             # Compare the output with the label
             label = batch['label'].to(device) # (B, seq_len)