在STM32上运行keyword spotting(三)模型搭建以及训练测试
搭建一个怎样的模型
KWS模型结构属于比较简单的模型结构,但是为了少走弯路,我计划使用现成的结构,我从这个演示视频参考而来:点击我跳转,这个KWS项目运行在AT32F403上,其网络模型结构为一个64个特征的普通卷积层,然后重复四次的DS-CNN卷积,在每次卷积后都进行一次relu,最后进行一次池化,用来减少全连接层的参数。
模型部分代码如下:
class Net(torch.nn.Module):
def __init__(self):
super(Net, self).__init__()
self.conv1 = torch.nn.Conv2d(in_channels=1, out_channels=64,
kernel_size=(4, 10), padding=(1, 4), stride=(2, 2))
self.fc1 = torch.nn.Linear(in_features=12*5*64, out_features=6)
self.avgpool = torch.nn.AvgPool2d(
kernel_size=(1*4), stride=(1, 1), padding=(0, 0))
self.depthwise = torch.nn.Conv2d(in_channels=64, out_channels=64, kernel_size=(
3, 3), padding=(1, 1), stride=(1, 1), groups=64)
self.pointwise = torch.nn.Conv2d(
in_channels=64, out_channels=64, kernel_size=1)
self.pooling = torch.nn.MaxPool2d(kernel_size=(2, 1))
def dw_conv(self, x):
x = self.depthwise(x)
x = F.relu(x)
x = self.pointwise(x)
x = F.relu(x)
return x
def forward(self, x):
x = self.conv1(x)
x = F.relu(x)
x = self.dw_conv(x)
x = self.dw_conv(x)
x = self.dw_conv(x)
x = self.dw_conv(x)
#x = self.avgpool(x)
x = self.pooling(x)
x = x.view(int(x.size(0)), -1)
x = self.fc1(x)
return x
为什么使用更换CNN为DS-CNN(深度可分离卷积)
对于一张5×5像素、三通道(shape为5×5×3),经过3×3卷积核的卷积层(假设输出通道数为4,则卷积核shape为3×3×3×4,最终输出4个Feature Map,如果有same padding则尺寸与输入层相同(5×5),如果没有则为尺寸变为3×3,可以看到参数数量成几何倍的增长,计算量也呈几何倍的增长,示意图如下:
更换为DS-CNN,分别对三个通道上提取特征,得到三张特征图,在对三张图进行逐点卷积,有几个卷积核就输出几张特征图,与常规来讲,降低了大量不必要的特征图,节省了参数数量与计算量,示意图如下:
回顾一下,常规卷积的参数个数为:
N_std = 4 × 3 × 3 × 3 = 108
Separable Convolution的参数由两部分相加得到:
N_depthwise = 3 × 3 × 3 = 27
N_pointwise = 1 × 1 × 3 × 4 = 12
N_separable = N_depthwise + N_pointwise = 39
相同的输入,同样是得到4张Feature map,Separable Convolution的参数个数是常规卷积的约1/3。因此,在参数量相同的前提下,采用Separable Convolution的神经网络层数可以做的更深,计算量和对内存占用更少,这在MCU设备上是相当重要的。
模型训练
我们已经知道了我们使用怎样的模型结构,就可以直接来训练了,迭代50次看看效果,并且保存loss值最小的模型到本地(其实不建议这么做,出来的模型可能是过拟合的模型)
*建立模型
train_model = Net()
*构建损失
criterion = torch.nn.CrossEntropyLoss()
*构建优化器
optimizer = torch.optim.RMSprop(train_model.parameters(),lr=0.001,alpha=0.9)
def train(model , epochs):
epoches = epochs
best_loss_val = 1
targets = torch.tensor([0, 1, 2, 3, 4,5])
for epoch in range(epoches):
for i in range(0, 1700):
inputs = torch.Tensor(([background_mfcc[i]],[one_mfcc[i]], [two_mfcc[i]], [
three_mfcc[i]], [four_mfcc[i]], [five_mfcc[i]]))
outputs = model(inputs)
loss = criterion(outputs, targets)
optimizer.zero_grad()
loss.backward()
optimizer.step()
loss_val = loss.data.numpy()
if loss_val<best_loss_val:
torch.save(model, 'best_model.pth')
best_loss_val = loss_val
print("saved as \"best_model.pth\"")
print("best_loss_val:",best_loss_val)
print('Epoch: ', epoch, '| train loss: %.10f' % loss_val)
train(train_model,50)
训练完成后生成了best_model.pth到本地,我们可以测试这个模型
模型测试
读取保存在本地的模型文件,使用300个测试样本进行测试,分别打印每个样本的错误数量,准确率,以及模型结构和参数大小
def test(model):
nothing = 0
one = 0
two = 0
three = 0
four = 0
five = 0
for i in range(1700, 2000):
inputs = torch.Tensor(([background_mfcc[i]],[one_mfcc[i]], [two_mfcc[i]], [
three_mfcc[i]], [four_mfcc[i]], [five_mfcc[i]]))
outputs = model(inputs)
pred_y = torch.max(outputs, 1)[1].data.numpy().squeeze()
if pred_y[0] != 0 :
nothing +=1
if pred_y[1] != 1 :
one +=1
if pred_y[2] != 2 :
two +=1
if pred_y[3] != 3 :
three +=1
if pred_y[4] != 4 :
four +=1
if pred_y[5] != 5 :
five +=1
print(nothing,one,two,three,four,five)
print((1-nothing/300),(1-one/300),(1-two/300),(1-three/300),(1-four/300),(1-five/300))
test_model = torch.load('best_model.pth')
test(test_model)
summary(test_model, input_size=[(1, 49, 10)],device="cpu")
打印的log:
0 10 19 38 12 33 测试过程中,每个样本有多少个错误
1.0 0.9666666666666667 0.9366666666666666 0.8733333333333333 0.96 0.89 测试结果准确率,可以看出来基本都在90%左右
Layer (type) Output Shape Param #
Conv2d-1 [-1, 64, 24, 5] 2,624
Conv2d-2 [-1, 64, 24, 5] 640
Conv2d-3 [-1, 64, 24, 5] 4,160
Conv2d-4 [-1, 64, 24, 5] 640
Conv2d-5 [-1, 64, 24, 5] 4,160
Conv2d-6 [-1, 64, 24, 5] 640
Conv2d-7 [-1, 64, 24, 5] 4,160
Conv2d-8 [-1, 64, 24, 5] 640
Conv2d-9 [-1, 64, 24, 5] 4,160
MaxPool2d-10 [-1, 64, 12, 5] 0
Linear-11 [-1, 6] 23,046
Total params: 44,870
Trainable params: 44,870
Non-trainable params: 0
Input size (MB): 0.00
Forward/backward pass size (MB): 0.56
Params size (MB): 0.17
Estimated Total Size (MB): 0.73
转换为ONNX模型
开放神经网络交换(Open Neural Network Exchange)简称ONNX,是微软和Facebook提出用来表示深度学习模型的开放格式。所谓开放就是ONNX定义了一组和环境,平台均无关的标准格式,来增强各种AI模型的可交互性。换句话说,无论你使用何种训练框架训练模型(比如TensorFlow/Pytorch/OneFlow/Paddle),在训练完毕后你都可以将这些框架的模型统一转换为ONNX这种统一的格式进行存储。注意ONNX文件不仅仅存储了神经网络模型的权重,同时也存储了模型的结构信息以及网络中每一层的输入输出和一些其它的辅助信息。后文中运行在STM32上,在cube中载入模型时可以选择onnx模型,我们来把我们的best_model.pth转换为best_model.onnx
def convert_onnx(model):
# set the model to inference mode
model.eval()
# Let's create a dummy input tensor
dummy_input = torch.randn(1, 1, 49, 10, requires_grad=True)
# Export the model
torch.onnx.export(model, # model being run
dummy_input, # model input (or a tuple for multiple inputs)
"best_model.onnx", # where to save the model
export_params=True, # store the trained parameter weights inside the model file
opset_version=11, # the ONNX version to export the model to
do_constant_folding=True, # whether to execute constant folding for optimization
input_names = ['input'], # the model's input names
output_names = ['output'], # the model's output names
dynamic_axes={'input' : {0 : 'batch_size'}, # variable length axes
'output' : {0 : 'batch_size'}})
print(" ")
print('Model has been converted to ONNX')
convert_onnx(test_model)
转换成ONNX模型格式,可以在这个网站更直观的查看我们的模型结构以及参数,可以看到完全符合我们的预期:
本章节结束,下一章节将阐述使用cube.ai在STM32上运行测试,以及优化时间
本作品采用 知识共享署名-相同方式共享 4.0 国际许可协议 进行许可。