这几天又在玩树莓派，先是搞了个物联网，又在尝试在树莓派上搞一些简单的神经网络，这次搞得是LSTM识别mnist手写数字识别
训练代码在电脑上，cpu就能训练，很快的：
import torch import torch.nn as nn import torchvision import numpy as np import os from PIL import Image # 定义LSTM模型 class LSTMModel(nn.Module): def __init__(self, input_size, hidden_size, num_layers, num_classes): super(LSTMModel, self).__init__() self.hidden_size = hidden_size self.num_layers = num_layers self.lstm = nn.LSTM(input_size, hidden_size, num_layers, batch_first=True) self.fc = nn.Linear(hidden_size, num_classes) def forward(self, x): h0 = torch.zeros(self.num_layers, x.size(0), self.hidden_size).to(x.device) c0 = torch.zeros(self.num_layers, x.size(0), self.hidden_size).to(x.device) out, (h_n,c_n) = self.lstm(x, (h0, c0)) out = self.fc(out[:, -1, :]) return out # 设置超参数 input_size = 28 sequence_length = 28 hidden_size = 128 num_layers = 2 num_classes = 10 batch_size = 100 num_epochs = 1 learning_rate = 0.001 # 加载MNIST数据集 train_dataset = torchvision.datasets.MNIST(root='./data', train=True, transform=torchvision.transforms.ToTensor(), download=True) test_dataset = torchvision.datasets.MNIST(root='./data', train=False, transform=torchvision.transforms.ToTensor()) # 创建数据加载器 train_loader = torch.utils.data.DataLoader(dataset=train_dataset, batch_size=batch_size, shuffle=True) test_loader = torch.utils.data.DataLoader(dataset=test_dataset, batch_size=batch_size, shuffle=False) # 创建LSTM模型 model = LSTMModel(input_size, hidden_size, num_layers, num_classes) # 定义损失函数和优化器 criterion = nn.CrossEntropyLoss() optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate) # 训练模型 total_step = len(train_loader) for epoch in range(num_epochs): for i, (images, labels) in enumerate(train_loader): images = images.reshape(-1, sequence_length, input_size) outputs = model(images) loss = criterion(outputs, labels) predictions = torch.argmax(outputs,dim=1) # acc = torch.eq(predictions,labels).sum().item() # print(acc) optimizer.zero_grad() loss.backward() optimizer.step() if (i + 1) % 100 == 0: print('Epoch [{}/{}], Step [{}/{}], Loss: {:.4f}'.format(epoch + 1, num_epochs, i + 1, total_step, loss.item())) # 保存模型 torch.save(model.state_dict(), 'model.pth') # 加载模型 model.load_state_dict(torch.load('model.pth')) with torch.no_grad(): for i, (images, labels) in enumerate(test_loader): images = images.reshape(-1, sequence_length, input_size) outputs = model(images) predictions = torch.argmax(outputs,dim=1) acc = torch.eq(predictions,labels).sum().item() print(acc) # folder_path = './mnist_pi' # 替换为图片所在的文件夹路径 # def infer_images_in_folder(folder_path): # with torch.no_grad(): # for file_name in os.listdir(folder_path): # file_path = os.path.join(folder_path, file_name) # if os.path.isfile(file_path) and file_name.endswith(('.jpg', '.jpeg', '.png')): # image = Image.open(file_path) # label = file_name.split(".")[0].split("_")[1] # image = np.array(image)/255.0 # image = np.expand_dims(image,axis=0) # image= torch.tensor(image).to(torch.float32) # logits = model(image) # predicted_class = torch.argmax(logits) # print("file_path:",file_path,"img size:",image.shape,"label:",label,'Predicted class:', predicted_class) # break # infer_images_in_folder(folder_path) # 保存模型参数为numpy的数组格式 model_params = {} # print(list(model.parameters())) for name, param in model.named_parameters(): model_params[name] = param.detach().numpy() print(name,param.shape) np.savez('model.npz', **model_params)
然后需要自己在dataset里导出一些图片：我保存在了mnist_pi文件夹下,“_”后面的是标签，主要是在pc端导出保存到树莓派下
树莓派推理端的代码，需要numpy手动重新搭建网络，并且需要手动实现双层的LSTM神经网络，然后加载那些保存的矩阵参数，做矩阵乘法和加法
import numpy as np import os from PIL import Image # 加载模型参数 model_data = np.load('model.npz') ''' weight_ih_l[k] : the learnable input-hidden weights of the :math:`\text{k}^{th}` layer `(W_ii|W_if|W_ig|W_io)`, of shape `(4*hidden_size, input_size)` for `k = 0`. Otherwise, the shape is `(4*hidden_size, num_directions * hidden_size)`. If ``proj_size > 0`` was specified, the shape will be `(4*hidden_size, num_directions * proj_size)` for `k > 0` weight_hh_l[k] : the learnable hidden-hidden weights of the :math:`\text{k}^{th}` layer `(W_hi|W_hf|W_hg|W_ho)`, of shape `(4*hidden_size, hidden_size)`. If ``proj_size > 0`` was specified, the shape will be `(4*hidden_size, proj_size)`. bias_ih_l[k] : the learnable input-hidden bias of the :math:`\text{k}^{th}` layer `(b_ii|b_if|b_ig|b_io)`, of shape `(4*hidden_size)` bias_hh_l[k] : the learnable hidden-hidden bias of the :math:`\text{k}^{th}` layer `(b_hi|b_hf|b_hg|b_ho)`, of shape `(4*hidden_size)` ''' # 提取模型参数 lstm_weight_ih_l0 = model_data['lstm.weight_ih_l0'] lstm_weight_hh_l0 = model_data['lstm.weight_hh_l0'] lstm_bias_ih_l0 = model_data['lstm.bias_ih_l0'] lstm_bias_hh_l0 = model_data['lstm.bias_hh_l0'] lstm_weight_ih_l1 = model_data['lstm.weight_ih_l1'] lstm_weight_hh_l1 = model_data['lstm.weight_hh_l1'] lstm_bias_ih_l1 = model_data['lstm.bias_ih_l1'] lstm_bias_hh_l1 = model_data['lstm.bias_hh_l1'] fc_weight = model_data['fc.weight'] fc_bias = model_data['fc.bias'] # print(lstm_weight_ih_l0.shape,lstm_weight_hh_l0.shape) # print(lstm_bias_ih_l0.shape,lstm_bias_hh_l0.shape) # 定义LSTM模型 def lstm_model(inputs): ''' 踩到两个坑，一个是矩阵形状都是这种(4*hidden_size, hidden_size)合并的，需要拆分。 另一个坑是，两层的lstm层需要每个时间步的输出都输入到下一层，而不是最后一个时间步的数据给下一层 ''' batch_size, sequence_length, input_size = inputs.shape hidden_size = lstm_weight_hh_l0.shape[1] num_classes = fc_weight.shape[0] h0 = np.zeros((batch_size, hidden_size)) c0 = np.zeros((batch_size, hidden_size)) # 第一层LSTM h_l0, c_l0 = np.zeros_like(h0), np.zeros_like(c0) out_0 = [] for t in range(sequence_length): x = inputs[:, t, :] ''' i_t = \sigma(W_{ii} x_t + b_{ii} + W_{hi} h_{t-1} + b_{hi}) \\ f_t = \sigma(W_{if} x_t + b_{if} + W_{hf} h_{t-1} + b_{hf}) \\ g_t = \tanh(W_{ig} x_t + b_{ig} + W_{hg} h_{t-1} + b_{hg}) \\ o_t = \sigma(W_{io} x_t + b_{io} + W_{ho} h_{t-1} + b_{ho}) \\ c_t = f_t \odot c_{t-1} + i_t \odot g_t \\ h_t = o_t \odot \tanh(c_t) \\ ''' # 输入门 i_t = sigmoid(np.dot(x, lstm_weight_ih_l0[:128].T) + np.dot(h_l0, lstm_weight_hh_l0[:128].T) + lstm_bias_ih_l0[:128] + lstm_bias_hh_l0[:128]) # 遗忘门 f_t = sigmoid(np.dot(x, lstm_weight_ih_l0[128:256].T) + np.dot(h_l0, lstm_weight_hh_l0[128:256].T) + lstm_bias_ih_l0[128:256] + lstm_bias_hh_l0[128:256]) # 候选向量 g_t = np.tanh(np.dot(x, lstm_weight_ih_l0[256:256+128].T) + np.dot(h_l0, lstm_weight_hh_l0[256:256+128].T) + lstm_bias_ih_l0[256:256+128] + lstm_bias_hh_l0[256:256+128]) # 输出门 o_t = sigmoid(np.dot(x, lstm_weight_ih_l0[256+128:512].T) + np.dot(h_l0, lstm_weight_hh_l0[256+128:512].T) + lstm_bias_ih_l0[256+128:512] + lstm_bias_hh_l0[256+128:512]) # 细胞状态 c_l0 = f_t * c_l0 + i_t * g_t # 隐藏状态 h_l0 = o_t * np.tanh(c_l0) out_0.append(h_l0) # 第二层LSTM h_l1, c_l1 = np.zeros_like(h0), np.zeros_like(c0) out_1 = [] for t in range(sequence_length): x = out_0[t] # 输入门 i_t = sigmoid(np.dot(x, lstm_weight_ih_l1[:128].T) + np.dot(h_l1, lstm_weight_hh_l1[:128].T) + lstm_bias_ih_l1[:128] + lstm_bias_hh_l1[:128]) # 遗忘门 f_t = sigmoid(np.dot(x, lstm_weight_ih_l1[128:256].T) + np.dot(h_l1, lstm_weight_hh_l1[128:256].T) + lstm_bias_ih_l1[128:256] + lstm_bias_hh_l1[128:256]) # 候选向量 g_t = np.tanh(np.dot(x, lstm_weight_ih_l1[256:256+128].T) + np.dot(h_l1, lstm_weight_hh_l1[256:256+128].T) + lstm_bias_ih_l1[256:256+128] + lstm_bias_hh_l1[256:256+128]) # 输出门 o_t = sigmoid(np.dot(x, lstm_weight_ih_l1[256+128:512].T) + np.dot(h_l1, lstm_weight_hh_l1[256+128:512].T) + lstm_bias_ih_l1[256+128:512] + lstm_bias_hh_l1[256+128:512]) # 细胞状态 c_l1 = f_t * c_l1 + i_t * g_t # 隐藏状态 h_l1 = o_t * np.tanh(c_l1) out_1.append(h_l1) # 全连接层 fc_output = np.dot(h_l1, fc_weight.T) + fc_bias predictions = np.argmax(fc_output, axis=1) return predictions # Sigmoid激活函数 def sigmoid(x): return 1 / (1 + np.exp(-x)) folder_path = './mnist_pi' # 替换为图片所在的文件夹路径 def infer_images_in_folder(folder_path): for file_name in os.listdir(folder_path): file_path = os.path.join(folder_path, file_name) if os.path.isfile(file_path) and file_name.endswith(('.jpg', '.jpeg', '.png')): image = Image.open(file_path) label = file_name.split(".")[0].split("_")[1] image = np.array(image)/255.0 image = np.expand_dims(image,axis=0) predicted_class = lstm_model(image) print("file_path:",file_path,"img size:",image.shape,"label:",label,'Predicted class:', predicted_class) infer_images_in_folder(folder_path)
这代码完全就是numpy推理，不需要安装pytorch，树莓派也装不动pytorch，太重了，下面是推理结果，比之前的MLP网络慢很多，主要是手动实现的LSTM网络全靠循环实现。
