import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.autograd import Variable
import math
import matplotlib.pyplot as plt
import numpy as np
import copy
import warnings
warnings.filterwarnings("ignore")

class Embedding(nn.Module):
    # embedding層
    def __init__(self, d_model, vocab):
        # d_model: 詞嵌入的維度（轉換後獲得的詞向量的維度）
        # vocab：詞表 的大小
        super(Embedding, self).__init__()
        # 定義Embedding層
        self.lut = nn.Embedding(vocab, d_model)
        # 將參數傳入類中
        self.d_model = d_model
    
    def forward(self, x):
        # x：代表輸入進模型的文本通過辭彙映射後的數字張量
        return self.lut(x) * math.sqrt(self.d_model)  # 陳尚d_model的開根號是為了對最後結果進行一個縮放

class PositionalEncoding(nn.Module):
    def __init__(self, d_model, dropout, max_len=5000):
        """
        位置編碼器類的初始化函數
        d_model：詞嵌入維度
        dropout：置0比率
        max_len：每個句子的最大長度
        """
        super(PositionalEncoding, self).__init__()
        
        # 實例化nn中預定義的Dropout層， 並將dropout傳入其中，獲得對象self.dropout
        self.dropout = nn.Dropout(p=dropout)
        
        # 初始化一個位置編碼矩陣，他是一個0矩陣，矩陣的大小是  max_len * d_model  
        pe = torch.zeros(max_len, d_model)
        
        # 初始化一個絕對位置矩陣，在我們這裡，辭彙的絕對位置就是用它的索引去表示
        # 所以我們首先使用arange方法獲得一個連續自然數向量，然後在使用unsqueeze方法拓展矩陣維度
        # 又因為參數傳的是1， 代表矩陣拓展的位置，會使向量編程一個max_len * 1的矩陣。
        position = torch.arange(0, max_len).unsqueeze(1)
        # 絕對位置矩陣初始化之後，接下來就是考慮如何將這個位置資訊加入到位置編碼矩陣中
        # 最簡單方法就是現將max_len * 1的絕對位置矩陣，變換成max_len * d_model形狀, 然後直接覆蓋原來的初始化位置編碼矩陣即可。
        # 要實現這種矩陣變化，就需要一個1 * d_model形狀的變化矩陣div_term，我們對這個變化矩陣的要求除了形狀外，
        # 還希望他能夠將自然數的絕對位置編碼縮放層足夠小的數字，有助於在之後的梯度下降過程中更快地收斂，這樣我們就可以開始初始化
        # 得到一個自然數矩陣，但是我們這裡並沒有按預計的一樣的初始化一個1 * d_model的矩陣
        # 只初始化了一半，即x*d_model/2的矩陣， 這是因為這裡並不是真正意義的初始化了一半的矩陣
        # 我們可以把它看作初始化了兩次，而兩次初始化的變化矩陣不同的處理
        # 並把這兩個矩陣分別填充在位置編碼矩陣的偶數和奇數位置上，組成最終的位置編碼矩陣。
        
        div_term = torch.exp(torch.arange(0, d_model, 2) * -(math.log(1000.0)/ d_model))
        pe[:, 0::2] = torch.sin(position * div_term)
        pe[:, 1::2] = torch.cos(position * div_term)
        # 這樣我們就得到了位置編碼矩陣pe， pe現在還只是一個二維矩陣，要想和embeddding的輸出（一個三維矩陣）
        # 就必須坨鎮一個維度，所以這裡使用unsequeeze拓展維度
        pe = pe.unsqueeze(0)  # 將二維蟑螂擴充為三維張量
        # 最後把pe位置編碼矩陣註冊層模型的buffer，什麼是buffer呢
        # 我們把它認為是對模型效果有幫助的，但是卻不是模型結構中超參數或者參數，不需要說著優化步驟進行更新的增益
        # 註冊之後我們就可以在模型保存後重新載入時盒模型結構與參數已通被載入
        self.register_buffer('pe',  pe)
        
    def forward(self, x):
        """
        forward函數的桉樹是x，表示文本序列的詞嵌入表示
        """
        # 在相加之前我們對pe做一些適配工作，將這個三維張量的第二位也就是句子最大長度的那一維度切片與輸入的x的第二維相同即x.size(1)
        # 因為我們默認max_len為5000，一般來講是在太大了，很難有一個句子包含5000辭彙，所以要進行與輸入轉給你來那個的適配
        # 最後使用Variable進行封裝，使其與x的樣式相同，但是他是不需要進行梯度求解的，因為這裡的位置資訊使用的是函數計算方式，且這一步中只是數字相加，因此把requuires_grad設置成False。
        
        x = x + Variable(self.pe[:, :x.size(1)], requires_grad=False)
        # 最後就使用self.dropout對象驚喜『丟棄』操作，並返回結果
        return self.dropout(x)

