RK3568笔记三：基于ResNet18的Cifar-10分类识别训练部署

若该文为原创文章，转载请注明原文出处。

本篇文章参考的是野火-lubancat的rk3568教程，本篇记录了在正点原子的ATK-DLK3568部署。

一、介绍

        ResNet18 是一种卷积神经网络，它有 18 层深度，其中包括带有权重的卷积层和全连接层。它是ResNet 系列网络的一个变体，使用了残差连接（residual connection）来解决深度网络的退化问题。 

       ResNet（Residual Neural Network）由微软研究院的 Kaiming He 等人在 2015 年提出，ResNet 的结 构可以极快的加速神经网络的训练，模型的准确率也有比较大的提升。 ResNet 是一种残差网络，可以把它理解为一个子网络，这个子网络经过堆叠可以构成一个很深的 网络。ResNet 系列有多种变体，如 ResNet18，ResNet34，ResNet50，ResNet101 和 ResNet152，其 网络结构如下：

这里我们主要看下 ResNet18，ResNet18 基本含义是网络的基本架构是 ResNet，网络的深度是 18层，是带有权重的 18 层，不包括 BN 层，池化层。ResNet18 使用的基本残差单元，每个单元由两 个 3x3 卷积层组成，中间有一个 BN 层和一个 ReLU 激活函数。

PyTorch 中的 ResNet18 源码实现：https://github.com/pytorch/vision/blob/main/torchvision/models/ resnet.py

二、环境安装

环境分为两部分：一是训练的环境；二是rknn环境；

rknn环境前面有介绍，自行安装；训练的环境是windows电脑无gpu,使用的是CPU

1、创建虚拟环境

conda create -n ResNet18_env python=3.8 -y

2、激活环境

conda activate ResNet18_env

3、安装环境

pip install torchvision
pip install onnxruntime

三、训练

        自定义一个 ResNet18 网络结构，并使用 CIFAR-10 数据集进行简单测试。CIFAR-10 数据集由 10 个类别的 60000 张 32x32 彩色图像组成，每个类别有 6000 张图像，总共分为 50000 张训练图像和 10000 张测试图像。

resnet18.py

import os
import torchvision
import torch
import torch.nn as nn

#device=torch.device("cuda" if torch.cuda.is_available() else "cpu")
device=torch.device("cpu")

# Transform configuration and data augmentation
transform_train=torchvision.transforms.Compose([
        torchvision.transforms.Pad(4),
        torchvision.transforms.RandomHorizontalFlip(), #图像一半的概率翻转，一半的概率不翻转
        torchvision.transforms.RandomCrop(32), #图像随机裁剪成32*32
        # torchvision.transforms.RandomVerticalFlip(),
        # torchvision.transforms.RandomRotation(15),
        torchvision.transforms.ToTensor(), #转为Tensor ，归一化
        torchvision.transforms.Normalize([0.5,0.5,0.5], [0.5,0.5,0.5])
        #torchvision.transforms.Normalize((0.4914, 0.4822, 0.4465),(0.2023, 0.1994, 0.2010))
        ])
transform_test=torchvision.transforms.Compose([
        torchvision.transforms.ToTensor(),
        torchvision.transforms.Normalize([0.5,0.5,0.5], [0.5,0.5,0.5])
        #torchvision.transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010))
        ])
# epoch时才对数据集进行以上数据增强操作

num_classes=10
batch_size=128
learning_rate=0.001
num_epoches=100
classes = ("plane","car","bird","cat","deer","dog","frog","horse","ship","truck")

# load downloaded dataset
train_dataset = torchvision.datasets.CIFAR10('./data', download=True, train=True, transform=transform_train)
test_dataset = torchvision.datasets.CIFAR10('./data', download=True, train=False, transform=transform_test)

# Data loader 
train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
test_loader = torch.utils.data.DataLoader(test_dataset, batch_size=batch_size, shuffle=False)

# Define 3*3 convolutional neural network
def conv3x3(in_channels, out_channels, stride=1):
    return nn.Conv2d(in_channels, out_channels, kernel_size=3, stride=stride, padding=1, bias=False)

class ResidualBlock(nn.Module):
    def __init__(self, in_channels, out_channels, stride=1, downsample=None):
        super(ResidualBlock, self).__init__()
        self.conv1 = conv3x3(in_channels, out_channels, stride)
        self.bn1 = nn.BatchNorm2d(out_channels)
        self.relu = nn.ReLU(inplace=True)
        self.conv2 = conv3x3(out_channels, out_channels)
        self.bn2 = nn.BatchNorm2d(out_channels)
        self.downsample = downsample
    def forward(self, x):
        residual=x
        out = self.conv1(x)
        out = self.bn1(out)
        out = self.relu(out)
        out = self.conv2(out)
        out = self.bn2(out)
        if(self.downsample):
            residual = self.downsample(x)
        out += residual
        out = self.relu(out)
        return out
 
