這一篇主要講解Pytorch搭建一個卷積神經網絡識別自己的數字數據集基本流程。
注:一開始接觸很多教程都是直接加載datasets已有的MNIST等,如果想要訓練自己的數據就可以采用這個方法。
基本步驟:獲取並讀取數據-->定義網絡模型和損失函數-->使用優化算法訓練模型-->利用驗證數據集求取網絡識別准確度
1、首先是獲取並讀取數據,其中最關鍵的就是Datasets這個類即torchvision.datasets,類里自帶很多數據集(包括mnist/coco/cifar10等)。就拿mnist手寫識別數據集來說,我們所需要提取的信息主要有兩個:一個是手寫圖片的所有像素,另一個就是這張圖片的標簽(即這個數字的大小是多少)。torchvision.datasets.MNIST()該函數所返回的就是這些信息。
datasets這個類是圖像數據集中非常重要的一個類,當我們需要定義自己的數據集進行加載訓練的時候,應該要繼承datasets這個父類,其中父類中的兩個私有成員函數必須被重載:
def __getitem__(self, index) def __len__(self)
len返回的是這個數據集的大小,而getitem主要用來編寫支持數據集索引的函數。
getitem()會接收一個index,然后返回圖片數據和標簽。這個index通常是指一個list的index,這個list的每個元素就包含了圖像所在的系統路徑和標簽信息。
2、那么,如果我們需要定義一個自己的數據集,就需要有一個這樣的list列表。方法就是將圖片的系統路徑和標簽信息放在一個txt文件當中,然后再從這個txt中讀取信息。所以,定義並讀取自己數據的基本流程如下:
(1)制作存儲了圖片系統路徑和標簽的txt文件
(2)將這些信息轉換成一個list,這個list的每一個元素對應一個樣本
(3)通過getitem函數讀取數據像素信息和標簽
3、比如我本來有這樣的10個數字信息,每一個數字取700張作為訓練集,200張作為驗證集,如圖所示


4、接下來需要編寫一個python腳本對這些圖像進行信息提取(提取圖片路徑和標簽)
import os
b = 0
dir = 'F:/ele/data/'
#os.listdir的結果就是一個list集合
#可以使用一個list的sort方法進行排序,有數字就用數字排序
files = os.listdir(dir)
files.sort()
#print("files:", files)
train = open('F:/ele/data/train.txt', 'a')
test = open('F:/ele/data/test.txt', 'a')
a = 0
a1 = 0
while(b < 20):#20是因為10個train文件夾+10個valid的文件夾
#這里采用的是判斷文件名的方式進行處理
if 'train' in files[b]:#如果文件名有train
label = a #設置要標記的標簽,比如sample001_train里面都是0的圖片,標簽就是0
ss = 'F:/ele/data/' + str(files[b]) + '/' #700張訓練圖片
pics = os.listdir(ss) #得到sample001_train文件夾下的700個文件名
i = 1
while i < 701:#一共有700張
name = str(dir) + str(files[b]) + '/' + pics[i-1] + ' ' + str(int(label)) + '\n'
train.write(name)
i = i + 1
a = a + 1
if 'valid' in files[b]:
label = a1
ss = 'F:/ele/data/' + str(files[b]) + '/' #200張驗證圖片
pics = os.listdir(ss)
j = 1
while j < 201:
name = str(dir) + str(files[b]) + '/' + pics[j-1] + ' ' + str(int(label)) + '\n'
test.write(name)
j = j + 1
a1 = a1 + 1
b = b + 1
5、通過上面的python文本處理就可以得到train和test兩個txt文件如圖所示。

