GPT-2复现学习

class Neuron(Module):
		# nin 为输入个数
    def __init__(self, nin, nonlin=True):
		    # 权重和偏置值随机初始化
        self.w = [Value(random.uniform(-1,1)) for _ in range(nin)]
        self.b = Value(0)
        self.nonlin = nonlin
		
		# __call__ 使得 x = [2.0, 3.0]; n = Neuron(2); n(x) 调用 __call__
    def __call__(self, x):
        act = sum((wi*xi for wi,xi in zip(self.w, x)), self.b)
        return act.relu() if self.nonlin else act

    def parameters(self): # 用于返回参数列表
        return self.w + [self.b]

class Layer(Module):
		# nout 为层中神经元个数
    def __init__(self, nin, nout, **kwargs):
        self.neurons = [Neuron(nin, **kwargs) for _ in range(nout)]

    def __call__(self, x):
        out = [n(x) for n in self.neurons] # 对每个神经元进行计算
        return out[0] if len(out) == 1 else out

    def parameters(self):
        return [p for n in self.neurons for p in n.parameters()]

class MLP(Module):
		# nouts 为 list(nout) 定义所有层尺寸
    def __init__(self, nin, nouts):
        sz = [nin] + nouts
        self.layers = [Layer(sz[i], sz[i+1], nonlin=i!=len(nouts)-1) for i in range(len(nouts))]

    def __call__(self, x):
        for layer in self.layers:
            x = layer(x)
        return x

    def parameters(self): # 获取所有参数
        return [p for layer in self.layers for p in layer.parameters()]

损失函数 L
  ↑
  |      情况1：梯度为正
  |     ↗   我们要往左走（减小w）
  |   ↗
  | ↗
  +----------→ 参数 w
  
损失函数 L  
  ↑
  |      情况2：梯度为负
  |   ↘     我们要往右走（增加w）
  |     ↘
  |       ↘
  +----------→ 参数 w

import torch
# 初始为 0 的 27*27 矩阵
N = torch.zeros((27, 27), dtype=torch.int32) 
# N[0, 0] = 1 赋值
# N[0].float() 取出第一行 求和为.sum()

for w in words:
	chs = ['.'] + list(w) + ['.'] # 收缩 .. 到 N[0,0]
	for ch1, ch2 in zip(chs, chs[1:]):
		ix1 = stoi[ch1]
		ix2 = stoi[ch2]
		N[ix1, ix2] += 1 # 第一列为以字母结尾, 第一行为以字母开头

# 绘图方法
import matplotlib.pyplot as plt
%matplotlib inline

plt.figure(figsize=(16, 16))
plt.imshow(N, cmap='Blues')
for i in range(27):
    for j in range(27):
        chstr = itos[i] + itos[j]
        plt.text(j, i, chstr, ha="center", va="bottom", color='gray')
        plt.text(j, i, N[i, j].item(), ha="center", va="top", color='gray')
plt.axis('off');

g = torch.Generator().manual_seed(xxx)

out = []
ix = 0
while True:
    
    # p = N[ix].float() # 遍历矩阵每一行 训练二元模型
    # p = p / p.sum() # 归一化
    # 优化上面两行
    p = P[ix]
    # p = torch.ones(27) / 27.0 均匀分布, 即完全未经训练的模型, 获取的名字效果很差
    ix = torch.multinomial(p, num_samples=1, replacement=True, generator=g).item() # 选取下一个索引行
    out.append(itos[ix])
    if ix == 0:
        break

print(''.join(out)) # 获取生成结果

log_likelihood = 0.0
n = 0 # 均值
for w in words[:3]:
    chs = ['.'] + list(w) + ['.']
    for ch1, ch2 in zip(chs, chs[1:]):
        ix1 = stoi[ch1]
        ix2 = stoi[ch2]
        prob = P[ix1, ix2]
        logprob = torch.log(prob) # log(a*b*c) = log(a) + log(b) + log(c)
        log_likelihood += logprob
        n += 1
        print(f'{ch1}{ch2}: {prob:.4f} {logprob:.4f}')
print(f'{-log_likelihood /n}') # 平均负对数似然

