A Minimalistic Auto-Diff Optimization Framework for Teaching and Understanding Pytorch
✔️ Teaching auto-diff and backpropagation
✔️ Introducing or learning libTorch/Pytorch
✔️ Showing that auto-diff optimization is not magic
✖️ Fitting models to your own data
✖️ Replacing Pytorch or any other optimization framework
- Forward and backward implementations of basic operations (+, -, *, sum, square)
- Automatic graph generation during forward
- Backpropagation through the generated graph
- Almost the same syntax as libTorch
- No dependencies except the std-library
#include "tiny_torch.h"
int main()
{
// The data tensors
tinytorch::Tensor observation = tinytorch::rand(10);
tinytorch::Tensor target = tinytorch::rand(10);
// The parameters of the model
std::vector<tinytorch::Tensor> params;
for (int i = 0; i < 4; ++i)
{
params.push_back(tinytorch::rand(10));
MakeParameter(params.back());
}
// The model itself
auto model = [&](tinytorch::Tensor x) -> tinytorch::Tensor
{
x = x * params[0];
x = x + params[1];
x = x * params[2];
x = x + params[3];
return x;
};
// Create a simple optimizer
tinytorch::SGDOptimizer optim(params, 0.1);
// Optimize the model for 50 iterations
for (int i = 0; i < 50; ++i)
{
optim.ZeroGrad();
auto prediction = model(observation);
auto loss = sum(square(prediction - target));
backward(loss);
optim.Step();
std::cout << "Step " << i << " Loss: " << loss << std::endl;
}
return 0;
}
Example Output:
Step 0 Loss: [Tensor s=1]: 5.73343
Step 1 Loss: [Tensor s=1]: 2.18133
Step 2 Loss: [Tensor s=1]: 1.17362
Step 3 Loss: [Tensor s=1]: 0.321065
Step 4 Loss: [Tensor s=1]: 0.0889282
Step 5 Loss: [Tensor s=1]: 0.0313245
Step 6 Loss: [Tensor s=1]: 0.0113702
Step 7 Loss: [Tensor s=1]: 0.0036413
Step 8 Loss: [Tensor s=1]: 0.00149557
Step 9 Loss: [Tensor s=1]: 0.000541403
Step 10 Loss: [Tensor s=1]: 0.000217643
Step 11 Loss: [Tensor s=1]: 8.67411e-05
Step 12 Loss: [Tensor s=1]: 3.2838e-05
Step 13 Loss: [Tensor s=1]: 1.26026e-05
...
cd TinyTorch
mkdir build && cd build
cmake ..
make
cd TinyTorch
./build/bin/tt_sample