Skip to content

PyTorch Example*

MNIST Model Compilation Workflow*

This document describes the overall workflow for exporting a PyTorch model, based on the MNIST model, covering model construction, model training, model validation, inference script construction, ScriptModule static model export, and model compilation.

1. Model Construction, Training, and Validation*

Model Construction

A custom MNIST model is constructed by inheriting from the nn.Module base class. The model structure includes convolutional layers, activation layers, and fully connected layers.

import torch
from torch import nn
from torch.nn import functional as F

class MNISTModel(nn.Module):
    def __init__(self):
        super(MNISTModel, self).__init__()
        self.conv1 = nn.Conv2d(1, 8, 1, 1)
        self.conv2 = nn.Conv2d(8, 32, 1, 1)
        self.dropout1 = nn.Dropout(0.25)
        self.fc1 = nn.Linear(25088, 32)
        self.fc2 = nn.Linear(32, 10)

    def forward(self, x):
        x = self.conv1(x)
        x = F.relu(x)
        x = self.conv2(x)
        x = F.relu(x)
        x = self.dropout1(x)
        x = torch.flatten(x, 1)
        x = self.fc1(x)
        x = F.relu(x)
        x = self.fc2(x)
        return x

Model Training

The MNIST training set is loaded using the torchvision library, and after normalization and standardization, it is fed into the network for training.

Model training parameters are as follows:

  • Batch size: 64

  • Number of epochs: 3

  • Optimizer: SGD

  • Learning rate: 0.01

  • Momentum: 0.5

  • Loss function: Negative log-likelihood loss

from torch import optim
from torchvision import datasets, transforms

def train(model, train_loader, optimizer, epoch):
    model.train()
    for batch_idx, (data, target) in enumerate(train_loader):
        optimizer.zero_grad()
        output = model(data)
        loss = criterion(output, target)
        loss.backward()
        optimizer.step()
        if batch_idx % 10 == 0:
            print('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format(
                epoch, batch_idx * len(data), len(train_loader.dataset),
                100. * batch_idx / len(train_loader), loss.item()))

# Data loading and preprocessing
train_loader = torch.utils.data.DataLoader(
    datasets.MNIST('../data', train=True, download=True,
                   transform=transforms.Compose([
                       transforms.ToTensor(),
                       transforms.Normalize((0.1307,), (0.3081,))
                   ])),
    batch_size=64, shuffle=True)

model = MNISTModel()
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(model.parameters(), lr=0.005, momentum=0.5)

for epoch in range(3):
    train(model, train_loader, optimizer, epoch)

torch.save(model.state_dict(), "./mnist.pth")

Model Validation

After model training is completed, the MNIST validation set is loaded to validate the model, and the average loss and result accuracy are output.

def test(model, test_loader):
    model.eval()
    test_loss = 0
    correct = 0
    with torch.no_grad():
        for data, target in test_loader:
            output = model(data)
            test_loss += criterion(output, target).item()  # Sum up batch loss
            pred = output.argmax(dim=1, keepdim=True)  # Get the index of the max log-probability
            correct += pred.eq(target.view_as(pred)).sum().item()

    test_loss /= len(test_loader.dataset)

    print('\nTest set: Average loss: {:.4f}, Accuracy: {}/{} ({:.0f}%)\n'.format(
        test_loss, correct, len(test_loader.dataset),
        100. * correct / len(test_loader.dataset)))

test_loader = torch.utils.data.DataLoader(
    datasets.MNIST('../data', train=False, transform=transforms.Compose([
                       transforms.ToTensor(),
                       transforms.Normalize((0.1307,), (0.3081,))
                   ])),
    batch_size=1000, shuffle=True)

test_model = MNISTModel()
criterion = nn.CrossEntropyLoss(reduction='sum')
test_model.load_state_dict(torch.load("./mnist.pth"))

test(model, test_loader)

2. Exporting the Specified Model File Format*

Inference Model Construction

Due to the limitations of exporting ScriptModule static model files, when the model structure uses specific control flow forms, the user needs to determine the fixed flow required for inference.

