network.py

from neuron import *

class Transformer:
    def __init__(self, number_of_layers: int,
                 model_dimension: int,
                 number_of_heads: int,
                 batch_size: int,
                 maximum_length: int,
                 output_vocabulary_size: int,
                 input_vocabulary_size: int,
                 learning_rate: float,
                 need_SGDR: bool=False,
                 warmup_steps: int=50,
                 learning_rate_min: float=0.0,
                 learning_rate_max: float=0.0,
                 cycle_duration: int=0):
        self.encoder_layers = []
        self.decoder_layers = []

        for i in range(number_of_layers):
            self.encoder_layers.append(TransformerEncoder(model_dimension, batch_size, number_of_heads, learning_rate, True, 0.1))
            self.decoder_layers.append(TransformerDecoder(model_dimension, batch_size, number_of_heads, learning_rate, True, 0.1))
        
        self.positional_encoding_encoder = TrainablePositionalEncoding(maximum_length, model_dimension, utils.Adam(learning_rate), batch_size)#SinusoidalPositionalEncoding(maximum_length, model_dimension, batch_size)
        self.positional_encoding_decoder = TrainablePositionalEncoding(maximum_length, model_dimension, utils.Adam(learning_rate), batch_size)#SinusoidalPositionalEncoding(maximum_length, model_dimension, batch_size)
        self.embedding_dropout_encoder = DropoutLayer(0.1)
        self.embedding_dropout_decoder = DropoutLayer(0.1)

        self.output_projection_layer = LinearLayer(model_dimension, output_vocabulary_size, 'none', 'xavier_normal',
                                                    utils.Adam(learning_rate), batch_size, False)
        self.softmax = SoftmaxLayer(output_vocabulary_size, batch_size)

        self.encoder_embedding = EmbeddingLayer(input_vocabulary_size, model_dimension, utils.Adam(learning_rate), batch_size)
        self.decoder_embedding = EmbeddingLayer(output_vocabulary_size, model_dimension, utils.Adam(learning_rate), batch_size)

        self.loss_layer = LossLayer(output_vocabulary_size, batch_size, 'cross_entropy', 'average')
        self.loss = 0.0

        self.SGDR_counter = 0
        self.SGDR_learning_rate_min = learning_rate_min
        self.SGDR_learning_rate_max = learning_rate_max
        self.SDGR = need_SGDR
        self.SGDR_cycle = cycle_duration

        self.step_count = 0
        self.warmup_steps = warmup_steps

        self.learning_rate = learning_rate


    def forward(self, input: np.typing.NDArray, target_decoder: np.typing.NDArray, target_loss: np.typing.NDArray, 
                input_masking: np.typing.NDArray, target_masking: np.typing.NDArray, loss_masking: np.typing.NDArray):
        #约定：masking会mask掉为1的地方
        embedded_input = self.encoder_embedding.forward(input)
        embedded_input += self.positional_encoding_encoder.forward(input)
        embedded_input = self.embedding_dropout_encoder.forward(embedded_input)

        encoder_output = embedded_input
        for each_encoder in self.encoder_layers:
            encoder_output = each_encoder.forward(encoder_output, input_masking)
        
        embedded_target = self.decoder_embedding.forward(target_decoder)
        embedded_target += self.positional_encoding_decoder.forward(target_decoder)
        embedded_target = self.embedding_dropout_decoder.forward(embedded_target)

        decoder_output = embedded_target
        
        for each_decoder in self.decoder_layers:
            decoder_output = each_decoder.forward(decoder_output, encoder_output, target_masking, input_masking)

        logits = self.output_projection_layer.forward(decoder_output)
        probabilities = self.softmax.forward(logits)
        '''
        #------------------
                # target_masking 中 1 表示填充，0 表示有效
                loss_mask = target_masking.copy()  # 复制一份，避免修改原始mask
                batch_size, seq_len = loss_mask.shape
                for i in range(batch_size):
                    # 找到当前样本所有有效位置（值为0）的索引
                    valid_indices = np.where(loss_mask[i] == 0)[0]
                    if len(valid_indices) > 0:
                        last_valid_idx = valid_indices[-1]   # 最后一个有效位置
                        loss_mask[i, last_valid_idx] = 1     # 将其设为填充（1），使其不参与损失计算
        #-----------------'''
        # 损失计算时，有效位置标记为 1 - loss_mask

        #计算loss用的是target_masking，这里要1-
        #现在是1 -> 有 0 -> 无
        self.loss = self.loss_layer.forward(probabilities, target_loss, 1 - loss_masking)

        return probabilities
    
    def backward(self):
        loss_gradient = self.loss_layer.backward()

        logits_gradient = self.output_projection_layer.backward(loss_gradient)

        decoder_gradient = logits_gradient

        encoder_decoder_gradient = []

        for each_decoder_reversed in reversed(self.decoder_layers):
            decoder_gradient, dL_dX_ = each_decoder_reversed.backward(decoder_gradient)
            encoder_decoder_gradient.append(dL_dX_)

        dL_dX_ = np.zeros_like(encoder_decoder_gradient[0])

        for each in encoder_decoder_gradient:
            dL_dX_ += each

        embedding_gradient_decoder = self.embedding_dropout_decoder.backward(decoder_gradient)
        self.decoder_embedding.backward(embedding_gradient_decoder)
        self.embedding_dropout_decoder.backward(embedding_gradient_decoder)

        last = True
        encoder_gradient = dL_dX_
        for each_encoder_reversed in reversed(self.encoder_layers):
            if last:
                encoder_gradient = each_encoder_reversed.backward(encoder_gradient)
            else:
                encoder_gradient = each_encoder_reversed.backward(encoder_gradient)
            last = False

        embedding_gradient_encoder = self.embedding_dropout_encoder.backward(encoder_gradient)
        self.encoder_embedding.backward(embedding_gradient_encoder)
        self.embedding_dropout_encoder.backward(embedding_gradient_encoder)
        

    def update(self):
        self.step_count += 1
        if self.SDGR:
            self.SGDR_update()
        self.loss_layer.update()
        self.softmax.update()
        self.output_projection_layer.update()
        for each in self.decoder_layers:
            each.update()

        self.decoder_embedding.update()

        for each in self.encoder_layers:
            each.update()
        
        self.encoder_embedding.update()

    def set_learning_rate(self, learning_rate: float):
        for i in self.encoder_layers:
            i.set_learning_rate(learning_rate)

        for i in self.decoder_layers:
            i.set_learning_rate(learning_rate)

        self.encoder_embedding.optimizer.learning_rate = learning_rate
        self.decoder_embedding.optimizer.learning_rate = learning_rate
        self.positional_encoding_decoder.optimizer.learning_rate = learning_rate
        self.positional_encoding_encoder.optimizer.learning_rate = learning_rate

    def SGDR_update(self):
        learning_rate = self.SGDR_learning_rate_min + 0.5 * (self.SGDR_learning_rate_max - self.SGDR_learning_rate_min) * \
            (1 - math.cos(math.pi * (self.SGDR_counter % (2 * self.SGDR_cycle)) / self.SGDR_cycle))
        
        if self.step_count <= self.warmup_steps:
            learning_rate *= 1e-2

        self.set_learning_rate(learning_rate)
        self.SGDR_counter += 1

        self.learning_rate = learning_rate

    def set_batch_size(self, batch_size: int):
        self.output_projection_layer.batch_size = batch_size
        self.encoder_embedding.batch_size = batch_size
        self.decoder_embedding.batch_size = batch_size
        self.positional_encoding_decoder.batch_size = batch_size
        self.positional_encoding_encoder.batch_size = batch_size
        for i in self.encoder_layers:
            i.set_batch_size(batch_size)

        for i in self.decoder_layers:
            i.set_batch_size(batch_size)
    
    def switch(self, mode: Literal['training', 'reasoning']):
        for i in self.encoder_layers:
            i.switch(mode)

        for i in self.decoder_layers:
            i.switch(mode)

        self.embedding_dropout_decoder.switch(mode)
        self.embedding_dropout_encoder.switch(mode)