# 自定义一个神经网络，使用nn.model，，通过__init__初始化每一层神经网络。
# 使用forward连接数据
class ResNet(nn.Module):
    def __init__(self, block, layers, num_classes):
        super(ResNet, self).__init__()
        self.in_channels = 16
        self.conv = conv3x3(3, 16)
        self.bn = torch.nn.BatchNorm2d(16)
        self.relu = torch.nn.ReLU(inplace=True)
        self.layer1 = self._make_layers(block, 16, layers[0])
        self.layer2 = self._make_layers(block, 32, layers[1], 2)
        self.layer3 = self._make_layers(block, 64, layers[2], 2)
        self.layer4 = self._make_layers(block, 128, layers[3], 2)
        self.avg_pool = torch.nn.AdaptiveAvgPool2d((1, 1))
        self.fc = torch.nn.Linear(128, num_classes)
        
    def _make_layers(self, block, out_channels, blocks, stride=1):
        downsample = None
        if (stride != 1) or (self.in_channels != out_channels):
            downsample = torch.nn.Sequential(
                conv3x3(self.in_channels, out_channels, stride=stride),
                torch.nn.BatchNorm2d(out_channels))
        layers = []
        layers.append(block(self.in_channels, out_channels, stride, downsample))
        self.in_channels = out_channels
        for i in range(1, blocks):
            layers.append(block(out_channels, out_channels))
        return torch.nn.Sequential(*layers)
    
    def forward(self, x):
        out = self.conv(x)
        out = self.bn(out)
        out = self.relu(out)
        out = self.layer1(out)
        out = self.layer2(out)
        out = self.layer3(out)
        out = self.layer4(out)
        out = self.avg_pool(out)
        out = out.view(out.size(0), -1)
        out = self.fc(out)
        return out

# Make model，使用cpu
model=ResNet(ResidualBlock, [2,2,2,2], num_classes).to(device=device)

# 打印model结构
# print(f"Model structure: {model}\n\n")

# 优化器和损失函数
criterion = nn.CrossEntropyLoss()  #交叉熵损失函数
optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate) #优化器随机梯度下降

if __name__ == "__main__":
    # Train the model
    total_step = len(train_loader)
    for epoch in range(0,num_epoches):
        for i, (images, labels) in enumerate(train_loader):
            images = images.to(device=device)
            labels = labels.to(device=device)
            # Forward pass
            outputs = model(images)
            loss = criterion(outputs, labels)
            # Backward and optimize
            optimizer.zero_grad()
            # 反向传播
            loss.backward()
            # 更新参数
            optimizer.step()
            #sum_loss += loss.item()
            #_, predicted = torch.max(outputs.data, dim=1)
            #total += labels.size(0)
            #correct += predicted.eq(labels.data).cpu().sum()
            if (i+1) % total_step == 0:
                print('Epoch [{}/{}], Step [{}/{}], Loss: {:.4f}'
                      .format(epoch+1, num_epoches, i+1, total_step, loss.item()))
    print("Finished Tranining")
    # 保存权重文件
    #torch.save(model.state_dict(), 'model_weights.pth')
    #torch.save(model, 'model.pt')

    print('\nTest the model')
    model.eval()
    with torch.no_grad():
        correct = 0
        total = 0
        for images, labels in test_loader:
            images = images.to(device=device)
            labels = labels.to(device=device)
            outputs = model(images)
            _, predicted = torch.max(outputs.data, 1)
            total += labels.size(0)
            correct += (predicted == labels).sum().item()
        print('在10000张测试集图片上的准确率:{:.4f} %'.format(100 * correct / total))

    # 导出onnx模型
    x = torch.randn((1, 3, 32, 32))
    torch.onnx.export(model, x, './resnet18.onnx', opset_version=12, input_names=['input'], output_names=['output'])

这里有个要注意的，数据集已经提前下载好了，所以没有在线下载

数据集下载是通过下面代码，数据集放在data目录下：

# load downloaded dataset
train_dataset = torchvision.datasets.CIFAR10('./data', download=True, train=True, transform=transform_train)
test_dataset = torchvision.datasets.CIFAR10('./data', download=True, train=False, transform=transform_test)

类型分类为10类

classes = ("plane","car","bird","cat","deer","dog","frog","horse","ship","truck")

等待大概1小时，训练结束后会在当前目录下生成resnet18.onnx模型

四、测试onnx模型

测试代码如下：

test_resnet18_onnx.py

import os, sys

sys.path.append(os.getcwd())
import onnxruntime
import torchvision.models as models
import torchvision.transforms as transforms
from PIL import Image


def to_numpy(tensor):
    return tensor.detach().cpu().numpy() if tensor.requires_grad else tensor.cpu().numpy()

