Pytorch學習筆記（二）自己加載單通道圖片用作數據集訓練

      作者：灰色橡皮擦

     pytorch 在torchvision包裏面有很多的的打包好的數據集，例如minist,Imagenet-12,CIFAR10 和CIFAR100。在torchvision的dataset包裏面，用的時候直接調用就行了。具體的調用格式可以去看文檔（目前好像只有英文的）。網上也有很多源代碼。

       不過，當我們想利用自己製作的數據集來訓練網絡模型時，就要有自己的方法了。pytorch在torchvision.dataset包裏面封裝過一個函數ImageFolder（）。這個函數功能很強大，只要你直接將數據集路徑保存爲例如“train/1/1.jpg ,rain/1/2.jpg ……  ”就可以根據根目錄“./train”將數據集裝載了。

dataset.ImageFolder(root="datapath", transfroms.ToTensor())

       但是後來我發現一個問題，就是這個函數加載出來的圖像矩陣都是三通道的，並且沒有什麼參數調用可以讓其變爲單通道。如果我們要用到單通道數據集（灰度圖）的話，比如自己加載Lenet-5模型的數據集，就只能自己寫numpy數組再轉爲pytorch的Tensor()張量了。接下來是我做的過程：

       首先，還是要用到opencv,用灰度圖打開一張圖片，省事。

#讀取圖片　這裏是灰度圖　
    for item in all_path:
        img = cv2.imread(item[1],0)
        img = cv2.resize(img,(28,28))
        arr = np.asarray(img,dtype="float32")
        data_x[i ,:,:,:] = arr
        i+=1
        data_y.append(int(item[0]))
        
    data_x = data_x / 255
    data_y = np.asarray(data_y)

      其次，pytorch有自己的numpy轉Tensor函數，直接轉就行了。

    data_x = torch.from_numpy(data_x)
    data_y = torch.from_numpy(data_y)

下一步利用torch.util和torchvision裏面的dataLoader函數，就能直接得到和torchvision.dataset裏面封裝好的包相同的數據集樣本了

    dataset = dataf.TensorDataset(data_x,data_y)
    loader = dataf.DataLoader(dataset, batch_size=batchsize, shuffle=True)

最後就是自己建網絡設計參數訓練了，這部分和文檔以及github中的差不多，就不贅述了。

下面是整個程序的源代碼，我利用的還是上次的車標識別的數據集，一共分四類，用的是2層卷積核兩層全連接。

源代碼：

# coding=utf-8
import os
import cv2
import numpy as np
import random

import torch
import torch.nn as nn
import torch.utils.data as dataf
from torch.autograd import Variable
import torch.nn.functional as F
import torch.optim as optim

#訓練參數
cuda   = False
train_epoch = 20
train_lr = 0.01
train_momentum = 0.5
batchsize = 5


#測試訓練集路徑
test_path = "/home/test/"
train_path = "/home/train/"

#路徑數據
all_path =[]

def load_data(data_path):
    signal = os.listdir(data_path)
    for fsingal in signal:    
        filepath = data_path+fsingal
        filename  = os.listdir(filepath)
        for fname in filename:
            ffpath = filepath+"/"+fname
            path = [fsingal,ffpath]
            all_path.append(path)
            
#設立數據集多大
    count = len(all_path)
    data_x = np.empty((count,1,28,28),dtype="float32")
    data_y = []
#打亂順序
    random.shuffle(all_path)
    i=0;

#讀取圖片　這裏是灰度圖　最後結果是i*i*i*i
#分別表示：batch大小　，　通道數，　像素矩陣
    for item in all_path:
        img = cv2.imread(item[1],0)
        img = cv2.resize(img,(28,28))
        arr = np.asarray(img,dtype="float32")
        data_x[i ,:,:,:] = arr
        i+=1
        data_y.append(int(item[0]))
        
    data_x = data_x / 255
    data_y = np.asarray(data_y)
