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.maxpool1 = nn.MaxPool2d(2, 2)
self.maxpool2 = nn.MaxPool2d(2, 2)
self.dropout1 = nn.Dropout(0.25)
self.fc1 = nn.Linear(1568, 32)
self.fc2 = nn.Linear(32, 10)
def forward(self, x):
x = self.conv1(x)
x = F.relu(x)
x = self.maxpool1(x)
x = self.conv2(x)
x = F.relu(x)
x = self.maxpool2(x)
x = self.dropout1(x)
x = torch.flatten(x, 1)
x = self.fc1(x)
x = F.relu(x)
x = self.fc2(x)
output = F.log_softmax(x, dim=1)
return output
模型训练
基于 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 = F.nll_loss(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()
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 += F.nll_loss(output, target, reduction='sum').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()
test_model.load_state_dict(torch.load("./mnist.pth"))
test(model, test_loader)
2. 导出指定格式模型文件*
推理模型构建
由于 ScriptModule 静态模型文件导出的限制,当模型结构中使用了特定的控制流形式时,那么需要用户确定推理时所需的固定流向;
这里解释下上述描述: 特定的控制流形式,即在模型定义中,nn.Module 子类 forward 方法中的存在如 if、assert 条件语句,且条件状态在推理时才能确定。
当出现该情况,用户需要在原始模型结构上,给出确定的计算流程,并消除这些控制流算子。
本例中,MNIST 模型中不存在控制流结构,因此没有特别改动,只将不参与推理计算的 dropout 算子删去,NPU 编译器可自行处理,用户可选择是否删除,不影响模型编译
NPU 编译器暂不支持 softmax 算子,故推理模型中使用 npu_softmax 替换
npu_softmax的构建方式,详见Softmax章节
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.maxpool1 = nn.MaxPool2d(2, 2)
self.maxpool2 = nn.MaxPool2d(2, 2)
self.fc1 = nn.Linear(1568, 32)
self.fc2 = nn.Linear(32, 10)
# for softmax
self.partitions = split_and_factorize(10)
def forward(self, x):
x = self.conv1(x)
x = F.relu(x)
x = self.maxpool1(x)
x = self.conv2(x)
x = F.relu(x)
x = self.maxpool1(x)
x = torch.flatten(x, 1)
x = self.fc1(x)
x = F.relu(x)
x = self.fc2(x)
x = npu_softmax(x, self.partitions)
return x
-
导出流程
-
步骤一: 基于推理模型类创建 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 配置项
CORE_NAME
gx8002 固定配置为 GRUS
NPU_UNIT
gx8002 固定配置为 NPU32
FRAMEWORK
本例为 PyTorch 模型,故配置 PT
MODEL_FILE
由 ScriptModule 静态模型文件导出中可知,NPU 编译器所需的模型文件名为 "./mnist.pth"
OUTPUT_TYPE
gx8002 固定配置为 c_code
OUTPUT_FILE
本例中配置输出文件名为 mnist.h
INPUT_OPS
MNIST 模型推理时的输入张量 shape 为 [1, 1, 28, 28]
INPUT_NCX_TO_NXC
本例中输入张量不进行数据排布格式转换
FP16_OUT_OPS
本例中无状态输出张量,不需要配置该参数
FUSE_BN, COMPRESS, CONV2D_COMPRESS
本例中,BN融合不开启,全连接层和卷积层权重压缩开启
EXCLUDE_COMPRESS_OPS, WEIGHT_MIN_MAX, WEIGHT_CACHE_SIZE
当前场景下暂不使用
具体配置文件如下
CORENAME: GRUS
NPU_UNIT: NPU32
FRAMEWORK: PT
MODEL_FILE: mnist_trace.pth
OUTPUT_TYPE: c_code
OUTPUT_FILE: mnist.h
INPUT_OPS:
0: [1, 1, 28, 28]
INPUT_NCX_TO_NXC: []
FP16_OUT_OPS: []
FUSE_BN: false
COMPRESS: true
CONV2D_COMPRESS: true
4. 模型编译*
使用 gxnpuc 工具编译
$ gxnpuc mnist_config.yaml
生成 NPU 文件mnist.h
,并且打印出模型需要的内存信息:
------------------------
Memory allocation info:
Mem0(ops): 0
Mem1(data): 28224
Mem2(instruction): 836
Mem3(in): 1568
Mem4(out): 40
Mem5(tmp content): 0
Mem6(weights): 56052
Total NPU Size (Mem0+Mem1+Mem2+Mem5+Mem6): 85112
Total Memory Size: 86720
------------------------
Compile OK.
各内存区域说明如下:
内存区域 | 说明 |
---|---|
Mem0(ops) | 暂未使用 |
Mem1(data) | 中间数据内存 |
Mem2(instruction) | 指令内存 |
Mem3(in) | 输入数据内存 |
Mem4(out) | 输出数据内存 |
Mem5(tmp content) | SRAM中权重内存 |
Mem6(weights) | 权重内存 |