基于CNN（LeNet）的垃圾分类（C语言实现）

---

一、先使用python训练模型

具体步骤参考基于pytorch的MNIST数据集的四层CNN，测试准确率99.77%这篇文章，各种步骤我写的很详细，只需要将MNIST数据集换成垃圾分类的数据集，再调整一下参数就好了。

二、提取参数

提取模型参数

权重和偏置
我们需要提取每一个具有学习的参数的训练层的权重和偏置，我使用了两层卷积和两层全连接，就要提取两个卷积层的权重和偏置，两个全连接层的权重和偏置。

# Extraction_Parameter.py
#引入库
#引用需要用到的库
import torch
import numpy as np
import torch.nn as nn
import torch.nn.functional as F
import numpy as np

#model
class CNNModel(nn.Module):
    def __init__(self):
        super(CNNModel, self).__init__()
        
        # Convolution layer 1
        self.conv1 = nn.Conv2d(in_channels = 3 , out_channels = 8, kernel_size = 3, stride = 1, padding = 0 )
        self.relu1 = nn.ReLU()
        self.maxpool1 = nn.MaxPool2d(kernel_size = 2, stride = 2)
        
        # Convolution layer 2
        self.conv2 = nn.Conv2d(in_channels =8 , out_channels = 16, kernel_size = 3, stride = 1, padding = 0 )
        self.relu2 = nn.ReLU()
        self.maxpool2 = nn.MaxPool2d(kernel_size = 2, stride = 2)
        
        # Fully-Connected layer 1
        self.fc1 = nn.Linear(400,40)
        
        # Fully-Connected layer 2
        self.fc2 = nn.Linear(40,4)
        
    
    def forward(self, x):
        # conv layer 1 的前向计算，3行代码
        out = self.conv1(x)
        #print(out.shape)
        out = self.relu1(out)
        out = self.maxpool1(out)
        #print(out.shape)
        
        # conv layer 2 的前向计算，3行代码
        out = self.conv2(out)
        #print(out.shape)
        out = self.relu2(out)
        out = self.maxpool2(out)
        #print(out.shape)
        
        #Flatten拉平操作
        out = out.view(out.size(0),-1)
        #print(out.shape)
        #FC layer的前向计算（2行代码）
        out = self.fc1(out)
        out = self.fc2(out)
  
        return F.log_softmax(out,dim = 1)


#实例化模型
network = CNNModel()

#加载模型
model_path = "model1.pth"
network.load_state_dict(torch.load(model_path, map_location = torch.device('cpu')))
#network.eval()

parm = {
   }
for name,parameters in network.state_dict().items():

    parm[name] = parameters.detach().numpy()
    print(name, parameters)

w1 = parm['conv1.weight']
b1 = parm['conv1.bias']
w2 = parm['conv2.weight']
b2 = parm['conv2.bias']

fc1_w = parm['fc1.weight']
fc1_b = parm['fc1.bias']
fc2_w = parm['fc2.weight']
fc2_b = parm['fc2.bias']

#print(type(w1))
#print(len(w1[0]))
#print(len(w1[0][0]))
#print(len(w1[0][0][0]))

#conv1_wb
with open("parameters1_wb.h","a") as f:
    print(type(w1))
    #new_str1 = str(np.transpose(w1).tolist())
    new_str1 = str(w1.tolist())
    new_str2 = new_str1.replace('[','')
    new_str3 = new_str2.replace(']','')
    f.write("float conv1_weight[8][3][9] = {" + new_str3 + "};\n\n")
    print("第一层卷积的权重保存成功")
    f.close()

with open("parameters1_wb.h","a") as f:
    print(type(b1))
    #new_str1 = str(np.transpose(b1).tolist())
    new_str1 = str(b1.tolist())
    new_str2 = new_str1.replace('[','')
    new_str3 = new_str2.replace(']','')
    f.write("float conv1_bias[8] = {" + new_str3 + "};\n\n")
    print("第一层卷积的偏置保存成功")
    f.close()