xs, ys = [], []
for w in words:
    chs = ['.'] + list(w) + ['.']
    for ch1, ch2 in zip(chs, chs[1:]):
        ix1 = stoi[ch1]
        ix2 = stoi[ch2]
        xs.append(ix1)
        ys.append(ix2)
xs = torch.tensor(xs)
ys = torch.tensor(ys)
num = xs.nelement()
print("number of examples: ", num)

g = torch.Generator().manual_seed(xx)
W = torch.randn((27, 27), generator=g, requires_grad=true)

# 梯度下降
for k in range(10):
	# 前向传播
	xenc = F.one_hot(xs, num_classes=27).float()
	logits = xenc @ W
	counts = (xenc @ W).exp()
	probs = counts / counts.sum(1, keepdims=True)
	# 加入L2正则化 0.01为正则化强度
	loss = -probs[torch.arange(5), ys].log().mean() + 0.01*(W**2).mean()
	
	# 反向传播
	W.grad = None
	loss.backward() # W.grad 查看情况, (27, 27)
	print(loss.item()) # 期望达到对原本名字集 words 计算得到的负平均对数似然损失相近的值
	
	# 更新
	W.data += -0.1 * W.grad # 0.1 学习率

g = torch.Generator().manual_seed(xx)

for i in range(5):
	out = []
	ix = 0
	while True:
		# BEFORE:
		# p = P[ix]
		
		# NOW: 取 ix 编码为 one-hot 编码
		xenc = F.one_hot(torch.tensor([ix]), num_classes=27).float() # (1, 27)
		logits = xenc @ W # (1, 27) * (27, 27) -> (1, 27)
		counts = logits.exp()
		p = counts / counts.sum(1, keepdims=True) # p: (1, 27)
		# 每个值表示当前值为 x , 下一个值为 27 个字符中每个的概率
		
		# 取出某个字符的索引值, 0 ~ 26, 输出形状 (1, 1)
		# 使用随机采样可以增加多样性, 如: q 可以结合 a,u, 但贪心则使得 q 可能一直和一个字母结合
		ix = torch.multinomial(p, num_samples=1, replacement=True, generator=g).item()
		out.append(itos[ix])
		if ix == 0:
			break
print(''.join(out))

block_size = 3 # 上下文长度, 使用多少来预测下一个
X, Y = [], []
for w in words[:5]: # 完整数据集: 228000 
	context = [0] * block_size
	for ch in w + '.':
		ix = stoi[ch]
		X.append(context)
		Y.append(ix)
		print(''.join(itos[i] for i in context), '--->', itos[ix])
		context = context[1:] + [ix]

X = torch.tensor(X) # 示例 shape (32, 3), 32 表示目前的数据集大小
Y = torch.tensor(Y) # 标签 shape (32)
# ... ---> e
# ..e ---> m
# .em ---> m
# emm ---> a
# mma ---> .

W2 = torch.randn((100, 27)) # 输出 27 个概率
b2 = torch.randn(27)
logits = h @ W2 + b2 # (32, 27)
counts = logits.exp()
prob = counts / counts.sum(1, keepdims=True) # shape: (32, 27)
prob[torch.arrange(32), Y)
# 获取给出权重设置下，由神经网络分配给序列中正确字符的当前概率, 每行 Y 标签对应的概率, 理想情况概率均为 1
# 损失获取
loss = -prob[torch.arrange(32), Y).log().mean()
# 优化损失 loss, 交叉熵得到结果一致, 优于上一行损失获取, 底层更加稳定
loss = F.cross_entropy(logits, Y)

# 数据集
X.shape, Y.shape

g = torch.Generator().manual_seed(xxx)
C = torch.randn((27, 2), generator=g) # 嵌入向量为 2 维
W1 = torch.randn((6, 100), generator=g)
b1 = torch.randn(100, generator=g)
W2 = torch.randn((100, 27), generator=g)
b2 = torch.randn(27, generator=g)

parameters = [C, W1, b1, W2, b2]
sum(p.nelement() for p in parameters) # 参数总数量

for p in parameters:
	p.requires_grad = True

