NumPy实现类PyTorch的动态计算图和神经网络框架(MLP, CNN, RNN, Transformer)
MIT License
前作:PyNet: Use NumPy to build neuron network。在那里我们基于求导规则实现了全连接网络。在这里,我们向当今的深度学习框架看齐,实现属于自己的DL框架。
PyDyNet已被多个技术公众号和社区分享:居然用Numpy实现了一个深度学习框架.
train_loader
函数;@no_grad()
和with no_grad()
,详见autograd.py;cu*.py
;2024.6.29: ver 0.1 增加了LayerNorm和Embedding; 修正了cuda训练的问题; 加入多个可复现测试: (MLP, LeNet, BN & Dropout, RNN, Transformer).
PyDyNet也是纯NumPy(0.0.7版本后加入CuPy,其用法和NumPy一致)实现的神经网络,语法受PyTorch的启发,大致结构如下:
graph BT
N(numpy/cupy.ndarray) ----> ds(Dataset) ----> Data(DataLoader)--> Mission
N --> A(Tensor) --Eager execution--> B(Basic operators: add, exp, etc)
B -.Autograd-.-> A
B --> CO(Complex operators:softmax,etc)
--> f(Function:linear, conv2d, etc)
--> M(Basic Module:Linear,Conv2d,etc)
--> CM(Advanced Module:CNN,RNN,Transformer,...)
--> Mission(PyDyNet)
N --> GD(Optimizer:SGD, Adam, etc) ----> LS(lr_scheduler:StepLR, etc)--> Mission
虚线表示用户可以通过no_grad
来关闭自动微分功能. 我们实现了:
将NumPy数组包装成具有梯度等信息的张量(Tensor):
from pydynet import Tensor
x = Tensor(1., requires_grad=True)
print(x.data) # 1.
print(x.ndim, x.shape, x.is_leaf) # 0, (), True
将NumPy数组的计算(包括数学运算、切片、形状变换等)抽象成基础算子(Basic operators),并对部分运算加以重载:
import pydynet as pdn
x = pdn.Tensor([1, 2, 3])
y = pdn.exp(x) + x
z = pdn.sum(x)
print(z.data) # 36.192...
手动编写基础算子的梯度,实现和PyTorch相同的动态图自动微分机制(Autograd),从而实现反向传播
import pydynet as pdn
from pydynet import Tensor
x = Tensor([1., 2., 3.], requires_grad=True)
y = pdn.log(x) + x
z = pdn.sum(y)
z.backward()
print(x.grad) # [2., 1.5, 1.33333333]
基于基础算子实现更高级的算子(Complex operators),它们不再需要手动编写导数:
import pydynet as pdn
def simple_sigmoid(x: pdn.Tensor):
return 1 / (1 + pdn.exp(-x))
实现了Mudule,包括激活函数,损失函数等,从而我们可以像下面这样定义神经网络,损失函数项:
import pydynet.nn as nn
import pydynet.nn.functional as F
n_input = 64
n_hidden = 128
n_output = 10
class Net(nn.Module):
def __init__(self) -> None:
super().__init__()
self.fc1 = nn.Linear(n_input, n_hidden)
self.fc2 = nn.Linear(n_hidden, n_output)
def forward(self, x):
x = self.fc1(x)
x = F.sigmoid(x)
return self.fc2(x)
net = Net()
loss = nn.CrossEntropyLoss()
l = loss(net(X), y)
l.backward()
实现了多种优化器和学习率衰减策略,从而实现神经网络的训练;其中优化器和PyTorch一样支持权值衰减,即正则化:
from pydynet.optim import Adam, StepLR
...
net = Net()
optimizer = Adam(net.parameters(), lr=0.01)
lr_scheduler = StepLR(optimizer, step_size=10)
for epoch in range(EPOCHES):
for data in data_loader:
train(...)
optimizer.step()
lr_scheduler.step()
实现了Dataset和DataLoader对数据集进行加载与划分:
from pydynet.data import Dataset, DataLoader
class TrainSet(Dataset):
def __init__(self, X, y) -> None:
self.data = X
self.target = y
def __getitem__(self, index):
return self.data[index], self.target[index]
def __len__(self):
return len(self.data)
data_loader = DataLoader(TrainSet(X, y), batch_size, shuffle)
Dropout机制,Batch Normalization机制,以及将网络划分成训练阶段和评估阶段;
基于im2col高效实现Conv1d, Conv2d, max_pool1d和max_pool2d,从而实现CNN;
支持多层的多层双向RNN,LSTM和GRU;
多种初始化方式,包括Kaiming和Xavier;
基于cupy实现了显卡计算和训练:
from pydynet import Tensor
x = Tensor([1., 2., 3.], device='cuda')
y = Tensor([1., 2., 3.], device='cuda')
z = (x * y).sum()
w = Tensor([1., 2., 3.]) # CPU上的Tensor
x * w # 报错
pip install pydynet
或本地安装
git clone https://github.com/Kaslanarian/PyDyNet
cd PyDyNet
python setup.py install
tests中是一些例子。运行python tests/XXX.py
即可:
autodiff1d.py利用自动微分,对一个一维凸函数进行梯度下降:
以及一个多元凸函数的例子: autodiff2d.py
mlp_cnn.py使用全连接网络(三层+残差)和LeNet对MNIST进行分类. 训练准确率和测试准确率:
mlp_dropout_bn.py使用三种网络对fetch_olivetti_faces
人脸(64×64)数据集进行分类并进行性能对比:
学习效果对比:
rnn_sin.py中是一个用RNN从$x=\sin(z)$学习$y=\cos(2z)$例子. 最后的训练结果:
transformer.py中是一个用Transformer训练文本分类模型的例子. 训练结果:
数据集 (CoLA) 链接: https://nyu-mll.github.io/CoLA/cola_public_1.1.zip
在训练batch size为128, 测试batch size为512情况下,模型在CPU和GPU上的训练速度比较:
Net | Dataset | CPU time (s) per Epoch | GPU time (s) per Epoch |
---|---|---|---|
ResidualMLP | MNIST (80000×574) | 20.256±0.138 | 2.903±.018 |
LeNet | MNIST (80000×574) | 239.664±2.108 | 10.148±0.026 |
1-layer Transformer | CoLA (8551×45×64) | 17.503±0.251 | 1.125±0.002 |
设备: Nvidia GeForce RTX 3090.