Quickstart

Install the SDK and train a model end to end.

Install

pip install deepgate

Requires Python 3.10+. The import name is dg (PyTorch and torchao are installed with it).

Train and export

This example trains a small convolutional classifier on CIFAR-10, quantizes it, and exports a compiler schema.

import json

import torch
import torch.nn as nn
import torch.nn.functional as F
import torchvision

import dg

device = "cuda" if torch.cuda.is_available() else "cpu"
EPOCHS = 30


# 1. Data: CIFAR-10 as uint8 tensors (dg.Norm normalizes inside the graph)
def cifar10(train):
    ds = torchvision.datasets.CIFAR10("~/data", train=train, download=True)
    x = torch.from_numpy(ds.data).permute(0, 3, 1, 2).contiguous()  # NHWC -> NCHW
    return x, torch.tensor(ds.targets)


(train_x, train_y), (test_x, test_y) = cifar10(train=True), cifar10(train=False)
train_ds = torch.utils.data.TensorDataset(train_x, train_y)


def augment(x):
    x = F.pad(x, (4, 4, 4, 4))
    i, j = torch.randint(0, 9, (2,)).tolist()
    x = x[:, :, i:i + 32, j:j + 32]
    return torch.where(torch.rand(len(x), 1, 1, 1, device=x.device) < 0.5, x.flip(-1), x)


# 2. Model with dg.Norm and dg.Flatten; the rest is stock PyTorch
def block(cin, cout, stride):
    return nn.Sequential(
        nn.Conv2d(cin, cout, 3, stride=stride, padding=1, bias=False),
        nn.BatchNorm2d(cout),
        nn.ReLU(),
    )


class ConvNet(nn.Module):
    def __init__(self, w=32):
        super().__init__()
        self.norm = dg.Norm(120.7, 1 / 64.0)
        self.features = nn.Sequential(
            block(3, w, 1), block(w, 2 * w, 2), block(2 * w, 4 * w, 2)
        )
        self.flatten = dg.Flatten()
        self.fc = nn.Linear(4 * w * 8 * 8, 10)

    def forward(self, x):
        return self.fc(self.flatten(self.features(self.norm(x))))


model = ConvNet().to(device)

# 3. Train with cosine-decaying learning rate
loader = torch.utils.data.DataLoader(train_ds, batch_size=128, shuffle=True)
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
schedule = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, EPOCHS * len(loader))
for epoch in range(EPOCHS):
    for x, y in loader:
        x, y = augment(x.to(device)), y.to(device)
        optimizer.zero_grad()
        loss = F.cross_entropy(model(x), y)
        loss.backward()
        optimizer.step()
        schedule.step()
    if epoch % 5 == 4:
        print(f"epoch {epoch + 1}: loss {loss.item():.3f}")


# 4. Quantize (calibrate on training set) and export
@torch.no_grad()
def accuracy(m, dev, batch=512):
    hits = sum((m(test_x[i:i + batch].to(dev)).argmax(1) == test_y[i:i + batch].to(dev))
               .sum().item() for i in range(0, len(test_x), batch))
    return 100 * hits / len(test_x)


quantized = dg.post_training_quantize(model.eval(), train_ds, num_samples=1024)
schema = dg.export(quantized)

print(f"accuracy: float {accuracy(model, device):.1f}%  quantized {accuracy(quantized, device):.1f}%")
json.dump(schema, open("schema.json", "w"))

Upload schema.json to the DeepGate developer platform to compile or benchmark your model.