大家好，我是紅色石頭！
在上三篇文章：
我用 PyTorch 復現了 LeNet-5 神經網路（MNIST 手寫資料集篇）！
我用 PyTorch 復現了 LeNet-5 神經網路（CIFAR10 資料集篇）！
詳細介紹了卷積神經網路 LeNet-5 的理論部分和使用 PyTorch 復現 LeNet-5 網路來解決 MNIST 資料集和 CIFAR10 資料集。然而大多數實際應用中，我們需要自己構建資料集，進行識別。因此，本文將講解一下如何使用 LeNet-5 訓練自己的資料。
正文開始！
三、用 LeNet-5 訓練自己的資料
下面使用 LeNet-5 網路來訓練本地的資料並進行測試。資料集是本地的 LED 數字 0-9，尺寸為 28×28 單通道，跟 MNIST 資料集類似。訓練集 0-9 各 95 張，測試集 0~9 各 40 張。圖片樣例如圖所示：
3.1 資料預處理
製作圖片資料的索引
對於訓練集和測試集，要分別製作對應的圖片資料索引，即 train.txt 和 test.txt兩個檔案，每個 txt 中包含每個圖片的目錄和對應類別 class。示意圖如下：
製作圖片資料索引的 python 指令碼程式如下：
import os
train_txt_path = os.path.join("data", "LEDNUM", "train.txt")
train_dir = os.path.join("data", "LEDNUM", "train_data")
valid_txt_path = os.path.join("data", "LEDNUM", "test.txt")
valid_dir = os.path.join("data", "LEDNUM", "test_data")
def gen_txt(txt_path, img_dir):
f = open(txt_path, 'w')
for root, s_dirs, _ in os.walk(img_dir, topdown=True): # 獲取 train檔案下各資料夾名稱
for sub_dir in s_dirs:
i_dir = os.path.join(root, sub_dir) # 獲取各類的資料夾 絕對路徑
img_list = os.listdir(i_dir) # 獲取類別資料夾下所有png圖片的路徑
for i in range(len(img_list)):
if not img_list[i].endswith('jpg'): # 若不是png檔案，跳過
continue
label = img_list[i].split('_')[0]
img_path = os.path.join(i_dir, img_list[i])
line = img_path + ' ' + label + '\n'
f.write(line)
f.close()
if __name__ == '__main__':
gen_txt(train_txt_path, train_dir)
gen_txt(valid_txt_path, valid_dir)
執行指令碼之後就在 ./data/LEDNUM/ 目錄下生成 train.txt 和 test.txt 兩個索引檔案。
構建Dataset子類
pytorch 載入自己的資料集，需要寫一個繼承自 torch.utils.data 中 Dataset 類，並修改其中的 init 方法、getitem 方法、len 方法。預設載入的都是圖片，init 的目的是得到一個包含資料和標籤的 list，每個元素能找到圖片位置和其對應標籤。然後用 getitem 方法得到每個元素的影像畫素矩陣和標籤，返回 img 和 label。
from PIL import Image
from torch.utils.data import Dataset
class MyDataset(Dataset):
def __init__(self, txt_path, transform = None, target_transform = None):
fh = open(txt_path, 'r')
imgs = []
for line in fh:
line = line.rstrip()
words = line.split()
imgs.append((words[0], int(words[1])))
self.imgs = imgs
self.transform = transform
self.target_transform = target_transform
def __getitem__(self, index):
fn, label = self.imgs[index]
#img = Image.open(fn).convert('RGB')
img = Image.open(fn)
if self.transform is not None:
img = self.transform(img)
return img, label
def __len__(self):
return len(self.imgs)
getitem 是核心函式。self.imgs 是一個 list，self.imgs[index] 是一個 str，包含圖片路徑，圖片標籤，這些資訊是從上面生成的txt檔案中讀取；利用 Image.open 對圖片進行讀取，注意這裡的 img 是單通道還是三通道的；self.transform(img) 對圖片進行處理，這個 transform 裡邊可以實現減均值、除標準差、隨機裁剪、旋轉、翻轉、放射變換等操作。
當 Mydataset構 建好，剩下的操作就交給 DataLoder，在 DataLoder 中，會觸發 Mydataset 中的 getiterm 函式讀取一張圖片的資料和標籤，並拼接成一個 batch 返回，作為模型真正的輸入。
pipline_train = transforms.Compose([
#隨機旋轉圖片
transforms.RandomHorizontalFlip(),
#將圖片尺寸resize到32x32
transforms.Resize((32,32)),
#將圖片轉化為Tensor格式
transforms.ToTensor(),
#正則化(當模型出現過擬合的情況時，用來降低模型的複雜度)
transforms.Normalize((0.1307,),(0.3081,))
])
pipline_test = transforms.Compose([
#將圖片尺寸resize到32x32
transforms.Resize((32,32)),
transforms.ToTensor(),
transforms.Normalize((0.1307,),(0.3081,))
])
train_data = MyDataset('./data/LEDNUM/train.txt', transform=pipline_train)
test_data = MyDataset('./data/LEDNUM/test.txt', transform=pipline_test)
#train_data 和test_data包含多有的訓練與測試資料，呼叫DataLoader批量載入
trainloader = torch.utils.data.DataLoader(dataset=train_data, batch_size=8, shuffle=True)
testloader = torch.utils.data.DataLoader(dataset=test_data, batch_size=4, shuffle=False)
3.2 搭建 LeNet-5 神經網路結構
class LeNet(nn.Module):
def __init__(self):
super(LeNet, self).__init__()
self.conv1 = nn.Conv2d(1, 6, 5)
self.relu = nn.ReLU()
self.maxpool1 = nn.MaxPool2d(2, 2)
self.conv2 = nn.Conv2d(6, 16, 5)
self.maxpool2 = nn.MaxPool2d(2, 2)
self.fc1 = nn.Linear(16*5*5, 120)
self.fc2 = nn.Linear(120, 84)
self.fc3 = nn.Linear(84, 10)
def forward(self, x):
x = self.conv1(x)
x = self.relu(x)
x = self.maxpool1(x)
x = self.conv2(x)
x = self.maxpool2(x)
x = x.view(-1, 16*5*5)
x = F.relu(self.fc1(x))
x = F.relu(self.fc2(x))
x = self.fc3(x)
output = F.log_softmax(x, dim=1)
return output
3.3 將定義好的網路結構搭載到 GPU/CPU，並定義優化器
#建立模型，部署gpu
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = LeNet().to(device)
#定義優化器
optimizer = optim.Adam(model.parameters(), lr=0.001)
3.4 定義訓練函式
def train_runner(model, device, trainloader, optimizer, epoch):
#訓練模型, 啟用 BatchNormalization 和 Dropout, 將BatchNormalization和Dropout置為True
model.train()
total = 0
correct =0.0
#enumerate迭代已載入的資料集,同時獲取資料和資料下標
for i, data in enumerate(trainloader, 0):
inputs, labels = data
#把模型部署到device上
inputs, labels = inputs.to(device), labels.to(device)
#初始化梯度
optimizer.zero_grad()
#儲存訓練結果
outputs = model(inputs)
#計算損失和
#多分類情況通常使用cross_entropy(交叉熵損失函式), 而對於二分類問題, 通常使用sigmod
loss = F.cross_entropy(outputs, labels)
#獲取最大概率的預測結果
#dim=1表示返回每一行的最大值對應的列下標
predict = outputs.argmax(dim=1)
total += labels.size(0)
correct += (predict == labels).sum().item()
#反向傳播
loss.backward()
#更新引數
optimizer.step()
if i % 100 == 0:
#loss.item()表示當前loss的數值
print("Train Epoch{} \t Loss: {:.6f}, accuracy: {:.6f}%".format(epoch, loss.item(), 100*(correct/total)))
Loss.append(loss.item())
Accuracy.append(correct/total)
return loss.item(), correct/total
3.5 定義測試函式
def test_runner(model, device, testloader):
#模型驗證, 必須要寫, 否則只要有輸入資料, 即使不訓練, 它也會改變權值
#因為呼叫eval()將不啟用 BatchNormalization 和 Dropout, BatchNormalization和Dropout置為False
model.eval()
#統計模型正確率, 設定初始值
correct = 0.0
test_loss = 0.0
total = 0
#torch.no_grad將不會計算梯度, 也不會進行反向傳播
with torch.no_grad():
for data, label in testloader:
data, label = data.to(device), label.to(device)
output = model(data)
test_loss += F.cross_entropy(output, label).item()
predict = output.argmax(dim=1)
#計算正確數量
total += label.size(0)
correct += (predict == label).sum().item()
#計算損失值
print("test_avarage_loss: {:.6f}, accuracy: {:.6f}%".format(test_loss/total, 100*(correct/total)))
3.6 執行
#呼叫
epoch = 5
Loss = []
Accuracy = []
for epoch in range(1, epoch+1):
print("start_time",time.strftime('%Y-%m-%d %H:%M:%S',time.localtime(time.time())))
loss, acc = train_runner(model, device, trainloader, optimizer, epoch)
Loss.append(loss)
Accuracy.append(acc)
test_runner(model, device, testloader)
print("end_time: ",time.strftime('%Y-%m-%d %H:%M:%S',time.localtime(time.time())),'\n')
print('Finished Training')
plt.subplot(2,1,1)
plt.plot(Loss)
plt.title('Loss')
plt.show()
plt.subplot(2,1,2)
plt.plot(Accuracy)
plt.title('Accuracy')
plt.show()
經歷 5 次 epoch 的 loss 和 accuracy 曲線如下：
3.7 模型儲存
torch.save(model, './models/model-mine.pth') #儲存模型
3.8 模型測試
下面使用上面訓練的模型對一張 LED 圖片進行測試。
from PIL import Image
import numpy as np
if __name__ == '__main__':
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model = torch.load('./models/model-mine.pth') #載入模型
model = model.to(device)
model.eval() #把模型轉為test模式
#讀取要預測的圖片
# 讀取要預測的圖片
img = Image.open("./images/test_led.jpg") # 讀取影像
#img.show()
plt.imshow(img,cmap="gray") # 顯示圖片
plt.axis('off') # 不顯示座標軸
plt.show()
# 匯入圖片，圖片擴充套件後為[1，1，32，32]
trans = transforms.Compose(
[
#將圖片尺寸resize到32x32
transforms.Resize((32,32)),
transforms.ToTensor(),
transforms.Normalize((0.1307,), (0.3081,))
])
img = trans(img)
img = img.to(device)
img = img.unsqueeze(0) #圖片擴充套件多一維,因為輸入到儲存的模型中是4維的[batch_size,通道,長，寬]，而普通圖片只有三維，[通道,長，寬]
# 預測
output = model(img)
prob = F.softmax(output,dim=1) #prob是10個分類的概率
print("概率：",prob)
value, predicted = torch.max(output.data, 1)
predict = output.argmax(dim=1)
print("預測類別：",predict.item())
概率：tensor([[7.2506e-11, 7.0065e-18, 7.1749e-06, 7.4855e-13, 7.3532e-08, 8.5405e-17,
2.5753e-15, 9.7887e-10, 2.7855e-05, 9.9996e-01]],
grad_fn=)
預測類別：9
模型預測結果正確！
以上就是 PyTorch 構建 LeNet-5 卷積神經網路並用它來識別自定義資料集的例子。全文的程式碼都是可以順利執行的，建議大家自己跑一邊。
總結：
是我們目前分別復現了 LeNet-5 來識別 MNIST、CIFAR10 和自定義資料集，基本上涵蓋了基於 PyToch 的 LeNet-5 實戰的所有內容。希望對大家有所幫助！
所有完整的程式碼我都放在 GitHub 上，GitHub地址為：
https://github.com/RedstoneWill/ObjectDetectionLearner/tree/main/LeNet-5
