关于rk3588s使用facenet-pytorch-main进行onnx的转换以及RKNN生成操作

MobilefaceNet特征提取，地址为：

https://github.com/bubbliiiing/facenet-pytorch

由于mobilefaceNet训练数据集很大，就使用预训练模型转换。先使用netron查看下pth，首先你得看你的pth模型是只有参数还是说既有参数，也有模型结构。如果模型是黑白的，就说明只有参数，没有模型结构，如下图。：

如果是彩色的，就是既有参数也有模型结构的，如下图：

很明显我们在官网的百度网盘下到的是不包含网络结构的pth，所以要加上模型架构进行onnx转换，上菜：

import torch
import nets.facenet as models
model_path='./facenet_mobilenet.pth' #模型路径

model=models.Facenet(mode='val')
device = torch.device('cpu')
model.load_state_dict(torch.load(model_path,map_location=device),strict=False) #模型加载
net=model.eval()
example=torch.rand(1,3,160,160) #给定一个输入
torch.onnx.export(model,(example),'./facenet_mobilenet1.onnx',verbose=True,opset_version=12) #导出onnx模型

接着我们运行代码进行转换，转换完获得一个onnx，但是通过工具看到里面有一个很奇怪的op:ReduceL2类型的ReduceL2_90，如下图：
显然这是一个网络前向传播结束时执行的L2正则化。所以我们回到facenet的的网络代码部分：
·

def forward(self, x, mode = "predict"):
    if mode == 'predict':
        x = self.backbone(x)
        x = self.avg(x)
        x = x.view(x.size(0), -1)
        x = self.Dropout(x)
        x = self.Bottleneck(x)
        x = self.last_bn(x)
        x = F.normalize(x, p=2, dim=1)
        return x
    x = self.backbone(x)
    x = self.avg(x)
    x = x.view(x.size(0), -1)
    x = self.Dropout(x)
    x = self.Bottleneck(x)
    before_normalize = self.last_bn(x)
    
    x = F.normalize(before_normalize, p=2, dim=1)
    cls = self.classifier(before_normalize)
    return x, cls

def forward_feature(self, x):
    x = self.backbone(x)
    x = self.avg(x)
    x = x.view(x.size(0), -1)
    x = self.Dropout(x)
    x = self.Bottleneck(x)
    before_normalize = self.last_bn(x)
    x = F.normalize(before_normalize, p=2, dim=1)
    return before_normalize, x

def forward_classifier(self, x):
    x = self.classifier(x)
    return x

显然，在经过last_bn的操作后又进行了一次normalize操作，所以出现了ReduceL2_90。我们把这行注释了就行：

# x = F.normalize(x, p=2, dim=1)

随后我们重新生成onnx，然后再看下结构：

可以看到ReduceL2_90已经不见了，至此onnx转换完毕。然后运行以下代码进行RKNN转换：

import numpy as np
import cv2
import os
from rknn.api import RKNN
from PIL import Image
from sklearn import preprocessing
from scipy.spatial.distance import pdist
os.environ['RKNN_DRAW_DATA_DISTRIBUTE']="1"

if __name__ == '__main__':

im_file = './9.jpg'
BUILD_QUANT = False
RKNN_MODEL_PATH = './mobilefacenet.rknn'
if BUILD_QUANT:
    RKNN_MODEL_PATH = './mobilefacenet_quant.rknn'

# Create RKNN object
rknn = RKNN()

NEED_BUILD_MODEL = True
if NEED_BUILD_MODEL:
    print('--> config model')

    rknn.config(mean_values=[[0, 0, 0]], std_values=[[255, 255, 255]], target_platform='rk3588')

    print('done')
    print('--> Loading model')



    ret = rknn.load_onnx(model='./weights/facenet_mobilenet1.onnx', )
    if ret != 0:
        print('Load retinaface failed!')
        exit(ret)


    print('done')

    # Build model
    print('--> Building model')
    ret = rknn.build(do_quantization=BUILD_QUANT, dataset='./dataset.txt')
    if ret != 0:
        print('Build model failed!')
        exit(ret)
    print('done')

    if BUILD_QUANT:
        print('--> Accuracy analysis')
        rknn.accuracy_analysis(inputs='./dataset.txt',output_dir="./result",target='rk3588')
        print('done')

    # Export rknn model


    print('--> Export RKNN model')
    ret = rknn.export_rknn(RKNN_MODEL_PATH)

    if ret != 0:
        print('Export rknn failed!')
        exit(ret)
    print('done')
else:
    # Direct load rknn model
    print('Loading RKNN model')
    ret = rknn.load_rknn(RKNN_MODEL_PATH)
    if ret != 0:
        print('load rknn model failed.')
        exit(ret)
    print('done')


print('--> Init runtime environment')
ret = rknn.init_runtime()
# ret = rknn.init_runtime(target='rk3588')
if ret != 0:
    print('Init runtime environment failed')
    exit(ret)
print('done')

# Set inputs
img = cv2.imread(im_file)
img = cv2.resize(img, (160,160))
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)


image_1 = Image.open(im_file)
image_1 = image_1.resize((160,160), Image.BICUBIC)
img = np.asarray(image_1, np.uint8)
print(img.shape)


# inference

print('--> inference')
outputs = rknn.inference(data_format='nhwc',inputs=[img])
print('done')



print(outputs)
image_1 = Image.open("1_001.jpg")
image_1 = image_1.resize((160,160), Image.BICUBIC)
img = np.asarray(image_1, np.uint8)
outputs0 = np.array(rknn.inference(data_format='nhwc', inputs=[img])[0])


image_1 = Image.open("1_002.jpg")
image_1 = image_1.resize((160,160), Image.BICUBIC)
img = np.asarray(image_1, np.uint8)
outputs1 = np.array(rknn.inference(data_format='nhwc', inputs=[img])[0])


l1 = np.linalg.norm(outputs1 - outputs0, axis=1)
print("l1 %f"%l1)
cosSim = 1 - pdist(np.vstack([outputs1, outputs0]), 'cosine')
print("pdist %f"%cosSim)
outputs1 = preprocessing.normalize(outputs1, norm='l2')
outputs0 = preprocessing.normalize(outputs0, norm='l2')
l1 = np.linalg.norm(outputs1 - outputs0, axis=1)
print("after l2 l1 %f" % l1)


rknn.release()