def subsequent_mask(size):
    """
    生成向後遮掩的掩碼張量，參數size是掩碼張量最後兩個維度的大小，最後兩維形成一個方陣
    """
    # 在函數中，首先定義掩碼張量的形狀
    attn_shape = (1, size, size)
    # 然後使用np.ones方法想這個形狀中添加1元素，形成上三角陣，最後為了節約空間
    # 再使用其中的數據類型變為無符號8為整型unit8
    sub_mask = np.triu(np.ones(attn_shape), k=1).astype('uint8')
    # 最後將numpy類型轉化為torch中的tensor，內部做一個1 - 的操作
    # 在這個其實是做了一個三角陣的反轉，subsequent_mask中的每個元素都會被1減，
    # 如果是0， subsequent_mask中的該位置由0變1
    # 如果是1， subsequent_mask中的該位置由1變0
    return torch.from_numpy(1 - sub_mask)

def attention(query, key, value, mask=None, dropout=None):
    """
    注意力機制的實現，輸入分別是query，key，value，mask：掩碼張量，
    dropout是nn.Dropout層的實例化對象，默認為None
    """
    # 在函數中，首先去query的最後一維的大小，一般情況下就等同於我們的詞嵌入維度，命名為d_k
    d_k = query.size(-1)
    # 按照注意力公式，將query與key的轉置相乘，這裡面key是將最後兩個維度進行轉置（轉置後才滿足矩陣乘法），key轉置之後shape為[1, 4, 6]，[1, 6, 4] * [1, 4, 6] = [1, 6, 6]
    # 再除以縮放係數，就得到注意力得分張量scores
    scores = torch.matmul(query, key.transpose(-2, -1)) // math.sqrt(d_k)
    # 接著判斷是否使用掩碼張量
    if mask is not None:
        # 使用tensor的masked_fill方法，將掩碼張量和scores張量每個位置一一比較，如果掩碼張量
        # 則對應的scores張量用-1e9這個值來替換
        scores = scores.masked_fill(mask == 0, -1e9)
    # 對scores的最後一維進行softmax操作，使用F.softmax方法，第一個參數softmax對象，第二個
    # 這樣獲得最終的注意力張量
    p_attn = F.softmax(scores, dim=-1)
    # 之後判斷是否使用dropout進行隨機置零
    if dropout is not None:
        # 將p_attn傳入dropout對象中進行「丟棄」處理
        p_attn = dropout(p_attn)
    #最後，根據公司將p_atten與value張量向曾獲得最終的query注意力表示，同時返回注意力權重張量
    return torch.matmul(p_attn, value), p_attn

query = key = value = pe_result

attn, p_attn = attention(query, key, value)

# 首先需要定義克隆函數，因為在多頭注意力機制的實現中，用到多個結構相同的線形層
# 我們將使用clone函數將他們已通初始化在同一個網路層列表對象中，之後的結構中也會用到該函數
def clone(model, N):
    # 用於生成相同網路層的克隆函數，它的參數module表示要克隆的目標網路層，N代表需要克隆的數量
    # 在函數中，我們通過for循環對module進行N次深度拷貝，使其每個module稱為獨立的層
    return nn.ModuleList([copy.deepcopy(model) for _ in range(N)])