# 自定义的数据增强
def get_test_transform(): 
    return transforms.Compose([
        transforms.Resize([32, 32]),
        transforms.ToTensor(),
        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
    ])

# 推理的图片路径
image = Image.open('./horse5.jpg').convert('RGB')

img = get_test_transform()(image)
img = img.unsqueeze_(0)  # -> NCHW, 1,3,224,224
# 模型加载
onnx_model_path = "resnet18.onnx"
resnet_session = onnxruntime.InferenceSession(onnx_model_path)
inputs = {resnet_session.get_inputs()[0].name: to_numpy(img)}
outs = resnet_session.run(None, inputs)[0]

print("onnx weights", outs)
print("onnx prediction", outs.argmax(axis=1)[0])

测试预测结果7，对应的是马。

五、RKNN模型转换

打开ATK搭建好的虚拟机，进入环境rknn2_env，RKNN环境要确保安装好。

转换代码在rknn-toolkit2目录下的example的pytorch里也有，参考代码如下：

rknn_transfer.py

import numpy as np
import cv2
from rknn.api import RKNN
import torchvision.models as models
import torch
import os

def softmax(x):
    return np.exp(x)/sum(np.exp(x))

def torch_version():
    import torch
    torch_ver = torch.__version__.split('.')
    torch_ver[2] = torch_ver[2].split('+')[0]
    return [int(v) for v in torch_ver]

if __name__ == '__main__':

    if torch_version() < [1, 9, 0]:
        import torch
        print("Your torch version is '{}', in order to better support the Quantization Aware Training (QAT) model,\n"
              "Please update the torch version to '1.9.0' or higher!".format(torch.__version__))
        exit(0)

    MODEL = './resnet18.onnx'

    # Create RKNN object
    rknn = RKNN(verbose=True)

    # Pre-process config
    print('--> Config model')
    rknn.config(mean_values=[127.5, 127.5, 127.5], std_values=[255, 255, 255], target_platform='rk3568')
    #rknn.config(mean_values=[123.675, 116.28, 103.53], std_values=[58.395, 58.395, 58.395], target_platform='rk3568')
    #rknn.config(mean_values=[125.307, 122.961, 113.8575], std_values=[51.5865, 50.847, 51.255], target_platform='rk3568')
    print('done')

    # Load model
    print('--> Loading model')
    #ret = rknn.load_pytorch(model=model, input_size_list=input_size_list)
    ret = rknn.load_onnx(model=MODEL)
    if ret != 0:
        print('Load model failed!')
        exit(ret)
    print('done')

    # Build model
    print('--> Building model')
    ret = rknn.build(do_quantization=False)
    if ret != 0:
        print('Build model failed!')
        exit(ret)
    print('done')

    # Export rknn model
    print('--> Export rknn model')
    ret = rknn.export_rknn('./resnet_18_100.rknn')
    if ret != 0:
        print('Export rknn model failed!')
        exit(ret)
    print('done')

    # Set inputs
    img = cv2.imread('./0_125.jpg')
    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
    img = cv2.resize(img,(32,32))
    #img = np.expand_dims(img, 0)
    
    # Init runtime environment
    print('--> Init runtime environment')
    ret = rknn.init_runtime()
    if ret != 0:
        print('Init runtime environment failed!')
        exit(ret)
    print('done')

    # Inference
    print('--> Running model')
    outputs = rknn.inference(inputs=[img])
    np.save('./pytorch_resnet18_qat_0.npy', outputs[0])
    #show_outputs(softmax(np.array(outputs[0][0])))
    print(outputs)
    print('done')

    rknn.release()

运行python rknn_transfer.py，正常生成rknn文件。

会在当前目录下生成rknn模型

六、部署

导出rknn后，把rknn和测试图片通过adb上传到开发板。

rknnlite_inference0.py

import numpy as np
import cv2
import os
from rknnlite.api import RKNNLite

IMG_PATH = '2_67.jpg'
RKNN_MODEL = './resnet18.rknn'
img_height = 32
img_width = 32
class_names = ["plane","car","bird","cat","deer","dog","frog","horse","ship","truck"]

# Create RKNN object
rknn_lite = RKNNLite()

# load RKNN model
print('--> Load RKNN model')
ret = rknn_lite.load_rknn(RKNN_MODEL)
if ret != 0:
    print('Load RKNN model failed')
    exit(ret)
print('done')

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

# load image
img = cv2.imread(IMG_PATH)
img = cv2.resize(img,(32,32))
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
img = np.expand_dims(img, 0)


# runing model
print('--> Running model')
outputs = rknn_lite.inference(inputs=[img])
print("result: ", outputs)
print(
    "This image most likely belongs to {}."
    .format(class_names[np.argmax(outputs)])
)

rknn_lite.release()

在开发板上运行结果

识别出来horse,验证正常。

如有侵权，或需要完整代码，请及时联系博主。