A lightweight automatic differentiation library with some simple machine learning concepts build on top of it.
We provide an API very similar to pytorch. See the following simple example on how to use the autograd system (from examples/gradient_descent.py
):
import lightgrad as light
# create tensors
a = light.uniform(-1, 1, shape=(10, 10))
b = light.uniform(-1, 1, shape=(10, 10))
c = light.uniform(-1, 1, shape=(10, 10))
# objective to minimize
f = lambda: (a.tanh() + b.sigmoid()) @ (c.relu() - a.sigmoid())
# optimization loop
for epoch in range(100):
# execute and compute gradients
# the argument allow_fill allows for backpropagation
# starting from non-item tensors
y = f()
y.backward(allow_fill=True)
# we need to disable gradients for the
# paremeter update
with light.no_grad():
# now we can easily access the gradients
# and update the parameters
a -= 0.1 * a.grad
b -= 0.1 * b.grad
c -= 0.1 * c.grad
# reset gradients for next iteration
a.zero_grad()
b.zero_grad()
c.zero_grad()
Alternatively to manually updating the parameters, we provides an optimizer module called optim
with some standard optimizers already implemented (SGD
, Adam
, etc.). Using this we can reduce the above example to the following:
...
# create optimizer
optim = light.optim.SGD([a, b, c], lr=0.1)
# optimization loop
for epoch in range(100):
# use the optimizer for gradient resetting
# and parameter updates
optim.zero_grad()
f().backward(allow_fill=True)
optim.step()
Modules yield an easy way of managing parameters and combining them with their functionality. Additionally they provide further features for easy parameter manipulation (loading, mapping, etc.). We can simplify our example using the Module
class as follows:
import lightgrad as light
# define a module class holding all the neccessary parameters
class SimpleModule(light.nn.Module):
""" Simple Module holding a few parameters """
def __init__(self):
nn.Module.__init__(self)
self.a = light.uniform(-1, 1, shape=(10, 10))
self.b = light.uniform(-1, 1, shape=(10, 10))
self.c = light.uniform(-1, 1, shape=(10, 10))
def forward(self):
return (self.a.tanh() + self.b.sigmoid()) @ (self.c.relu() - self.a.sigmoid())
# create module and optimizer
f = SimpleModule()
optim = light.optim.SGD(f.parameters(), lr=0.1)
...
Combining the autograd system with both Optimizers and Modules results in a simple interface for defining, training and evaluating neural networks. The lightgrad.nn
package already provides some common neural network components (Linear
, Conv2d
, etc.). A simple example of a convolutional neural network is shown below (from examples/mnist.py
).
import lightgrad.nn as nn
class CNN(nn.Module):
""" Simple convolutional neural network """
def __init__(self):
nn.Module.__init__(self)
# create convolutional layers
self.c1 = nn.Conv2d(1, 8, kernelsize=3)
self.c2 = nn.Conv2d(8, 16, kernelsize=3)
# create classification layer
self.l1 = nn.Linear(5 * 5 * 16, 10)
def forward(self, x):
# apply convolutions
y = self.c1(x).max_pool().relu()
y = self.c2(y).max_pool().relu()
# classify
y = self.l1(y.reshape(-1, 5 * 5 * 16))
return y
Lightgrad currently supports the following backends:
- CPU (Default)
- OpenCL
Pushing a tensor to another accelerator can be done using one function call.
# create a cpu tensor and push it to opencl
light.uniform(-1, 1, shape=(10, 10)).opencl()
# push each parameter of the module to opencl
module.map_parameters(lambda p: p.opencl())
Create a new Tensor class inheriting from autograd.tensor.AbstractTensor
. The AbstractTensor
takes care of all gradient managment.
import numpy as np
from lightgrad.autograd.tensor import AbstractTensor
class NewTensor(AbstractTensor):
""" Tensor class for new accelerator """
def __init__(self, *args, dtype:type =np.float32, requires_grad:bool =True):
# create raw tensor data and initialize abstract tensor
AbstractTensor.__init__(self, data=data, requires_grad=requires_grad)
@property
def dtype(self):
raise NotImplementedError()
@property
def shape(self) -> tuple:
raise NotImplementedError()
@staticmethod
def empty(shape, *args, **kwargs) -> "NewTensor":
raise NotImplementedError()
@staticmethod
def zeros(shape, *args, **kwargs) -> "NewTensor":
raise NotImplementedError()
@staticmethod
def ones(shape, *args, **kwargs) -> "NewTensor":
raise NotImplementedError()
@staticmethod
def uniform(low, high, shape, *args, **kwargs) -> "NewTensor":
raise NotImplementedError()
@staticmethod
def from_numpy(a:np.ndarray, requires_grad:bool =True) -> "NewTensor":
raise NotImplementedError()
def numpy(self) -> np.ndarray:
raise NotImplementedError()
def copy(self, requires_grad:bool =True) -> "NewTensor":
raise NotImplementedError()
Further functionality is provided by registering subclasses of lightgrad.autograd.func.Function
to the new tensor class.
from lightgrad.autograd.func import Function
@NewTensor.register_op(name="exp")
class exp(Function):
def forward(ctx, t):
# compute exp on raw tensor data
raw_x_data = t.data
raw_y_data = raw_x_data.exp()
# save result for backpropagation
ctx.save_for_backward(raw_y_data)
# return tensor instance holding raw data
return NewTensor(raw_y_data)
def backward(ctx, out_grad):
# get raw data
raw_out_grad_data = out_grad.data
raw_y_data, = ctx.get_saved_tensors()
# compute gradient of input
raw_in_grad_data = raw_y_data * raw_out_grad_data
return NewTensor(raw_in_grad_data)
The minimum set of operations needed for the autograd system are
neg, add, mul, pow, fill