這幾天又在玩樹莓派，先是搞了個物聯網，又在嘗試在樹莓派上搞一些簡單的神經網絡，這次搞得是mlp識別mnist手寫數字識別
訓練代碼在電腦上，cpu就能訓練，很快的：
1 import torch 2 import torch.nn as nn 3 import torch.optim as optim 4 from torchvision import datasets, transforms 5 6 # 設置隨機種子 7 torch.manual_seed(42) 8 9 # 定義MLP模型 10 class MLP(nn.Module): 11 def __init__(self): 12 super(MLP, self).__init__() 13 self.fc1 = nn.Linear(784, 256) 14 self.fc2 = nn.Linear(256, 128) 15 self.fc3 = nn.Linear(128, 10) 16 17 def forward(self, x): 18 x = x.view(-1, 784) 19 x = torch.relu(self.fc1(x)) 20 x = torch.relu(self.fc2(x)) 21 x = self.fc3(x) 22 return x 23 24 # 加載MNIST數據集 25 transform = transforms.Compose([ 26 transforms.ToTensor(), 27 # transforms.Normalize((0.1307,), (0.3081,)) 28 ]) 29 30 train_dataset = datasets.MNIST(root='./data', train=True, download=True, transform=transform) 31 test_dataset = datasets.MNIST(root='./data', train=False, download=True, transform=transform) 32 33 train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=64, shuffle=True) 34 test_loader = torch.utils.data.DataLoader(test_dataset, batch_size=64, shuffle=False) 35 36 # 創建模型實例 37 model = MLP() 38 39 # 定義損失函數和優化器 40 criterion = nn.CrossEntropyLoss() 41 optimizer = optim.SGD(model.parameters(), lr=0.01, momentum=0.5) 42 43 # 訓練模型 44 def train(model, train_loader, optimizer, criterion, epochs): 45 model.train() 46 for epoch in range(1, epochs + 1): 47 for batch_idx, (data, target) in enumerate(train_loader): 48 optimizer.zero_grad() 49 output = model(data) 50 loss = criterion(output, target) 51 loss.backward() 52 optimizer.step() 53 54 if batch_idx % 100 == 0: 55 print('Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format( 56 epoch, batch_idx * len(data), len(train_loader.dataset), 57 100. * batch_idx / len(train_loader), loss.item())) 58 59 # 訓練模型 60 train(model, train_loader, optimizer, criterion, epochs=5) 61 62 # 保存模型爲NumPy格式 63 numpy_model = {} 64 numpy_model['fc1.weight'] = model.fc1.weight.detach().numpy() 65 numpy_model['fc1.bias'] = model.fc1.bias.detach().numpy() 66 numpy_model['fc2.weight'] = model.fc2.weight.detach().numpy() 67 numpy_model['fc2.bias'] = model.fc2.bias.detach().numpy() 68 numpy_model['fc3.weight'] = model.fc3.weight.detach().numpy() 69 numpy_model['fc3.bias'] = model.fc3.bias.detach().numpy() 70 71 # 保存爲NumPy格式的數據 72 import numpy as np 73 np.savez('mnist_model.npz', **numpy_model)
然後需要自己倒出一些圖片在dataset裏：我保存在了mnist_pi文件夾下,“_”後面的是標籤，主要是在pc端導出保存到樹莓派下
樹莓派推理端的代碼，需要numpy手動重新搭建網絡，然後加載那些保存的矩陣參數，做矩陣乘法和加法
1 import numpy as np 2 import os 3 from PIL import Image 4 5 # 加載模型 6 model_data = np.load('mnist_model.npz') 7 weights1 = model_data['fc1.weight'] 8 biases1 = model_data['fc1.bias'] 9 weights2 = model_data['fc2.weight'] 10 biases2 = model_data['fc2.bias'] 11 weights3 = model_data['fc3.weight'] 12 biases3 = model_data['fc3.bias'] 13 14 # 進行推理 15 def predict(image, weights1, biases1,weights2, biases2,weights3, biases3): 16 image = image.flatten()/255 # 將輸入圖像展平並進行歸一化 17 output = np.dot(weights1, image) + biases1 18 output = np.dot(weights2, output) + biases2 19 output = np.dot(weights3, output) + biases3 20 predicted_class = np.argmax(output) 21 return predicted_class 22 23 24 25 26 folder_path = './mnist_pi' # 替換爲圖片所在的文件夾路徑 27 def infer_images_in_folder(folder_path): 28 for file_name in os.listdir(folder_path): 29 file_path = os.path.join(folder_path, file_name) 30 if os.path.isfile(file_path) and file_name.endswith(('.jpg', '.jpeg', '.png')): 31 image = Image.open(file_path) 32 label = file_name.split(".")[0].split("_")[1] 33 image = np.array(image) 34 print("file_path:",file_path,"img size:",image.shape,"label:",label) 35 predicted_class = predict(image, weights1, biases1,weights2, biases2,weights3, biases3) 36 print('Predicted class:', predicted_class) 37 38 infer_images_in_folder(folder_path)
結果:
效果還不錯：
這次內容就到這裏了，下次爭取做一個卷積的神經網絡在樹莓派上推理，然後爭取做一個目標檢測的模型在樹莓派上