for _ in range(10):
	# 前向传播
	emb = C[X]
	h = torch.tanh(emb.view(-1, 6) @ W1 + b1)
	logits = h @ W2 + b2
	loss = F.cross_entropy(logits, Y)
	print(loss.item()) # 最终稳定在 0.25 左右 实际可能导致过拟合
	
	# 反向传播
	for p in parameters:
		p.grad = None
	loss.backward()
	
	# 更新
	for p in parameters:
		p.data += -0.1 * p.grad

for _ in range(100):
	# 小批量
	ix = torch.randint(0, X.shape[0], (32, ))
	
	# 前向传播
	emb = C[X[ix]] # 只获取 32 行 (32, 3, 2)
	h = torch.tanh(emb.view(-1, 6) @ W1 + b1)
	logits = h @ W2 + b2
	loss = F.cross_entropy(logits, Y[ix])
	
	# 反向传播
	for p in parameters:
		p.grad = None
	loss.backward()
	
	# 更新
	for p in parameters:
		p.data += -0.1 * p.grad
print(loss.item()) # 最终损失值在 2.5 左右徘徊, 该损失值只是小批次的损失值

# 所有数据集损失
emb = C[X]
h = torch.tanh(emb.view(-1, 6) @ W1 + b1)
logits = h @ W2 + b2
loss = F.cross_entropy(logits, Y)

lri = []
lossi = []

for i in range(1000):
	# 小批量
	ix = torch.randint(0, X.shape[0], (32, ))
	
	# 前向传播
	emb = C[X[ix]] # 只获取 32 行 (32, 3, 2)
	h = torch.tanh(emb.view(-1, 6) @ W1 + b1)
	logits = h @ W2 + b2
	loss = F.cross_entropy(logits, Y[ix])
	
	# 反向传播
	for p in parameters:
		p.grad = None
	loss.backward()
	
	# 更新
	lr = lrs[i] # 学习率
	for p in parameters:
		p.data += -lr * p.grad
	
	# 记录
	# lri.append(lr) 第一种
	lri.append(lre[i]) # 第二种, 指数队列
	lossi.append(loss)

# 构建数据集
def build_dataset(words):
	block_size = 3
	X, Y = [], []
	for w in words:
		context = [0] * block_size
		for ch in w + '.':
			ix = stoi[ch]
			X.append(context)
			Y.append(ix)
			context = context[1:] + [ix]
	X = torch.tensor(X)
	Y = torch.tensor(Y)
	return X, Y
	
# 拆分数据集
import random
random.seed(xx)
random.shuffle(words)
n1 = int(0.8*len(words))
n2 = int(0.9*len(words))

Xtr, Ytr = build_dataset(words[:n1])
Xdev, Ydev = build_dataset(words[n1:n2])
Xte, Yte = build_dataset(words[n2:])

lri = []
lossi = []
stepi = []

for i in range(30000):
	# 小批量
	ix = torch.randint(0, Xtr.shape[0], (32, ))
	
	# 前向传播
	emb = C[Xtr[ix]] # 只获取 32 行 (32, 3, 2)
	h = torch.tanh(emb.view(-1, 6) @ W1 + b1)
	logits = h @ W2 + b2
	loss = F.cross_entropy(logits, Ytr[ix])
	
	# 反向传播
	for p in parameters:
		p.grad = None
	loss.backward()
	
	# 更新
	lr = 0.1 # 学习率
	for p in parameters:
		p.data += -lr * p.grad
	
	# 记录
	stepi.append(i)
	lossi.append(loss.item())
plt.plot(stepi, lossi)

# 通过模型采样
g = torch.Generator().manual_seed(2147483647 + 10)
for _ in range(20):
    out = []
    context = [0] * block_size
    while True:
		    # 上下文字符索引转换为嵌入向量, 只处理一个样本
		    # 批量大小 1，上下文长度 block_size，嵌入维度 d
        emb = C[torch.tensor([context])] # (1, block_size, d)
        h = torch.tanh(emb.view(1, -1) @ W1 + b1) # 将嵌入展平为 (1, block_size*d)
        logits = h @ W2 + b2
        probs = F.softmax(logits, dim=1)
        ix = torch.multinomial(probs, num_samples=1, generator=g).item()
        context = context[1:] + [ix]
        out.append(ix)
        if ix == 0:
            break
    print(''.join(itos[i] for i in out))

