本節內容的源代碼倉庫。
我們在前面的部分搭建了一個多層感知機字符級的語言模型，現在是時候把它的結構變得更複雜了。現在的目標是，輸入序列能夠輸入更多字符，而不是現在的 3 個。除此之外，我們不想把它們都放到一個隱藏層中，避免壓縮太多信息。這樣得到一個類似WaveNet的更深的模型。
WaveNet#
發表於 2016 年，基本上也是一種語言模型，只不過預測對象是音頻序列，而不是字符級或單詞級的序列。但從根本上說建模設置是相同的 —— 都是自回歸模型 (Autoregressive Model)，試圖預測序列中的下一個字符。
論文中使用了這種樹狀的層次結構來預測，本節將實現這個模型。
nn.Module#
把上節的內容封裝到類中，模仿 PyTorch 中 nn.Module 的 API。這樣可以把 “Linear”、“1 維 Batch Norm” 和 “Tanh” 這些模塊想象成樂高積木塊，然後用這些積木堆出神經網絡：
class Linear:
def __init__(self, fan_in, fan_out, bias=True):
self.weight = torch.randn((fan_in, fan_out), generator=g) / fan_in**0.5
self.bias = torch.zeros(fan_out) if bias else None
def __call__(self, x):
self.out = x @ self.weight
if self.bias is not None:
self.out += self.bias
return self.out
def parameters(self):
return [self.weight] + ([] if self.bias is None else [self.bias])
Linear，線性層這個模塊的作用就是 forward pass 的過程中做一個矩陣乘法。
class BatchNorm1d:
def __init__(self, dim, eps=1e-5, momentum=0.1):
self.eps = eps
self.momentum = momentum
self.training = True
# 使用反向傳播訓練的參數
self.gamma = torch.ones(dim)
self.beta = torch.zeros(dim)
# 使用“動量更新”進行訓練的緩衝區
self.running_mean = torch.zeros(dim)
self.running_var = torch.ones(dim)
def __call__(self, x):
# 計算前向傳播
if self.training:
xmean = x.mean(0, keepdim=True) # 批次平均
xvar = x.var(0, keepdim=True) # 批次方差
else:
xmean = self.running_mean
xvar = self.running_var
xhat = (x - xmean) / torch.sqrt(xvar + self.eps) # 將數據標準化為單位方差
self.out = self.gamma * xhat + self.beta
# 更新緩衝區
if self.training:
with torch.no_grad():
self.running_mean = (1 - self.momentum) * self.running_mean + self.momentum * xmean
self.running_var = (1 - self.momentum) * self.running_var + self.momentum * xvar
return self.out
def parameters(self):
return [self.gamma, self.beta]
Batch-Norm：
- 具有在 back prop 外部訓練的 running mean & variance
self.training = True，這是由於 batch norm 在訓練和評估這兩個階段的行為不同，需要有這樣一個 training flag 來跟蹤 batch norm 在狀態- 批次內處理元素耦合計算，來控制激活的統計特徵，減少內部協變量偏移（Internal Covariate Shift）
class Tanh:
def __call__(self, x):
self.out = torch.tanh(x)
return self.out
def parameters(self):
return []
與以前局部設置一個 g 的torch.Generator相比，後面直接設置全局隨機種子
torch.manual_seed(42);
下面的內容應該很眼熟，包括 embedding table C，和我們的 layer 結構：
n_embd = 10 # 字符嵌入向量的維度
n_hidden = 200 # MLP的隱藏層中神經元的數量
C = torch.randn((vocab_size, n_embd))
layers = [
Linear(n_embd * block_size, n_hidden, bias=False),
BatchNorm1d(n_hidden),
Tanh(),
Linear(n_hidden, vocab_size),
]
# 初始化參數
with torch.no_grad():
layers[-1].weight *= 0.1 # 按比例縮小最後一層(這裡是輸出層)，減少初期模型對預測的自信度
parameters = [C] + [p for layer in layers for p in layer.parameters()]
'''
列表推導式，相當於:
for layer in layers:
for p in layer.parameters():
p...
'''
print(sum(p.nelement() for p in parameters)) # number of parameters in total
for p in parameters:
p.requires_grad = True
優化訓練部分先不做修改，繼續往下看到我們的損失函數曲線波動較大，這是因為 32 的 batch size 太小了，每個批次中你的預測可能非常幸運或不幸（噪聲很大）。
在評估階段，我們要將所有層的 training flag 設置為 False（目前只影響 batch norm 層）：
# 將layer置於評估狀態
for layer in layers:
layer.training = False
我們先解決損失函數圖像的問題：
lossi 是包含所有損失的列表，我們現在要做的基本就是把裡面的值取平均，得到一個更有代表性的值。
複習一下torch.view()的使用：
等同於
view(5, -1)
這可以很方便的將一些列表中的值展開。
torch.tensor(lossi).view(-1, 1000).mean(1)
現在看起來好多了，圖中還能觀察到學習率減少達到了局部最小值。
接下來，我們把下面所示的原先的 Embedding 和 Flattening 操作也變為模塊：
emb = C[Xb]
x = emb.view(emb.shape[0], -1)
class Embedding:
def __init__(self, num_embeddings, embedding_dim):
self.weight = torch.randn((num_embeddings, embedding_dim))
# 現在C成為了embedding的權值
def __call__(self, IX):
self.out = self.weight[IX]
return self.out
def parameters(self):
return [self.weight]
class FlattenConsecutive:
def __call__(self, x):
self.out = x.view(x.shape[0], -1)
return self.out
def parameters(self):
return []
PyTorch 中還有一個容器的概念，基本上是一種將 layer 組織為列表或字典等的方式。其中有一個叫Sequential，基本作用就是把給定的輸入按順序在所有層中傳遞：
class Sequential:
def __init__(self, layers):
self.layers = layers
def __call__(self, x):
for layer in self.layers:
x = layer(x)
self.out = x
return self.out
def parameters(self):
# 獲取所有圖層的參數並將它們拉伸成一個列表。
return [p for layer in self.layers for p in layer.parameters()]
現在我們有了一個 Model 的概念：
model = Sequential([
Embedding(vocab_size, n_embd),
Flatten(),
Linear(n_embd * block_size, n_hidden, bias=False),
BatchNorm1d(n_hidden), Tanh(),
Linear(n_hidden, vocab_size),
])
parameters = model.parameters()
print(sum(p.nelement() for p in parameters)) # 總參數數量
for p in parameters:
p.requires_grad = True
因此得到了更一步的簡化：
# forward pass
logits = model(Xb)
loss = F.cross_entropy(logits, Yb) # loss function
# evaluate the loss
logits = model(x)
loss = F.cross_entropy(logits, y)
# sample from the model
# forward pass the neural net
logits = model(torch.tensor([context]))
probs = F.softmax(logits, dim=1)
實現層狀結構#
我們不希望像現在的模型一樣，在一個步驟中就把信息都壓到一層中了，我們希望像 WaveNet 中預測序列中的下一個字符時，把兩個字符融合成一種雙字符表示，然後再合成四個字符級別的小塊，用這樣的樹狀分層結構慢慢把信息融合到網絡中。
在 WaveNet 的例子中，這張圖是 "Dilated causal convolution layer"（擴張因果卷積層）的可視化，不用管它具體是啥，我們學習它的核心思想 “Progressive fusion（漸進式融合）” 即可。
增加上下文輸入，將這 8 個輸入字符以樹形結構進行處理
# block_size = 3
# train 2.0677597522735596; val 2.1055991649627686
block_size = 8
僅僅將上下文長度擴大就得到了性能提升：
為了弄清我們在做什麼，現在觀察經過各個 layer 過程中 tensor 的形狀：
輸入 4 個隨機數，在模型中的形狀就是 4x8（block_size=8）。
- 經過第一層（embedding），得到了 4x8x10 的輸出，意義就是我們的 embedding table 對於每個字符都有一個要學習的 10 維向量；
- 經過第二層（flatten），就像前面提到的那樣會變成 4x80，這個層的效果是將這 8 個字符的 10 維嵌入拉伸成一長行，就像是連接運算。
- 第三層（linear）就是將這個 80 通過矩陣乘法創建 200 個通道 (channel)
再次總結一下，Embedding 層最終完成的工作。
這個回答中說的非常好:
1. 將稀疏矩陣經過線性變換（查表）變成一個密集矩陣
2. 這個密集矩陣用了 N 個特徵來表示所有的詞。密集矩陣中表象上是一個詞和特徵的關係系數，實際上蘊含了大量的詞與詞之間的內在關係。
3. 它們之間的權重參數，用的是嵌入層學習來的參數進行表徵的編碼。在神經網絡反向傳播優化的過程中，這個參數也會不斷的更新優化。
而線性層在 forward pass 中接受輸入 X 將其與權重相乘，然後可選地添加一個偏差：
def __init__(self, fan_in, fan_out, bias=True):
self.weight = torch.randn((fan_in, fan_out)) / fan_in**0.5 # note: kaiming init
self.bias = torch.zeros(fan_out) if bias else None
這裡的權重是二維的，偏差是一維的
根據輸入輸出的形狀，這個線性層內部的樣子如下：
(torch.randn(4, 80) @ torch.randn(80, 200) + torch.randn(200)).shape
輸出是 4x200，最後加的偏差這裡發生的是廣播語義。
補充一點，PyTorch 中的矩陣乘法運算符十分強大，支持傳入高維 tensor，而矩陣乘法只在最後一個維度上起作用，而其他所有的維度則被視作批處理維度（batch dimensions）
這非常利於我們後面要做的事情：並行的批處理維度。我們不希望一下子輸入 80 個數字，而是在第一層有兩個融合在一起的字符，也就是說只想要 20 個數字輸入，如下所示：
# (1 2) (3 4) (5 6) (7 8)
(torch.randn(4, 4, 20) @ torch.randn(20, 200) + torch.randn(200)).shape
這樣就變成了四組 bigram，bigram 組中的每一個都是 10 維向量
為了實現這樣的結構，Python 中有這樣一個便捷的方法能夠獲取列表中的偶數、奇數部分：
e = torch.randn(4, 8, 10)
torch.cat([e[:, ::2, :], e[:, 1::2, :]], dim=2)
# torch.Size([4, 4, 20])
這樣明確地提取出了偶數、奇數部分，然後將這兩個 4x4x10 的部分連接在一起。
強大的
view()也能完成等效的工作
現在來完善我們的 Flatten 層，創建一個構造函數，並在輸出的最後一個維度中獲取我們想要連接的連續元素的數量，基本上就是將 n 個連續的元素平展並將他們放到最後一個維度中。
class FlattenConsecutive:
def __init__(self, n):
self.n = n
def __call__(self, x):
B, T, C = x.shape
x = x.view(B, T//self.n, C*self.n)
if x.shape[1] == 1:
x = x.squeeze(1)
self.out = x
return self.out
def parameters(self):
return []
- B: Batch size（批大小），代表了批處理中包含的樣本數量。
- T: Time steps（時間步長），表示序列中的元素數量，即序列的長度。
- C: Channels or Features（通道或特徵），代表每個時間步中數據的特徵數量。
-
輸入張量: 輸入
x是一個三維張量，形狀為(B, T, C)。 -
扁平化操作: 通過調用
x.view(B, T//self.n, C*self.n)，這個類將原始數據中連續的時間步合併起來。這裡self.n表示要合併的時間步數。操作的結果是將每n個連續的時間步合併為一個更寬的特徵向量。因此，時間維度T被減少了n倍，而特徵維度C則增加了n倍。新的形狀變為(B, T//n, C*n)，這樣每個新的時間步就包含了原來n個時間步的信息。 -
去除單一時間步維度: 如果合併後的時間步長為 1，即
x.shape[1] == 1，則通過x.squeeze(1)操作去除這一維度，也就是我們之前面對的二維向量情況。
修改後檢查中間各層的形狀：
我們希望 batch norm 中，只維護 68 個通道的均值和方差，而不是 32x4 維的，因此改變現有的 BatchNorm1D 的實現：
class BatchNorm1d:
def __call__(self, x):
# calculate the forward pass
if self.training:
if x.ndim == 2:
dim = 0
elif x.ndim == 3:
dim = (0,1) # torch.mean()可以接受tuple，也就是多個維度的dim
xmean = x.mean(dim, keepdim=True) # batch mean
xvar = x.var(dim, keepdim=True) # batch variance
現在 running_mean.shape 就是 [1, 1, 68] 了
擴大神經網絡#
以及完成了上述改進，我們現在通過增加網絡的大小來進一步提高性能。
n_embd = 24 # 嵌入向量維度
n_hidden = 128 # MLP隱藏層神經元數量
現在的參數量達到了 76579 個，性能也突破了 2.0 的大關：
到目前為止，訓練神經網絡所需的時間增長了很多，儘管性能提升了，但是我們對於學習率等超參數的正確設置都是茫然的，只是盯著訓練的 loss 而不斷 debug 和修改。
卷積#
在本節中，我們實現了 WaveNet 的主要架構，但並沒有實現其中涉及的特定的 forward pass，也就是一個更複雜的線性層：門控線性層 (gated linear layer)，還有殘差連接 (Residual connection) 和跳躍連接 (Skip connection)
這裡簡單了解一下我們實現的樹狀結構與 WaveNet 論文中使用的卷積神經網絡相關的地方。
基本上，我們在這裡使用卷積 (Convolution) 是為了提高效率。卷積允許我們在輸入序列上滑動模型，讓這部分的 for 循環（指卷積核滑動和計算）在 CUDA 內核中完成
我們只是實現了單一的圖中所示的黑色結構並得到一個輸出，但卷積允許你通過這個黑色的結構放到輸入序列上，像線性濾波器一樣同時計算出所有的橙色輸出。
效率提升的原因如下：
- for 循環在 CUDA 核心中完成；
- 重複利用變量，比如第二層的一個白點既是一個第三層白點的左子節點，又是另一個白點的右子節點，這個節點和它的值被使用了兩次。
總結#
本節過後，torch.nn 模塊已經被解鎖了，後面會把模型的實現轉為使用它。
回想一下本節的工作，很多時間都在嘗試讓各個 layer 的形狀正確。因此 Andrej 總是在 Jupyter Notebook 中進行形狀調試，滿意後再複製到 vscode 中。