Implement a minimal library from scratch to help understand the Dynamic Computaional Graph of PyTorch.
- support CUDA
-
Create virtual environment
python3 -m venv minitorch-env
-
Activate virtual environment
source minitorch-env/bin/activate
-
Install dependencies
pip install -r requirements.txt
-
Clone the codebase
git clone git@github.com:zhouzaida/minitorch.git
-
Install or develop
python setup.py install # or python setup.py develop
-
create Tensor
from minitorch import Tensor t1 = Tensor(2.0) t2 = Tensor(3.0) t3 = t1 + t2 print(t3) # Tensor(3.0, requires_grad=False)
-
autograd
from minitorch import Tensor t1 = Tensor(2.0, requires_grad=True) t2 = Tensor(3.0) t3 = t1 + t2 t4 = t1 * t3 t4.backward() print(f"t1 grad: {t1.grad}") # t1 grad: Tensor(7.0, requires_grad=False) print(f"t2 grad: {t2.grad}") # t2 grad: None
-
gradient for broadcast
from minitorch import Tensor t1 = Tensor([1.0, 2.0], requires_grad=True) t2 = Tensor(2.0, requires_grad=True) t3 = t1 + t2 t3.backward(Tensor([1.0, 1.0])) print(f"t1 grad: {t1.grad}") # t1 grad: Tensor([1., 1.], requires_grad=False) print(f"t2 grad: {t2.grad}") # t2 grad: Tensor(2.0, requires_grad=False)
-
create neural network
import minitorch import minitorch.nn as nn input = minitorch.rand(2, 3) linear = nn.Linear(3, 5, bias=True) output = nn.linear(input) print(f"output: {output}") class Model(nn.Module): def __init__(self): super().__init__() self.linear_1 = nn.Linear(3, 5, bias=True) self.linear_2 = nn.Linear(5, 6) def forward(self, input): output = self.linear_1(input) output = self.linear_2(output) return output input = minitorch.rand(2, 3) model = Model() output = model(input) print(f"output: {output}") for name, module in model.named_modules(prefix='model'): print(f"{name}: {module}")