This is a fully NumPy-based CNN library with PyTorch-style APIs. I tried to replicate some functions in the torch library like nn.Sequential, nn.Conv2d, a naive version of Autograd engine, etc. You might find some of the expressions in this programme quite familiar. Btw my favourite one is loss.backward() :)
I used this library to build a LeNet, and trained with MNIST datasets. (Note that I replaced the sigmoid in the original LeNet with ReLU, and used MaxPooling rather than AveragePooling. Other model specs are completely the same.)
model = Sequential(
Conv2d(in_channels=1, out_channels=6, kernel_size=5, stride=1, padding=2),
ReLU(),
MaxPool2d(kernel_size=2, stride=2),
Conv2d(in_channels=6, out_channels=16, kernel_size=5, stride=1, padding=0),
ReLU(),
MaxPool2d(kernel_size=2, stride=2),
Flatten(),
Linear(400, 120),
ReLU(),
Linear(120, 84),
ReLU(),
Linear(84, 10),
)
Training this model for 1 epoch takes roughly 5 minutes on my Apple M1 chip, while training an equivalent model using PyTorch takes 20 seconds.
# forward
loss = Cross_entropy(model, images, labels)
# backward
loss.backward()
model.save_model('pretrained_model')
model.load_model('pretrained_model.npy')
On average, a full forward-backward-update process takes 0.05 second (LeNet, batch_size=10, on Apple M1 chip).
forward backward
Conv2d | 19.606s 157.369s
ReLU | 0.633s 1.039s
MaxPool2d | 11.346s 49.476s
Conv2d | 10.881s 44.958s
ReLU | 0.281s 0.852s
MaxPool2d | 3.415s 9.228s
Flatten | 0.186s 0.026s
Linear | 0.904s 1.489s
ReLU | 0.072s 0.142s
Linear | 0.195s 0.338s
ReLU | 0.044s 0.131s
Linear | 0.059s 0.189s
----------------------------------------
Total forward + backward time : 312.8575s
Iterations : 6000
Like PyTorch, we can adjust our CNN models to arbitrary layer arrangements, channel sizes, conv kernel and maxpooling kernel sizes, strides, zeropaddings, etc. The autograd engine (naive version, of course) will do the backprop automatically.