File size: 1,498 Bytes
a03c9b4 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 |
# Copyright 2024 The YourMT3 Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Please see the details in the LICENSE file.
"""lm_head.py"""
import torch
from torch import nn
from typing import Optional, Dict
class LMHead(nn.Module):
"""Language Model Head with tied weights."""
def __init__(self, decoder_config: Dict, init_factor: float = 1.0, tie_word_embeddings: bool = True):
super().__init__()
self.d_model = decoder_config["d_model"]
self.init_factor = init_factor
self.tie_word_embeddings = tie_word_embeddings
self.lm_head = nn.Linear(decoder_config["d_model"], decoder_config["vocab_size"], bias=False)
self._init_weights()
def _init_weights(self):
if self.tie_word_embeddings is False:
self.lm_head.weight.data.normal_(mean=0.0, std=self.init_factor * 1.0)
def forward(self, decoder_hs: torch.FloatTensor) -> torch.FloatTensor:
if self.tie_word_embeddings is True:
# Rescale output before projecting on vocab
# See https://github.com/tensorflow/mesh/blob/fa19d69eafc9a482aff0b59ddd96b025c0cb207d/mesh_tensorflow/transformer/transformer.py#L586
decoder_hs = decoder_hs * (self.d_model**-0.5)
lm_logits = self.lm_head(decoder_hs)
return lm_logits
|