Torchtrace
Suit for the new learners to have a basic understanding what is happending behind the library with auto-grad supported like tensorflow, torch.
The basic data type in torchtrace is tracer, inherited from torch.Tensor.Thanks to that, all the arithmetic operation, type conversion are not in our concern.
Note that, all the auto-grad function is implemented in torchtrace itself, we only use a few traits of torch, see the torch apis we used in Whitelist.txt(excludetorchtrace.test, which uses auto-grad traits in torch for unit-test ) .
What is good in torchtrace
- Thanks to the
torch-like api's fashion, we can easily migrate atorchcode into atorchtracecode, if all the operations are supported intorchtrace. - For the operation that are not supported in
torchtraceright now, we can easily implement one. Please see the implementation intorchtrace.nnfor details. -
torchtraceis easy and easy to read. Right now,torchtraceonly consider the sequential neural network, which makes the dynamic computation graph pretty straightforward. See more intorchtrace.base
Code snippets
A.
import torchtrace
from torchtrace import nn
a = torchtrace.tracer([1,2,3,4])
a.requires_grad_()
b = a.Sum()
b.backward()
print(a.grad)
# output
# tensor([1., 1., 1., 1.])B.
import torch
import torchtrace
from torchtrace.nn as Linear, ReLU
import torchtrace.optim as optim
model = nn.Sequential(
Linear(8,256),
ReLU(),
Linear(256,128),
ReLU(),
Linear(128,64),
ReLU(),
Linear(64, 4),
)
print(model)
optimizer = optim.RMSprop(model.parameters(), lr=5e-4)
mse = nn.MSE()
x = torch.rand(16,8).float()
target = torch.rand(16,1)
for i in range(10):
out = model(x)
_, action = out.max(1)
gather = nn.Gather_last(action)
out = gather(out)
loss = mse(loss, target)
print(loss)
optimizer.zero_grad()
loss.backward()
optimizer.step()C.
import torch
from torchtrace.nn import Conv2d, ReLU, Linear, Model, Sequential
class myConv(Model):
def __init__(self):
self.conv = Sequential(
Conv2d(3, 8, kernel_size=8, stride=4), ReLU(),
Conv2d(8, 4, kernel_size=4, stride=2), ReLU(),
Conv2d(4, 2, kernel_size=3, stride=1), ReLU(),
)
self.fc = Linear(7*7*2, 1)
super(myConv, self).__init__()
def construct(self):
ops = [self.conv, self.fc]
return ops
def forward(self, x):
x = self.conv(x)
x = x.View(x.shape[0], -1)
return self.fc(x)
model_trace = myConv()
print(model_trace)
for para in model_trace.parameters():
para.data.fill_(0.001)
x = torch.rand(3, 3, 84, 84).float()
out = model_trace(x).Sum()
out.backward(show=True)
# output
Sequence(
(0):Sequence(
(0):Conv2d 3 -> 8
(1):ReLU()
(2):Conv2d 8 -> 4
(3):ReLU()
(4):Conv2d 4 -> 2
(5):ReLU()
)
(1):Linear layer (98 X 1)
)
BackProp:
↘︎ Sum()
↘︎ Linear layer (98 X 1)
↘︎ View((3, -1))
↘︎ ReLU()
↘︎ Conv2d 4 -> 2 GRAD:torch.Size([3, 2, 7, 7]) -> torch.Size([3, 4, 9, 9])
↘︎ ReLU()
↘︎ Conv2d 8 -> 4 GRAD:torch.Size([3, 4, 9, 9]) -> torch.Size([3, 8, 20, 20])
↘︎ ReLU()
↘︎ Conv2d 3 -> 8 GRAD:torch.Size([3, 8, 20, 20]) -> torch.Size([3, 3, 84, 84])
↘︎ Inputsee more in torchtrace.test and examples