#conv2_wb
with open("parameters1_wb.h","a") as f:
    print(type(w2))
    #new_str1 = str(np.transpose(w2).tolist())
    new_str1 = str(w2.tolist())
    new_str2 = new_str1.replace('[','')
    new_str3 = new_str2.replace(']','')
    f.write("float conv2_weight[16][8][9] = {" + new_str3 + "};\n\n")
    print("第二层卷积的权重保存成功")
    f.close()

with open("parameters1_wb.h","a") as f:
    print(type(b2))
    #new_str1 = str(np.transpose(b2).tolist())
    new_str1 = str(b2.tolist())
    new_str2 = new_str1.replace('[','')
    new_str3 = new_str2.replace(']','')
    f.write("float conv2_bias[16] = {" + new_str3 + "};\n\n")
    print("第二层卷积的偏置保存成功")
    f.close()


#fc1_wb
with open("parameters1_wb.h","a") as f:
    print(type(fc1_w))
    new_str1 = str(np.transpose(fc1_w).tolist())
    #new_str1 = str(fc1_w.tolist())
    new_str2 = new_str1.replace('[','')
    new_str3 = new_str2.replace(']','')
    f.write("float fc1_weight[" + str(400*40) + "] = {" + new_str3 + "};\n\n")
    print("第一层全连接的权重保存成功")
    f.close()

with open("parameters1_wb.h","a") as f:
    print(type(fc1_b))
    #new_str1 = str(np.transpose(fc1_b).tolist())
    new_str1 = str(fc1_b.tolist())
    new_str2 = new_str1.replace('[','')
    new_str3 = new_str2.replace(']','')
    f.write("float fc1_bias[40] = {" + new_str3 + "};\n\n")
    print("第一层全连接的偏置保存成功")
    f.close()


#fc2_wb
with open("parameters1_wb.h","a") as f:
    print(type(fc2_w))
    new_str1 = str(np.transpose(fc2_w).tolist())
    #new_str1 = str(fc2_w.tolist())
    new_str2 = new_str1.replace('[','')
    new_str3 = new_str2.replace(']','')
    f.write("float fc2_weight[" + str(40*4) + "] = {" + new_str3 + "};\n\n")
    print("第二层全连接的权重保存成功")
    f.close()

with open("parameters1_wb.h","a") as f:
    print(type(fc2_b))
    #new_str1 = str(np.transpose(fc2_b).tolist())
    new_str1 = str(fc2_b.tolist())
    new_str2 = new_str1.replace('[','')
    new_str3 = new_str2.replace(']','')
    f.write("float fc2_bias[4] = {" + new_str3 + "};\n\n")
    print("第二层全连接的偏置保存成功")
    f.close()

提取成功后会得到一个parameters1_wb.h文件，如图所示

提取图片

将测试的图片同样提取为.h文件

# Extract_Image.py
from torchvision import transforms
import torch
import numpy as np
from PIL import Image
from itertools import chain

# 在训练模型时对图片进行怎样的预处理
# 在提取图片参数时需要先进行同样的处理再提取，不然维度数据对不上
data_transform = transforms.Compose(
    [transforms.ToTensor()
     #transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
     ])

img = Image.open("./test/Others1.jpg")#预测图片
img = data_transform(img)
img = torch.unsqueeze(img, dim=0)
img = img.numpy()
#img = img * (1.0/255.0)
#img = img.tolist()

print(img)

#参数提取
with open("pic1.h","a") as f:
    #new_str1 = str(np.transpose(img).tolist())
    new_str1 = str(img.tolist())
    new_str2 = new_str1.replace('[','')
    new_str3 = new_str2.replace(']','')
    f.write("float Others2" + "[3][28][28] = {" + new_str3 + "};\n\n")
    print("图片Others1参数读取完成")
    f.close()

三、编写CNN算子

在windows中实现

#include <stdio.h>
#