跳转至

PyTorch示例*

MNIST 模型编译流程*

本文基于 MNIST 模型,从模型搭建、模型训练、模型验证、推理脚本构建、ScriptModule静态模型导出、模型编译等方面,简述 PyTorch 模型导出的整体流程

1. 模型构建、训练及验证*

模型构建

通过继承 nn.Module 基类的方式来构建自定义 MNIST 模型,模型结构包含卷积层、激活层以及全连接层。

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

模型训练

基于 torchvision 库加载 MNIST 训练集,并进行归一化、标准化处理后输入至网络训练

模型训练操参数如下:

  • 批次大小: 64

  • 训练轮数: 3

  • 优化器: SGD

  • 学习率: 0.01

  • 动量: 0.5

  • 损失函数: 负对数似然损失函数

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()))

# 数据加载和预处理
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")

模型验证

模型训练完成后,加载 MNIST 验证集对模型进行验证处理,并输出平均损失和结果精度

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()  # 将批次的损失相加
            pred = output.argmax(dim=1, keepdim=True)  # 获取概率最高的预测结果
            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. 导出指定格式模型文件*

推理模型构建

由于 ScriptModule 静态模型文件导出的限制,当模型结构中使用了特定的控制流形式时,那么需要用户确定推理时所需的固定流向;

这里解释下上述描述: 特定的控制流形式,即在模型定义中,nn.Module 子类 forward 方法中的存在如 if、assert 条件语句,且条件状态在推理时才能确定。

当出现该情况,用户需要在原始模型结构上,给出确定的计算流程,并消除这些控制流算子。

本例中,MNIST 模型中不存在控制流结构,因此没有特别改动,只将不参与推理计算的 dropout 算子删去,NPU 编译器可自行处理,用户可选择是否删除,不影响模型编译

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 静态模型文件导出

  • 导出流程

    • 步骤一: 基于推理模型类创建 PyTorch Module,并加载训练后的模型权重数据

    • 步骤二: 将 PyTorch Module 通过 tracing 方式转换为 TorchScript Module

    • 步骤三: 序列化输出 TorchScript Module 到文件 "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. 编译配置文件的编写*

阅读本章节前,请提前熟悉 NPU 编译器使用文档中 PyTorch 配置项

CORENAME

gx830X 固定配置为 APUS

FRAMEWORK

本例为 PyTorch 模型,故配置 PT

MODEL_FILE

ScriptModule 静态模型文件导出中可知,NPU 编译器所需的模型文件名为 "./mnist_trace.pth"

IN_FEATS_FILE

本例中配置输入特征文件名为feats.txt

QUANT_FILE

本例中配置输输出量化文件名为quant.txt

OUTPUT_TYPE

gx830X 固定配置为 c_code

OUTPUT_FILE

本例中配置输出文件名为 mnist.h

INPUT_OPS

MNIST 模型推理时的输入张量 shape 为 [1, 1, 28, 28]

FUSE_BN, COMPRESS

本例中,BN融合不开启,全连接层权重压缩开启

MAX_CACHE_SIZE, USE_DATA_CACHE

当前场景下暂不使用

具体配置文件如下

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. 模型编译*

首先,使用 gxnpuc 工具编译生成量化文件

$ gxnpuc mnist_config.yaml -q

其次,使用 gxnpuc 工具编译生成 NPU 文件mnist.h

$ gxnpuc mnist_config.yaml

打印出模型需要的内存信息:

------------------------
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.

各内存区域说明如下:

内存区域 说明
Mem1(data) 中间数据内存
Mem2(instruction) 指令内存
Mem3(in) 输入数据内存
Mem4(out) 输出数据内存
Mem5(cache) SRAM中权重和data内存
Mem6(weights) 权重内存