Getting started¶
After installing torch_tvarant, verify the device and run a small model.
Verify installation¶
import torch
import torch_tvarant
print(torch.tvarant.is_available()) # True
print(torch.tvarant.backend()) # 'sim' or 'opencl'
x = torch.ones(4, device="tvarant")
print(x.device) # tvarant:0
Run a small MLP¶
import torch
import torch.nn as nn
import torch_tvarant
model = nn.Sequential(
nn.Linear(16, 32),
nn.ReLU(),
nn.Linear(32, 4),
).to("tvarant")
x = torch.randn(2, 16, device="tvarant")
y = model(x)
print(y.shape) # torch.Size([2, 4])
Compile for inference¶
import torch_tvarant
compiled = torch_tvarant.compiler.compile(model)
y = compiled(x)
See JIT Compiler for fusion details.
Learn more¶
| Topic | Page |
|---|---|
| Device helpers | torch.tvarant |
| Supported ops | Supported ops |
| Kernels | Kernels |
| C++ extension | C++ API |
| FPGA path | FPGA / OpenCL |