# 我們使用一個類來實現多頭注意力機制的處理
class MultiHeadAttention(nn.Module):
    def __init__(self, head, embedding_dim, dropout=0.1):
        """
        在類的初始化時，會傳入三個參數，head代表頭數，embedding_dim代表詞嵌入的維度，dropout代表進行dropout操作時置零比率，默認是0.1
        """
        super(MultiHeadAttention, self).__init__()
        # 在函數中，首先使用了一個測試中常用的assert語句，判斷h是否能被d_model整除
        # 這是因為我們之後要給每個頭分配等量的詞特徵，也就是embedding_dim//head個
        assert embedding_dim % head == 0
        self.head = head  # 傳入頭數
        self.embedding_dim = embedding_dim
        self.dropout = nn.Dropout(p=dropout)
        self.d_k = embedding_dim // head  # 得到每個頭獲得的分割詞向量維度d_k
        # 然後獲得線形層對象，通過nn的Linear實例化，它的內部變化矩陣是embedding_dim * Embedding_dim
        self.linears = clone(nn.Linear(embedding_dim, embedding_dim), 4)
        self.attn = None
    
    def forward(self, query, key, value, mask=None):
        """前向邏輯函數，它的輸入參數有4個，前三個就是注意力機制需要的Q、K、V，
        最後一個是注意力機制中可能需要的mask掩碼張量，默認是None"""
        if mask is not None:  # 如果存在掩碼張量
            mask = mask.unsqueeze(0)  # 使用unsqueeze拓展維度，代表多頭中的第n頭
        # 接著，我們獲得一個batch_size的變數，它是query尺寸的第1個數字，代表有多少條樣本
        batch_size = query.size(0)
        # 之後就進入多頭處理環節
        # 首先利用zip將輸入QKV與三個線形層組到一起，然後使用for循環，將輸入QKV分別傳入到線形層中
        # 完成線性變換後，開始為每個頭分割輸入，這裡使用view方法對線性變化的結果進行維度重塑
        # 這樣就意味著每個頭可以獲得一部分詞特徵組成的句子，其中的-1代表自適應維度
        # 電腦會根據這種變換自動計算這裡的值，然後對第二維和第三維進行轉置操作
        # 為了讓代表句子長度維度和詞向量維度能夠相鄰，這樣注意力機制才能找到遲疑與句子位置的關係
        # 從attention函數中可以看到，利用的是原始輸入的倒數第一和第二，這樣我們就得到了每個頭的
        print('----------------------------------')
        print('query-before transpose: ', query.shape)
        print('key-before transpose: ', key.shape)
        print('value-before transpose: ', value.shape)
        # 此時，query, key, value的shape為[1, 6, 4]， [batch size, sequence length, embedding dimension]
        query, key, value = \
               [model(x).view(batch_size, -1, self.head, self.d_k).transpose(1, 2) 
                for model, x in zip(self.linears, (query, key, value))]
        print('----------------------------------')
        print('query-after transpose: ', query.shape)
        print('key-after transpose: ', key.shape)
        print('value-after transpose: ', value.shape)
        print('----------------------------------')
        # 此時，query, key, value的shape修改為：[1, 2, 6, 2] ， [batch size, sequence length, h, embedding dimension / h]
        # 得到每個頭的輸入後，接下來就是將他們傳入attention中
        # 這裡直接調用我們之前實現的attention函數，同時也將mask和dropout傳入其中
        x, self.attn = attention(query, key, value, mask, self.dropout)
        print('x- after attention: ', x.shape)  # 經過注意力機制後，輸出x的shape為：[1, 2, 6, 2]， [batch size, h, sequence length, embedding dimension / h]
        # 通過多頭注意力計算後，我們就得到了每個頭計算結果組成的4維張量，我們需要將其轉換為輸入的
        # 因此這裡開始進行第一步處理環節的逆操作，先對第二和第三維驚喜轉置，然後使用contiguous方法
        # 這個方法的作用就是能夠讓轉置後的張量應用view方法，否則將無法直接使用
        # 所以，下一步就是使用view方法重塑形狀，變成和輸入形狀相同
        x = x.transpose(1, 2).contiguous().view(batch_size, -1, self.head * self.d_k)  # 對x進行reshape
        print('x', x.shape)  # x的shape回到[1, 6, 4]，  [batch size, sequence length, embedding dimension]
        return self.linears[-1](x)