#     lener = len(all_path)
    data_x = torch.from_numpy(data_x)
    data_y = torch.from_numpy(data_y)
    dataset = dataf.TensorDataset(data_x,data_y)
    
    loader = dataf.DataLoader(dataset, batch_size=batchsize, shuffle=True)
     
    return  loader
#     print data_y



train_load = load_data(train_path)
test_load  = load_data(test_path)

class L5_NET(nn.Module):
    def __init__(self):
        super(L5_NET  ,self).__init__();
        #第一層輸入1，20個卷積核　每個5*5
        self.conv1 = nn.Conv2d(1 , 20 , kernel_size=5)
        #第二層輸入20，30個卷積核　每個5*5
        self.conv2 = nn.Conv2d(20 , 30 , kernel_size=5)
        #drop函數
        self.conv2_drop = nn.Dropout2d()
        #全鏈接層１，展開30*4*4，連接層50個神經元
        self.fc1 = nn.Linear(30*4*4,50)
        #全鏈接層１，50-4 ,4爲最後的輸出分類
        self.fc2 = nn.Linear(50,4)
    
    #前向傳播
    def forward(self,x):
        #池化層１ 對於第一層卷積池化，池化核2*2
        x = F.relu(F.max_pool2d(   self.conv1(x)     ,2 )   )
        #池化層2 對於第二層卷積池化，池化核2*2
        x = F.relu(F.max_pool2d(   self.conv2_drop( self.conv2(x) )  ,  2 )   )
        #平鋪軸30*4*4個神經元
        x = x.view(-1 , 30*4*4)
        #全鏈接１
        x = F.relu( self.fc1(x) )
        #dropout鏈接
        x = F.dropout(x ,   training= self.training)
        #全鏈接w
        x = self.fc2(x)
        #softmax鏈接返回結果
        return F.log_softmax(x)

model = L5_NET()
if cuda :
    model.cuda()
     

optimizer  =  optim.SGD(model.parameters()     ,  lr =train_lr , momentum = train_momentum )

#預測函數
def train(epoch):
    model.train()
    for batch_idx, (data, target) in enumerate(train_load):
        if cuda:
            data, target = data.cuda(), target.cuda()
        data, target = Variable(data), Variable(target)
        #求導
        optimizer.zero_grad()
        #訓練模型，輸出結果
        output = model(data)
        #在數據集上預測loss
        loss = F.nll_loss(output, target)
        #反向傳播調整參數pytorch直接可以用loss
        loss.backward()
        #SGD刷新進步
        optimizer.step()
        #實時輸出
        if batch_idx % 10 == 0:
            print('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format(
                epoch, batch_idx * len(data), len(train_load.dataset),
                100. * batch_idx / len(train_load), loss.data[0]))
#             
            
#測試函數
def test(epoch):
    model.eval()
    test_loss = 0
    correct = 0
    for data, target in test_load:
        
        if cuda:
            data, target = data.cuda(), target.cuda()
            
        data, target = Variable(data, volatile=True), Variable(target)
        #在測試集上預測
        output = model(data)
        #計算在測試集上的loss
        test_loss += F.nll_loss(output, target).data[0]
        #獲得預測的結果
        pred = output.data.max(1)[1] # get the index of the max log-probability
        #如果正確，correct+1
        correct += pred.eq(target.data).cpu().sum()
    
    #loss計算
    test_loss = test_loss
    test_loss /= len(test_load)
    #輸出結果
    print('\nThe {} epoch result : Average loss: {:.6f}, Accuracy: {}/{} ({:.2f}%)\n'.format(
        epoch,test_loss, correct, len(test_load.dataset),
        100. * correct / len(test_load.dataset)))

for epoch in range(1, train_epoch+ 1):
    train(epoch)
    test(epoch)

        
        
        

最後的訓練結果和在keras下差不多，不過我訓練的時候好像把訓練集和測試集弄反了，數目好像測試集比訓練集還多，有點尷尬，不過無傷大雅。結果圖如下：