Construct the embeddings from word, position and token_type embeddings.
| 285 | |
| 286 | |
| 287 | class BertEmbeddings(nn.Module): |
| 288 | """Construct the embeddings from word, position and token_type embeddings. |
| 289 | """ |
| 290 | |
| 291 | def __init__(self, config): |
| 292 | super(BertEmbeddings, self).__init__() |
| 293 | self.word_embeddings = nn.Embedding(config.vocab_size, config.hidden_size) |
| 294 | # self.word_embeddings = mpu.VocabParallelEmbedding( |
| 295 | # config.vocab_size, config.hidden_size, |
| 296 | # init_method=normal_init_method(mean=0.0, |
| 297 | # std=config.initializer_range)) |
| 298 | self.position_embeddings = nn.Embedding(config.max_position_embeddings, config.hidden_size) |
| 299 | self.token_type_embeddings = nn.Embedding(config.type_vocab_size, config.hidden_size) |
| 300 | |
| 301 | # self.LayerNorm is not snake-cased to stick with TensorFlow model variable name and be able to load |
| 302 | # any TensorFlow checkpoint file |
| 303 | self.fp32_layernorm = config.fp32_layernorm |
| 304 | self.fp32_embedding = config.fp32_embedding |
| 305 | self.fp32_tokentypes = config.fp32_tokentypes |
| 306 | self.LayerNorm = BertLayerNorm(config.hidden_size, eps=config.layernorm_epsilon) |
| 307 | self.dropout = nn.Dropout(config.hidden_dropout_prob) |
| 308 | |
| 309 | def forward(self, input_ids, token_type_ids=None): |
| 310 | seq_length = input_ids.size(1) |
| 311 | position_ids = torch.arange(seq_length, dtype=torch.long, device=input_ids.device) |
| 312 | position_ids = position_ids.unsqueeze(0).expand_as(input_ids) |
| 313 | if token_type_ids is None: |
| 314 | token_type_ids = torch.zeros_like(input_ids) |
| 315 | |
| 316 | words_embeddings = self.word_embeddings(input_ids) |
| 317 | position_embeddings = self.position_embeddings(position_ids) |
| 318 | token_type_embeddings = self.token_type_embeddings(token_type_ids) |
| 319 | if not self.fp32_tokentypes: |
| 320 | |
| 321 | embeddings = words_embeddings + position_embeddings + token_type_embeddings |
| 322 | if self.fp32_embedding and not self.fp32_layernorm: |
| 323 | embeddings = embeddings.half() |
| 324 | previous_type = embeddings.type() |
| 325 | if self.fp32_layernorm: |
| 326 | embeddings = embeddings.float() |
| 327 | embeddings = self.LayerNorm(embeddings) |
| 328 | if self.fp32_layernorm: |
| 329 | if self.fp32_embedding: |
| 330 | embeddings = embeddings.half() |
| 331 | else: |
| 332 | embeddings = embeddings.type(previous_type) |
| 333 | else: |
| 334 | embeddings = words_embeddings.float() + position_embeddings.float() + token_type_embeddings.float() |
| 335 | if self.fp32_tokentypes and not self.fp32_layernorm: |
| 336 | embeddings = embeddings.half() |
| 337 | previous_type = embeddings.type() |
| 338 | if self.fp32_layernorm: |
| 339 | embeddings = embeddings.float() |
| 340 | embeddings = self.LayerNorm(embeddings) |
| 341 | if self.fp32_layernorm: |
| 342 | if self.fp32_tokentypes: |
| 343 | embeddings = embeddings.half() |
| 344 | else: |