6、接下來就要定義數據集的類MyDataset如下所示:
class MyDataset(Dataset):
#初始化一些需要傳入的參數和數據集的調用
def __init__(self, txt, transform=None,
target_transform=None, loader=default_loader):
super(MyDataset, self).__init__()
imgs = []
fh = open(txt, 'r')
#按照傳入的路徑和txt文本參數,以只讀的形式打開這個文本
for line in fh:
#迭代該列表,按行循環txt文本
line = line.strip('\n')
line = line.rstrip('\n')
#刪除本行string字符串末尾的指定字符,
words = line.split()
#用split的方式將該行分割成列表
#split的默認參數是空格,所以不傳遞任何參數時分割空格
imgs.append((words[0], int(words[1])))
#把txt的內容讀入imgs列表保存
#word[0]是圖片信息,words[1]是label
self.imgs = imgs
self.transform = transform
self.target_transform = target_transform
self.loader = loader
#對數據進行預處理並返回想要的信息
#這個方法是必須要有的,用於按照索引讀取每個元素的具體內容
def __getitem__(self, index):
fn, label = self.imgs[index]
img = self.loader(fn) #按照標簽里面的地址讀取圖片的RGB像素
if self.transform is not None:
img = self.transform(img)
return img, label
#return哪些內容,那么我們在訓練時循環讀取每個batch時,就能獲取哪些內容
#初始化一些需要傳入的參數和數據集的調用
#這個函數必須寫,返回數據集長度,也就是多少張圖片,和Loader長度區分
def __len__(self):
return len(self.imgs)
7、其中比較關鍵的一個就是讀取圖像信息的方式default_loader,這個是定義的一個函數:
#定義讀取文件的格式
def default_loader(path):
return Image.open(path).convert('L')
L表示讀取灰度信息,只返回一個通道的數據
RGB表示返回三個通道的數據
接下來就是正式讀取信息返回到列表里面了
因為我們所讀取到的圖片的尺度是圖片原本的尺度,在設計神經網絡的過程中,輸入尺度往往有一定的要求,比如我文章使用的這個LeNet輸入必須是28*28,因此可以自己定義一個train_transform進行尺度變換如下:
train_transforms = transforms.Compose(
[transforms.RandomResizedCrop((28, 28)),
transforms.ToTensor()]
)
test_transforms = transforms.Compose(
[transforms.RandomResizedCrop((28, 28)),
transforms.ToTensor()]
)
然后加載MyDataset這個類:
train_data = MyDataset(txt=root+'train.txt', transform=train_transforms)
test_data = MyDataset(txt=root+'test.txt', transform=test_transforms)
8、當我們把所有的數字數據打包好了之后就是要放到DataLoader里面進行批量打包,batch_size可以一次讀取多張圖片的信息。
train_loader = DataLoader(dataset=train_data, batch_size=10, shuffle=True, num_workers=1)
test_loader = DataLoader(dataset=test_data, batch_size=10, shuffle=False, num_workers=1)
這樣一個自己的數字數據集就定義好了。
9、接下來就是要定義自己的網絡模型,這里舉例最簡單的LeNet網絡結構的類定義,其他網絡只需更改類即可。
class LeNet(nn.Module):
def __init__(self):
super(LeNet, self).__init__()
self.conv = nn.Sequential(
nn.Conv2d(1, 6, 5),
nn.Sigmoid(),
nn.MaxPool2d(2, 2),
nn.Conv2d(6, 16, 5),
nn.Sigmoid(),
nn.MaxPool2d(2, 2)
)
self.fc = nn.Sequential(
nn.Linear(16*4*4, 120),
nn.Sigmoid(),
nn.Linear(120, 84),
nn.Sigmoid(),
nn.Linear(84, 10)
)
def forward(self, img):
feature = self.conv(img)
output = self.fc(feature.view(img.shape[0], -1))
return output
__init__用於對網絡元素進行初始化,forward用於定義網絡前向傳播的規則。
10、接下來創建一個這樣的網絡,定義學習率,學習epoch次數,optimizer優化器類型:
net = LeNet() lr, num_epoch = 0.001, 20 optimizer = torch.optim.Adam(net.parameters(), lr=lr)
11、接下來編寫訓練代碼,在訓練過程中
def train(net, train_iter, test_iter, start_epoch, optimizer, device, num_epochs):
net = net.to(device)
print("training on:", device)
loss = torch.nn.CrossEntropyLoss()#定義損失函數
batch_count = 0 #第幾個batch,如果7000張圖片的batch_size是10,那么共有700個batch
nb = len(train_iter)#訓練數據一共有多少
for epoch in range(start_epoch, num_epochs):
#這里之所以會有start_epoch是為了后面直接加載上次未訓練完的信息
train_l_sum = 0.0#訓練損失值
train_acc_sum = 0.0#訓練精度
n, start = 0, time.time()
pbar = tqdm(enumerate(train_iter), total=nb)
#tqmd可以更直觀地觀察訓練集加載的過程
for i, (imgs, targets) in pbar:
imgs = imgs.to(device)
targets = targets.to(device)
y_hat = net(imgs)#把像素信息傳入網絡得出預測結果
l = loss(y_hat, targets)#計算預測結果和標簽的損失值
optimizer.zero_grad()#梯度清零
l.backward()#反向傳播
optimizer.step()#優化器作用
train_l_sum += l.cpu().item()
#這里使用y_hat.argmax(dim=1)是因為該網絡返回的是一個包含10個結果的向量
# 這10個結果分別是所屬類別的概率
train_acc_sum += (y_hat.argmax(dim=1) == targets).sum().cpu().item()
#10個類別里面取出最大值的索引作為結果
n += targets.shape[0]
batch_count += 1
s = '%g/%g %g' % (epoch, num_epochs - 1, len(targets))
pbar.set_description(s) # 這個就是進度條顯示
mean_loss = train_l_sum/batch_count
train_acc = train_acc_sum/n
test_acc = test(net, test_iter, device)
#下面這三個列表作為全局變量用於后面的繪圖
mean_loss_list.append(mean_loss)
train_acc_list.append(train_acc)
test_acc_list.append(test_acc)
print('loss %.4f, train_acc %.3f, test_acc %.3f' % (mean_loss, train_acc, test_acc))
#在所有的epoch訓練完之后創建節點列表保存到.pt文件里面
#這樣創建的好處是可以把當前未訓練完的epoch也保存進去
chkpt = {'epoch': epoch,
'model': net.state_dict(),
'optimizer': optimizer.state_dict()}
torch.save(chkpt, PATH)
del chkpt
12、訓練過程中的test是表示采用驗證集計算精度,考察模型泛化能力
def test(net, test_iter, device):
acc_sum, n = 0.0, 0
with torch.no_grad():
for imgs, targets in test_iter:
net.eval()
y_hat = net(imgs.to(device)).argmax(dim=1)
acc_sum += (y_hat == targets.to(device)).float().sum().cpu().item()
net.train()
n += targets.shape[0]
return acc_sum/n
13、訓練完成后將結果使用PLT庫畫出來
train(net, train_loader, test_loader, start_epoch=start_epoch, optimizer=optimizer, device=device, num_epochs=num_epoch) plot_use_plt(mean_loss_list, train_acc_list, test_acc_list, num_epoch)
14、首先要加載import matplotlib.pyplot as plt
def plot_use_plt(mean_loss_list, train_acc_list, test_acc_list, num_epoch):
x1 = range(0, num_epoch)
x2 = range(0, num_epoch)
x3 = range(0, num_epoch)
plt.subplot(1, 3, 1) #一行三列的第一列
plt.plot(x1, mean_loss_list, 'o-')
plt.title('Train_loss vs.epochs')
plt.ylabel('Train loss')
plt.subplot(1, 3, 2)
plt.plot(x2, train_acc_list, '.-')
plt.title('Train_acc vs.epochs')
plt.ylabel('Train acc')
plt.subplot(1, 3, 3)
plt.plot(x3, test_acc_list, '.-')
plt.title('Test_acc vs.epochs')
plt.ylabel('Test acc')
plt.savefig("F:/ele/show.jpg")#這一句話一定要放在plt.show()前面
plt.show()
15、訓練結果如圖所示

本項目完整代碼已上傳至github:https://github.com/logic03/Digital-recognition-with-ConvNet