X = torch.randn(1000, 10)
W = torch.randn(10, 200) # 需要在此基础上乘上一个值, 可以改变标准差缩小增大, / 10**0.5
	# 基本上乘以多少后标准差为多少倍, 输入归一化
Y = X @ W
print(X.mean(), X.std())
print(Y.mean(), Y.std())
plt.figure(figsize=(20, 5))
plt.subplot(121)
plt.hist(X.view(-1).tolist(), 50, density=True);
plt.subplot(122)
plt.hist(Y.view(-1).tolist(), 50, density=True);

hpreact = embcat @ W1 + b1 # (32, 200)
bnmeani = hpreact.mean(0, keepdim=True)
bnstdi = hpreact.std(0, keepdim=True)
hpreact = bngain * (hpreact - bnmeani) / bnstdi + bnbias
# 上述代码中 bnstdi 也可加入一个极小值 epsilon 防止 bnstdi 为 0, 默认 10 的 -5 次方

... 
# 不使用梯度下降更新, 平滑均值方式更新
with torch.no_grad():
	bnmean_running = 0.999 * bnmean_running + 0.001 * bnmeani
	bnstd_running = 0.999 * bnstd_running + 0.001 * bnstdi

# 参数初始化阶段 仍使用 b1
# bngain = torch.ones((1, n_hidden))
# bnbias = torch.zeros((1, n_hidden))
bngain = torch.randn((1, n_hidden))*0.1 + 1.0
bnbias = torch.randn((1, n_hidden))*0.1

# bnmean_running = torch.zeros((1, n_hidden))
# bnstd_running = torch.ones((1, n_hidden))

# 小批量提出作为单独单元格
batch_size = 32
n = batch_size
ix = torch.randint(0, Xtr.shape[0], (batch_size, ), generator=g)
Xb, Yb = Xtr[ix], Ytr[ix]

# 前向传播
emb = C[Xb] # 小批量
embcat = emb.view(emb.shape[0], -1) # 嵌入向量展平, 准备输入到全连接层
# 线性层 1
hprebn = embcat @ W1 + b1
# 归一化层
bnmeani = 1/n*hprebn.sum(0, keepdim=True) # 计算当前批次均值
bndiff = hprebn - bnmeani # 每个激活值与均值的差值
bndiff2 = bndiff**2 # 差值平方
bnvar = 1/(n-1)*(bndiff2).sum(0, keepdim=True) # 当前批次方差(无偏估计, 除以 n-1)
# 原论文奇怪的点: 训练时使用偏差估计, 测试时使用无偏估计
bnvar_inv = (bnvar + 1e-5)**-0.5 # 标准差的倒数(小常量防止除零)
bnraw = bndiff * bnvar_inv # 激活值归一化处理, 减去均值除以标准差
hpreact = bngain * bnraw + bnbias # 应用可学习的缩放和偏移参数
# 非线性层
h = torch.tanh(hpreact)
# 线性层 2
logits = h @ W2 + b2
# 损失计算
logit_maxes = logits.max(1, keepdim=True).values # 每个样本最大 logit 值
norm_logits = logits - logit_maxes # 所有 logits 中减去最大值, 防止指数运算溢出
counts = norm_logits.exp() # 计算指数得到未归一化的概率
counts_sum = counts.sum(1, keepdim=True) # 计算每个样本概率总和, 归一化
counts_sum_inv = counts_sum**-1 # 计算总和的倒数, 计算 softmax 概率
probs = counts * counts_sum_inv # 每个 count 除以总和, 计算 softmax 概率
logprobs = probs.log() # 计算概率对数 
loss = -logprobs[range(n), Yb].mean()
# 选择了每个样本真实正确类别对应的对数概率，计算交叉熵损失

# pytorch 反向传播
for p in parameters:
    p.grad = None
