File size: 2,863 Bytes
5c6c847 | 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 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 | """优化器。"""
import numpy as np
class Optimizer:
def __init__(self, params, lr):
self.params = list(params)
self.lr = lr
def zero_grad(self):
for p in self.params:
p.grad = None
def step(self):
raise NotImplementedError
class SGD(Optimizer):
def __init__(self, params, lr=0.01, momentum=0.0):
super().__init__(params, lr)
self.momentum = momentum
self.v = [np.zeros_like(p.data) for p in self.params]
def step(self):
for p, v in zip(self.params, self.v):
if p.grad is None:
continue
v *= self.momentum
v += p.grad
p.data -= self.lr * v
class AdamW(Optimizer):
"""Adam + 解耦权重衰减。weight_decay=0 时就是普通 Adam。"""
def __init__(self, params, lr=1e-3, betas=(0.9, 0.999), eps=1e-8, weight_decay=0.0):
super().__init__(params, lr)
self.b1, self.b2 = betas
self.eps = eps
self.weight_decay = weight_decay
self.t = 0
self.m = [np.zeros_like(p.data) for p in self.params]
self.v = [np.zeros_like(p.data) for p in self.params]
def step(self):
self.t += 1
c1, c2 = 1 - self.b1 ** self.t, 1 - self.b2 ** self.t
for p, m, v in zip(self.params, self.m, self.v):
if p.grad is None:
continue
g = p.grad
m *= self.b1
m += (1 - self.b1) * g
v *= self.b2
v += (1 - self.b2) * g * g
if self.weight_decay:
p.data -= self.lr * self.weight_decay * p.data
p.data -= self.lr * (m / c1) / (np.sqrt(v / c2) + self.eps)
def state_dict(self):
"""优化器状态（步数 + 一阶/二阶动量）。和 model.state_dict() 一起存，才能精确续训。"""
state = {"t": np.array(self.t), "lr": np.array(self.lr)}
for i, (m, v) in enumerate(zip(self.m, self.v)):
state[f"m.{i}"], state[f"v.{i}"] = m.copy(), v.copy()
return state
def load_state_dict(self, state):
assert sum(k.startswith("m.") for k in state) == len(self.params), "参数个数和存档对不上"
self.t, self.lr = int(state["t"]), float(state["lr"])
for i in range(len(self.params)):
self.m[i][...] = state[f"m.{i}"]
self.v[i][...] = state[f"v.{i}"]
Adam = AdamW
def clip_grad_norm(params, max_norm):
"""按全局 L2 范数裁剪梯度，返回裁剪前的范数。"""
grads = [p.grad for p in params if p.grad is not None]
total = float(np.sqrt(sum((g.astype(np.float64) ** 2).sum() for g in grads)))
if total > max_norm:
scale = max_norm / (total + 1e-6)
for p in params:
if p.grad is not None:
p.grad = p.grad * scale
return total
|