class PositionalwiseFeedForward(nn.Module):
    def __init__(self, d_model, d_ff, dropout=0.1):
        """
        d_model: 線形層的輸入維，因為我們希望輸入通過前饋全連接層後輸入和輸出的維度不變
        d_ff: 第二個線形層
        droupout：置零比率
        """
        super(PositionalwiseFeedForward, self).__init__()
        # 首先按照我們預期使用nn實例化了兩個線形層對象
        # 他們的參數分別是d_model， d_ff和d_ff， d_model
        self.w1 = nn.Linear(d_model, d_ff)
        self.w2 = nn.Linear(d_ff, d_model)
        # 然後我們使用nn的Dropout實例化了對象self.dropout
        self.dropout = nn.Dropout(p=dropout)
    
    def forward(self, x):
        """輸入參數為x，代表來自上一層的輸出"""
        # 首先經過第一個線性層，然後使用Funtional中relu函數進行激活,
        # 之後再使用dropout進行隨機置0，最後通過第二個線性層w2，返回最終結果.
        return self.w2(self.dropout(F.relu(self.w1(x))))

class LayerNorm(nn.Module):
    def __init__(self, features, eps=1e-6):
        # features：詞嵌入的維度
        # eps：一個足夠小的數，在規範化公式的分母中出現，房子分母為0。
        super(LayerNorm, self).__init__()
        # 根據features的形狀初始化兩個參數張量a2和b2，a2為元素全為1的張量，b2為元素全為0的張量。這兩個張量就是規範化層的參數
        # 因為直接對上一層得到的結果做規範化公式計算，將改變結果的正常表徵，因此就需要有參數作為調節因子
        # 使其既能滿足規範化要求，又能不改變針對目標的表徵，最後使用nn.parameter封裝，代表它們是模型的參數
        self.eps = eps
        self.a2 = nn.Parameter(torch.ones(features))
        self.b2 = nn.Parameter(torch.zeros(features))
        
    def forward(self, x):
        # x來自於上一層的輸出
        # 在函數中，首先對輸入變數x求其最後一個維度的均值，並保持輸出維度與輸入維度一致
        # 接著再求最後一個維度的標準差，然後就是根據規範化公式，用x減去均值除以標準差獲得規範化的結果
        # 最後對結果乘以我們的縮放茶樹，即a2， *號代表同型點乘，即對應位置進行乘法操作，加上位移參數b2返回即可
        # 
        mean = x.mean(-1, keepdim=True)
        std = x.std(-1, keepdim=True)
        return self.a2 * (x - mean) / (std + self.eps) + self.b2

class SubLayerConnection(nn.Module):
    def __init__(self, size, dropout=0.1):
        # size: 詞嵌入的維度大小
        # dropout：隨機置零比率
        super(SubLayerConnection, self).__init__()
        self.size = size
        self.dropout = nn.Dropout(p=dropout)
        # 實例化了規範化對象self.norm
        self.norm = LayerNorm(size)
    
    def forward(self, x, sublayer):
        # 前向邏輯函數中，接收上一層或者子層的輸入作為第一個參數
        # 將該子層連接中的子層函數作為第二個函數
        # 我們首先對輸出進行規範化，然後將結果傳給子層處理，之後再對子層進行dropout操作
        # 隨機停止一些網路中神經元的作用，來房子過擬合，最後還有一個add操作
        # 因為存在跳躍連接，所以是將輸入x與dropout後的子層輸出結果相加作為最終的子層連接輸出
        return x + self.dropout(sublayer(self.norm(x)))

class EncoderLayer(nn.Module):
    def __init__(self, size, self_attn, feed_forward, dropout):
        # size：詞嵌入維度的大小
        # self_attn：傳入多頭注意力子層實例化對象，並且是自注意力機制
        # feed_forward：前饋全連接層實例化對象
        # dropout：置零比率
        super(EncoderLayer, self).__init__()
        self.size = size
        # 首先將self_attn和feed_forward傳入其中
        self.self_attn = self_attn
        self.feed_forward = feed_forward
        self.dropout = nn.Dropout(p=dropout)
        # 編碼器層中有兩個子層連接結構，所以使用clones函數進行克隆
        self.sublayer = clone(SubLayerConnection(size, dropout), 2)
    
    def forward(self, x, mask):
        # x：上一層輸出
        # mask：掩碼張量
        # 裡面就是按照結構圖左側的流程，首先通過第一個子層連接結構，其中包含多頭注意力子層，然後通過第二個子層連接結構，
        #其中包含前饋全連接子層，最後返回結果
        x = self.sublayer[0](x, lambda x: self.self_attn(x, x, x, mask))
        return self.sublayer[1](x, self.feed_forward)

