A simple (and probably quite inefficient) autodiff engine built on NumPy, offering a PyTorch-like API for some basic neural net operations (linear transformation, 2D convolution and pooling). This is only meant to be a proof-of-concept built for fun, inspired by Andrej Karpathy's micrograd.
npgrad
is based on an Array
class wrapping a NumPy ndarray
(and implementing the __array__
method), with some extra bits to handle the gradient computation. Some of the main NumPy functions (add
, multiply
, log
, reshape
, ...) have an implementation which works with Array
and can be called from either npgrad
or numpy
itself. Array
also supports some of the same attributes and methods provided by PyTorch's Tensor
(requires_grad
, grad
, backward()
) and can be used with the few functions from npgrad.nn
.
import numpy as np
import npgrad as npg
a = npg.array([1, 2, 3])
# both b and c are Array instances
b = npg.log(a)
c = np.log(a)
print(b.__class__, c.__class__) # prints <class 'npgrad.array.Array'> for both
# start tracking operations on a for grad computation
a.requires_grad = True
(a ** 2).sum().backward() # populate a.grad
# nn example
n, c = 4, 3
data = np.random.rand(n, 3, 16, 16)
lbls = np.random.randint(0, c, n)
w1 = npg.asarray(np.random.rand(2, 3, 4, 4)).requires_grad_()
x = npg.nn.functional.conv2d(data, w1, stride=2)
x = npg.nn.functional.max_pool2d(x, 4, padding=1)
x = x.reshape((n, -1)) # -> (n, 8)
w2 = npg.asarray(np.random.rand(8, c)).requires_grad_()
x = x @ w2 # -> (n, c)
loss = npg.nn.functional.cross_entropy(x, lbls).mean()
loss.backward() # populate w1.grad and w2.grad
# support modules too
linear = npg.nn.Linear(16, 5)
conv = npg.nn.Conv2d(4, 3, kernel_size=5, stride=3, padding=2)