checkpoint

Author: Guanheng Zhang

beginner_source/transformer_tutorial.py

@@ -53,7 +53,6 @@ def __init__(self, ntoken, ninp, nhead, nhid, nlayers, dropout=0.5):
         super(TransformerModel, self).__init__()
         from torch.nn import TransformerEncoder, TransformerEncoderLayer
         self.model_type = 'Transformer'
-        self.src_mask = None
         self.pos_encoder = PositionalEncoding(ninp, dropout)
         encoder_layers = TransformerEncoderLayer(ninp, nhead, nhid, dropout)
         self.transformer_encoder = TransformerEncoder(encoder_layers, nlayers)
@@ -63,7 +62,7 @@ def __init__(self, ntoken, ninp, nhead, nhid, nlayers, dropout=0.5):
 
         self.init_weights()
 
-    def _generate_square_subsequent_mask(self, sz):
+    def generate_square_subsequent_mask(self, sz):
         mask = (torch.triu(torch.ones(sz, sz)) == 1).transpose(0, 1)
         mask = mask.float().masked_fill(mask == 0, float('-inf')).masked_fill(mask == 1, float(0.0))
         return mask
@@ -74,18 +73,12 @@ def init_weights(self):
         self.decoder.bias.data.zero_()
         self.decoder.weight.data.uniform_(-initrange, initrange)
 
-    def forward(self, src):
-        if self.src_mask is None or self.src_mask.size(0) != len(src):
-            device = src.device
-            mask = self._generate_square_subsequent_mask(len(src)).to(device)
-            self.src_mask = mask
-
+    def forward(self, src, src_mask):
         src = self.encoder(src) * math.sqrt(self.ninp)
         src = self.pos_encoder(src)
-        output = self.transformer_encoder(src, self.src_mask)
+        output = self.transformer_encoder(src, src_mask)
         output = self.decoder(output)
-        return F.log_softmax(output, dim=-1)
-
+        return output
 
 ######################################################################
 # ``PositionalEncoding`` module injects some information about the
@@ -113,7 +106,6 @@ def forward(self, x):
         x = x + self.pe[:x.size(0), :]
         return self.dropout(x)
 
-
 ######################################################################
 # Load data
 # ---------
@@ -200,15 +192,14 @@ def get_batch(batch_data):
 # equal to the length of the vocab object.
 # 
 
-ntokens = len(TEXT.vocab.stoi) # the size of vocabulary
+ntokens = len(vocab.stoi) # the size of vocabulary
 emsize = 200 # embedding dimension
 nhid = 200 # the dimension of the feedforward network model in nn.TransformerEncoder
 nlayers = 2 # the number of nn.TransformerEncoderLayer in nn.TransformerEncoder
-nhead = 2 # the number of heads in the multiheadattention models 
+nhead = 2 # the number of heads in the multiheadattention models
 dropout = 0.2 # the dropout value
 model = TransformerModel(ntokens, emsize, nhead, nhid, nlayers, dropout).to(device)
 
-
 ######################################################################
 # Run the model
 # -------------