class Encoder(nn.Module):
    def __init__(self, layer, N):
        # layer：編碼器層
        # N: 編碼器層的個數
        super(Encoder, self).__init__()
        # 首先使用clones函數克隆N個編碼器層放在self.layers中
        self.layers = clone(layer, N)
        # 再初始化一個規範化層，將用在編碼器的最後面
        self.norm = LayerNorm(layer.size)

    def forward(self, x, mask):
        # forward函數的輸入和編碼器層相同，x代表上一層的輸出，mask代表掩碼張量
        # 首先就是對我們克隆的編碼器層進行循環，每次都會得到一個新的x
        # 這個循環的過程，就相當於輸出的x經過N個編碼器層的處理
        # 最後在通過規範化層的對象self.norm進行處理，最後返回結果
        for layer in self.layers:
            x = layer(x, mask)
        return self.norm(x)

class DecoderLayer(nn.Module):
    def __init__(self, size, self_attn, src_attn, feed_forward, dropout=0.1):
        # size：詞嵌入的維度大小，同時也代表解碼器層的尺寸
        # self_attn：多頭注意力對象，也就是說這個注意力機制需要Q=K=V
        # src_attn：多頭注意力對象，這裡Q!=K=V
        # feed_forward：前饋全連接層對象
        super(DecoderLayer, self).__init__()
        self.size = size
        self.self_attn = self_attn
        self.src_attn = src_attn
        self.feed_forward = feed_forward
        self.dropout = nn.Dropout(p=dropout)
        # 按照結構圖使用clones函數克隆三個子層連接對象
        self.sublayer = clone(SubLayerConnection(size, dropout), 3)
    
    def forward(self, x, memory, source_mask, target_mask):
        # x：來自上一層的輸出
        # mermory：來自編碼器層的語義存儲變數
        # 源數據掩碼張量和目標數據掩碼張量
        m = memory
        # 將x傳入第一個子層結構，第一個子層結構分別是x和self-attn函數，因為是自注意力機制，所以Q,K,V都是x
        # 最後一個參數的目標數據掩碼張量，這是要對目標數據進行遮掩，因為此時模型可能還沒有生成任何目標數據
        # 比如紮起解碼器準備生成第一個字元或者辭彙時，我們其實已經傳入了第一個字元以便計算損失
        # 但是我們不希望在生成第一個字元時模型還能利用這個資訊，因此我們會將其遮掩，同樣生成第二個字元或辭彙時
        # 模型只能使用第一個字元或者辭彙資訊，第二個字元以及之後的資訊都不允許模型使用
        x = self.sublayer[0](x, lambda x: self.self_attn(x, x, x, target_mask))
        # 接著進入第二個子層，這個職稱中常規的注意力機制，q是輸入x；k,v是編碼層輸出的memory
        # 同樣也傳入source_mask，但是進行源數據遮掩的原因並非是抑制資訊泄露，而是這筆掉對結果沒有意義的字元而陳勝的注意力值
        # 以此提升模型效果和訓練速度，這樣就完成了第二個子層的處理
        x = self.sublayer[1](x, lambda x: self.src_attn(x, m, m, source_mask))
        # 最後一個子層就是前饋全連接子層，經過它的處理後，就可以返回結果，這就是我們的解碼器層結構。
        return self.sublayer[2](x, self.feed_forward)

class Decoder(nn.Module):
    def __init__(self, layer, N):
        # layer：解碼器層layer
        # N：解碼器層的個數N
        super(Decoder, self).__init__()
        # 首先使用clones方法克隆了N個layer，然後實例化了一個規範化層
        # 因為數據走過了所有的解碼器層後，最後要做規範化處理
        self.N = N
        self.layers = clone(layer, N)
        self.norm = LayerNorm(layer.size)
    
    def forward(self, x, memory, source_mask, target_mask):
        # x：數據的嵌入表示
        # memory：編碼器層的輸出
        # source_mask：源數據掩碼張量
        # target_mask：目標數據掩碼張量
        
        # 對每個層進行循環，淡然這個循環就是變數x通過每一個層的處理
        # 得出最後的結果，再進行一次規範化返回即可
        for layer in self.layers:
            x = layer(x, memory, source_mask, target_mask)
        return self.norm(x)

