title: How initialize weights of a `torch.nn.Transformer` module?
body: <p>I am using a vanilla transformer architecture from the &quot;Attention Is All You Need&quot; paper for a sequence-to-sequence task. As shown in the following code.
  Assuming that I would like to use the <code>torch.nn.init.kaiming_uniform_</code> initialization method, how would one go about initializing the weights of the <code>nn.Transformer</code> ?
  Is it necessary to use a custom encoder and decoder class in order for that to happen?</p>
  <pre class="lang-py prettyprint-override"><code>import torch
  import torch.nn as nn
  import math
  
  # helper Module that adds positional encoding to the token embedding to introduce a notion of word order.
  import torch.nn as nn
  class PositionalEncoding(nn.Module):
      def __init__(self,
                   emb_size: int,
                   dropout: float,
                   maxlen: int = 20):
          super(PositionalEncoding, self).__init__()
  
          den = torch.exp(- torch.arange(0, emb_size, 2)* math.log(10000) / emb_size)
          pos = torch.arange(0, maxlen).reshape(maxlen, 1)
          pos_embedding = torch.zeros((maxlen, emb_size))
          pos_embedding[:, 0::2] = torch.sin(pos * den)
          pos_embedding[:, 1::2] = torch.cos(pos * den)
          pos_embedding = pos_embedding.unsqueeze(-2)
  
          self.dropout = nn.Dropout(dropout)
          self.register_buffer('pos_embedding', pos_embedding)
  
      def forward(self, token_embedding: Tensor):
          return self.dropout(token_embedding + self.pos_embedding[:token_embedding.size(0), :])
  
  # helper Module to convert tensor of input indices into corresponding tensor of token embeddings
  class TokenEmbedding(nn.Module):
      def __init__(self, vocab_size: int, emb_size):
          super(TokenEmbedding, self).__init__()
          self.embedding = nn.Embedding(vocab_size, emb_size)
          # Initialize weights with He initialization
          self.embedding.weight = nn.init.kaiming_uniform_(self.embedding.weight)
          self.emb_size = emb_size
  
      def forward(self, tokens: Tensor):
          return self.embedding(tokens.long()) * math.sqrt(self.emb_size)
  
  # Seq2Seq Network 
  class Transformer(nn.Module):
      def __init__(self,
                   src_vocab_size: int,
                   tgt_vocab_size: int,
                   num_encoder_layers: int = 1,
                   num_decoder_layers: int = 1,
                   emb_size: int = 300,
                   nhead: int = 3,
                   dim_feedforward: int = 512,
                   dropout: float = 0.1,
                   activation_function='relu'):
          super().__init__()
          self.src_tok_emb = TokenEmbedding(src_vocab_size, emb_size)
          self.tgt_tok_emb = TokenEmbedding(tgt_vocab_size, emb_size)
          self.positional_encoding = PositionalEncoding(
              emb_size, dropout=dropout)
          
          self.transformer = torch.nn.Transformer(d_model=emb_size,
                                         nhead=nhead,
                                         num_encoder_layers=num_encoder_layers,
                                         num_decoder_layers=num_decoder_layers,
                                         dim_feedforward=dim_feedforward,
                                         dropout=dropout,
                                         batch_first=True,
                                         activation=activation_function)
          self.generator = nn.Linear(emb_size, tgt_vocab_size)
      
      def init_weights(self):
        self.generator.weight = nn.init.kaiming_uniform_(self.generator.weight)
        
      def forward(self,
                  src: Tensor,
                  trg: Tensor,
                  src_mask: Tensor,
                  tgt_mask: Tensor,
                  src_padding_mask: Tensor,
                  tgt_padding_mask: Tensor,
                  memory_key_padding_mask: Tensor):
          # the .permute() is necessary since the positional-encoder expects tensors to be of shape
          # (seq_len, batch_size, emb_length)
          src_emb = self.positional_encoding(self.src_tok_emb(src).permute(1,0,2)).permute(1,0,2)
          tgt_emb = self.positional_encoding(self.tgt_tok_emb(trg).permute(1,0,2)).permute(1,0,2)
          outs = self.transformer(src_emb, tgt_emb, src_mask, tgt_mask, None, 
                                  src_padding_mask, tgt_padding_mask, memory_key_padding_mask)
          return self.generator(outs)
  
      def encode(self, src: Tensor, src_mask: Tensor):
          return self.transformer.encoder(self.positional_encoding(
                              self.src_tok_emb(src)), src_mask)
  
      def decode(self, tgt: Tensor, memory: Tensor, tgt_mask: Tensor):
          return self.transformer.decoder(self.positional_encoding(
                            self.tgt_tok_emb(tgt)), memory,
                            tgt_mask)
  </code></pre>