for t in [logprobs, probs, counts, counts_sum, counts_sum_inv,
          norm_logits, logit_maxes, logits, h, hpreact, bnraw, 
          bnvar_inv, bnvar, bndiff2, bndiff, hprebn, bnmeani,
          embcat, emb]:
    t.retain_grad()
loss.backward() # 执行反向传播，计算所有参数的梯度
loss

probs = counts * counts_sum_inv # counts(32, 27), counts_sum_inv(32, 1)
# 简化
c = a * b
a[3x3] * b[3,1] ->
a11*b1 a12*b1 a13*b1
a21*b2 a22*b2 a23*b2
a31*b3 a32*b3 a33*b3
c[3x3] # 乘法求导 dc/db = a

# 通过 sum 来处理重复的 b
dcounts_sum_inv = (counts * dprobs).sum(1, keepdim=True) # (32, 1)
cmp('counts_sum_inv', dcounts_sum_inv, counts_sum_inv)
# counts_sum_inv  | exact: True  | approximate: True  | maxdiff: 0.0

# 对于 counts_sum = counts.sum(1, keepdim=True) 沿着行的总和
# counts_sum (32, 1)   counts (32, 27)

# a11 a12 a13 ----> b1 (= a11 + a12 + a13)   
# a21 a22 a23 ----> b2 (= a21 + a22 + a23)
# a31 a32 a33 ----> b3 (= a31 + a32 + a33)

# 导数均匀分配 db1/da11 = 1, db1/da12 = 1, db1/da13 = 1, 而b1 对第二三行所有元素导数为 0
dcounts += torch.ones_like(counts) * dcounts_sum # dcounts 由不同计算影响相加获取
cmp('counts', dcounts, counts)
# counts          | exact: True  | approximate: True  | maxdiff: 0.0

# logits = h @ W2 + b2
# dlogits (32, 27)  h (32, 64)  W2 (64, 27)  b2 (27)
dh = dlogits @ W2.T
dW2 = h.T @ dlogits
db2 = dlogits.sum(0)
cmp('h', dh, h)
cmp('W2', dW2, W2)
cmp('b2', db2, b2)
# h               | exact: True  | approximate: True  | maxdiff: 0.0
# W2              | exact: True  | approximate: True  | maxdiff: 0.0
# b2              | exact: True  | approximate: True  | maxdiff: 0.0

d = a @ b + c

d11 d12   a11 a12   b11 b12   c1 c2
d21 d22 = a21 a22 * b21 b22 + c1 c2

# 即
d11 = a11b11 + a12b21 + c1
d12 = a11b12 + a12b22 + c2
d21 = a21b11 + a22b21 + c1
d22 = a21b12 + a22b22 + c2

dL/da11 = dL/dd11 * b11 + dL/dd12 * b12
dL/da12 = dL/dd11 * b21 + dL/dd12 * b22
dL/da21 = dL/dd21 * b11 + dL/dd22 * b12
dL/da22 = dL/dd21 * b21 + dL/dd22 * b22

dL/da11 dL/da12   dL/dd11 dL/dd12   b11 b21
dL/da21 dL/da22 = dL/dd21 dL/dd22 * b12 b22 = dL/dd @ b^T

# 最终
dL/da = dL/dd @ b^T
dL/dc = dL/dd * sum(0)
dL/db = a^T @ dL/dd

# hpreact (32, 64)  bngain (1, 64)  bnraw (32, 64)  bnbias (1, 64)
# hpreact = bngain * bnraw + bnbias
dbngain = (bnraw * dhpreact).sum(0, keepdim=True) # (32, 64) -> (1, 64)
dbnraw = bngain * dhpreact
dbnbias = dhpreact.sum(0, keepdim=True) # (32, 64) -> (1, 64)
cmp('bngain', dbngain, bngain)
cmp('bnraw', dbnraw, bnraw)
cmp('bnbias', dbnbias, bnbias)
# bngain          | exact: True  | approximate: True  | maxdiff: 0.0
# bnraw           | exact: True  | approximate: True  | maxdiff: 0.0
# bnbias          | exact: True  | approximate: True  | maxdiff: 0.0