# nn.functional工具包裝載了網路層中那些只驚喜年計算，而沒有參數的層
# 將線形層和softmax計算層一起實現，因為二者的共同目標是生成最後的結構
# 因此把類的名字叫做Generator，生成器類
class Generator(nn.Module):
    def __init__(self, d_model, vocab_size):
        # d_model：詞嵌入維度
        # vocab_size：詞表大小
        super(Generator, self).__init__()
        # 首先就是使用nn中的預定義線形層進行實例化，得到一個對象self.project等待使用
        # 這個線性層的參數有兩個，計時初始化函數時傳進來的 d_model和vocab_size
        self.vocab_size = vocab_size
        self.embedding_dim = embedding_dim
        self.project = nn.Linear(d_model, vocab_size)
    
    def forward(self, x):
        # 前向邏輯函數中輸入是上一層輸出張量x
        # 在函數中，首先使用上一步得到的self.project對x進行線性變化
        # 然後使用F中已經實現的log_softmax進行的softmax處理
        return F.log_softmax(self.project(x), dim=-1)

# 編碼器-解碼器結構實現
class EncoderDecoder(nn.Module):
    def __init__(self, encoder, decoder, source_embed, target_embed, generator):
        # encoder：編碼器對象
        # decoder：解碼器對象
        # source_embed：源數據嵌入函數
        # target_embed：目標數據嵌入函數
        # generator：輸出部分的類別生成器對象
        super(EncoderDecoder, self).__init__()
        # 將參數傳入類中
        self.encoder = encoder
        self.decoder = decoder
        self.src_embed = source_embed
        self.tgt_embed = target_embed
        self.generator = generator
    
    def encode(self, source, source_mask):
        # 使用src_embed對source做處理，然後和source_mask一起傳給self.encoder
        return self.encoder(self.src_embed(source), source_mask)
    
    def decode(self, memory, source_mask, target, target_mask):
        # 解碼函數
        # 使用tgt_embed對target做處理，然後和source_mask、target_mask、memory一起傳給self.decoder
        return self.decoder(self.tgt_embed(target), memory, source_mask, 
                            target_mask)
    
    def forward(self, source, target, source_mask, target_mask):
        # source：源數據
        # target：目標數據
        # source_mask、target_mask：對應的掩碼張量
        
        # 在函數中，將source、source_mask傳入編碼函數，得到結果後，與source-mask、target、target_mask一同傳給解碼函數
        return self.decode(self.encode(source, source_mask), source_mask, 
                           target, target_mask)

def make_model(source_vocab, target_vocab, N=6, d_model=512,
               d_ff=2048, head=8, dropout=0.1):
    # source_vocab：源數據特徵（辭彙）總數
    # target_vocab：目標數據特徵（辭彙）總數
    # N：編碼器和解碼器堆疊數
    # d_model：詞向量映射維度
    # d_ff：前饋全連接網路中變化矩陣的維度
    # head：多頭注意力機制中多頭數
    # dropout：置零比率
    
    # 首先得到一個深度拷貝命令，接下來很多結構都要進行深度拷貝，從而保證它們彼此之間相互獨立，不受干擾
    c = copy.deepcopy
    # 實例化了多頭注意力機制類，
    attn = MultiHeadAttention(head, d_model, dropout)
    # 然後實例化前饋全連接類
    ff = PositionalwiseFeedForward(d_model, d_ff, dropout)
    # 實例化位置編碼器類
    position = PositionalEncoding(d_model, dropout)
    # 根據結構圖，最外層是EncoderDecode，在EncoderDecoder中
    # 分別是編碼器層，解碼器層，源數據Embedding層和位置編碼組成的有序結構
    # 目標數據Embedding層和位置編碼組成的有序結構，以及類別生成器層
    # 在編碼器層中有attention子層以及前饋全連接子層
    # 在解碼器層中有兩個attention子層以及前饋全連接層
    model = EncoderDecoder(
            Encoder(EncoderLayer(d_model, c(attn), c(ff), dropout), N),
            Decoder(DecoderLayer(d_model, c(attn), c(attn), c(ff), dropout), N),
            nn.Sequential(Embedding(d_model, source_vocab), c(position)),
            nn.Sequential(Embedding(d_model, target_vocab), c(position)),
            Generator(d_model, target_vocab))
    # 模型構建完成後，接下來就是初始化模型中的參數，比如線形層的變化矩陣
    # 這裡一旦判斷參數的維度大於1，則會將其初始化成為一個服從均勻分布的矩陣
    for p in model.parameters():
        if p.dim() > 1:
            nn.init.xavier_uniform(p)
    return model