(self, num_embeddings, embedding_dim,
init_method=init.xavier_normal_)
| 79 | init_method: method to initialize weights. |
| 80 | """ |
| 81 | def __init__(self, num_embeddings, embedding_dim, |
| 82 | init_method=init.xavier_normal_): |
| 83 | super(VocabParallelEmbedding, self).__init__() |
| 84 | # Keep the input dimensions. |
| 85 | self.num_embeddings = num_embeddings |
| 86 | self.embedding_dim = embedding_dim |
| 87 | # Set the detauls for compatibility. |
| 88 | self.padding_idx = None |
| 89 | self.max_norm = None |
| 90 | self.norm_type = 2. |
| 91 | self.scale_grad_by_freq = False |
| 92 | self.sparse = False |
| 93 | self._weight = None |
| 94 | # Divide the weight matrix along the vocaburaly dimension. |
| 95 | self.vocab_start_index, self.vocab_end_index = \ |
| 96 | VocabUtility.vocab_range_from_global_vocab_size( |
| 97 | self.num_embeddings, get_model_parallel_rank(), |
| 98 | get_model_parallel_world_size()) |
| 99 | self.num_embeddings_per_partition = self.vocab_end_index - \ |
| 100 | self.vocab_start_index |
| 101 | |
| 102 | # Allocate weights. |
| 103 | self.weight = Parameter(torch.Tensor(self.num_embeddings_per_partition, |
| 104 | self.embedding_dim)) |
| 105 | self.weight.model_parallel = True |
| 106 | # And initialize. |
| 107 | _initialize_affine_weight( |
| 108 | self.weight, self.num_embeddings, self.embedding_dim, |
| 109 | self.num_embeddings_per_partition, 0, init_method) |
| 110 | |
| 111 | def forward(self, input_): |
| 112 | # Build the mask. |
nothing calls this directly
no test coverage detected