Here is an explanation of the above description: specific control flow forms refer to conditional statements such as if and assert in the forward method of an nn.Module subclass, where the condition state can only be determined during inference.

When this situation occurs, the user needs to define a deterministic computation flow based on the original model structure and eliminate these control flow operators.

In this example, the MNIST model does not contain control flow structures, so no special modifications are required. Only the dropout operator, which is not involved in inference computation, is removed. The NPU compiler can handle this on its own; the user may choose whether to remove it, as it does not affect model compilation.

class MNISTModelInference(nn.Module):
    def __init__(self):
        super(MNISTModelInference, self).__init__()
        self.conv1 = nn.Conv2d(1, 8, 1, 1)
        self.conv2 = nn.Conv2d(8, 32, 1, 1)
        self.fc1 = nn.Linear(25088, 32)
        self.fc2 = nn.Linear(32, 10)

    def forward(self, x):
        x = self.conv1(x)
        x = F.relu(x)
        x = self.conv2(x)
        x = F.relu(x)
        x = torch.flatten(x, 1)
        x = self.fc1(x)
        x = F.relu(x)
        x = self.fc2(x)
        return x

ScriptModule Static Model File Export

  • Export Workflow

    • Step 1: Create a PyTorch Module based on the inference model class and load the trained model weights.

    • Step 2: Convert the PyTorch Module to a TorchScript Module using tracing.

    • Step 3: Serialize and output the TorchScript Module to the file "mnist_trace.pth".

inference_model = MNISTModelInference()
inference_model.load_state_dict(torch.load("./mnist.pth"))

input_tensor = torch.randn([1, 1, 28, 28])

traced_model = torch.jit.trace(inference_model, [input_tensor])
traced_model.save("./mnist_trace.pth")

3. Writing the Compilation Configuration File*

Before reading this chapter, please familiarize yourself with the PyTorch Configuration Items in the NPU compiler usage documentation.

CORENAME

For gx830X, this is fixed to APUS.

FRAMEWORK

This example uses a PyTorch model, so configure it as PT.

MODEL_FILE

As described in ScriptModule Static Model File Export, the model file required by the NPU compiler is "./mnist_trace.pth".

IN_FEATS_FILE

In this example, the input feature file name is configured as feats.txt.

QUANT_FILE

In this example, the output quantization file name is configured as quant.txt.

OUTPUT_TYPE

For gx830X, this is fixed to c_code.

OUTPUT_FILE

In this example, the output file name is configured as mnist.h.

INPUT_OPS

The input tensor shape for MNIST model inference is [1, 1, 28, 28].

FUSE_BN, COMPRESS

In this example, BN fusion is disabled, and fully connected layer weight compression is enabled.

MAX_CACHE_SIZE, USE_DATA_CACHE

Not used in the current scenario.

The specific configuration file is as follows

mnist_config.yaml
CORENAME: APUS
FRAMEWORK: PT
MODEL_FILE: mnist_trace.pth
IN_FEATS_FILE: feats.txt
QUANT_FILE: quant.yaml
OUTPUT_TYPE: c_code
OUTPUT_FILE: mnist.h

INPUT_OPS:
    0: [1, 1, 28, 28]

FUSE_BN: false
COMPRESS: true

4. Model Compilation*

First, use the gxnpuc tool to compile and generate the quantization file.

$ gxnpuc mnist_config.yaml -q

Next, use the gxnpuc tool to compile and generate the NPU file mnist.h.

$ gxnpuc mnist_config.yaml

The memory information required by the model is printed:

------------------------
Memory allocation info:
Mem1(data): 112896
Mem2(instruction): 260
Mem3(in): 1568
Mem4(out): 20
Mem5(cache): 0
Mem6(weights): 803828
Total NPU Size (Mem0+Mem1+Mem2+Mem5+Mem6): 916984
Total Memory Size: 918572
------------------------
Compile OK.

The memory regions are described as follows:

Memory Region Description
Mem1(data) Intermediate data memory
Mem2(instruction) Instruction memory
Mem3(in) Input data memory
Mem4(out) Output data memory
Mem5(cache) Weight and data memory in SRAM
Mem6(weights) Weight memory