Pytorch自己加載單通道圖片用作數(shù)據(jù)集訓練的實例

更新時間：2020年01月18日 15:06:56 作者：灰色橡皮擦

今天小編就為大家分享一篇Pytorch自己加載單通道圖片用作數(shù)據(jù)集訓練的實例，具有很好的參考價值，希望對大家有所幫助。一起跟隨小編過來看看吧

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

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

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

但是后來我發(fā)現(xiàn)一個問題，就是這個函數(shù)加載出來的圖像矩陣都是三通道的，并且沒有什么參數(shù)調(diào)用可以讓其變?yōu)閱瓮ǖ?。如果我們要用到單通道?shù)據(jù)集（灰度圖）的話，比如自己加載Lenet-5模型的數(shù)據(jù)集，就只能自己寫numpy數(shù)組再轉為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函數(shù)，直接轉就行了。

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

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

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

最后就是自己建網(wǎng)絡設計參數(shù)訓練了，這部分和文檔以及github中的差不多，就不贅述了。

下面是整個程序的源代碼，我利用的還是上次的車標識別的數(shù)據(jù)集，一共分四類，用的是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
 
#訓練參數(shù)
cuda = False
train_epoch = 20
train_lr = 0.01
train_momentum = 0.5
batchsize = 5
 
 
#測試訓練集路徑
test_path = "/home/test/"
train_path = "/home/train/"
 
#路徑數(shù)據(jù)
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)
   
#設立數(shù)據(jù)集多大
 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大小　，　通道數(shù)，　像素矩陣
 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函數(shù)
  self.conv2_drop = nn.Dropout2d()
  #全鏈接層１，展開30*4*4，連接層50個神經(jīng)元
  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個神經(jīng)元
  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 )
 
#預測函數(shù)
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)
  #在數(shù)據(jù)集上預測loss
  loss = F.nll_loss(output, target)
  #反向傳播調(diào)整參數(shù)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]))
#    
   
#測試函數(shù)
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下差不多，不過我訓練的時候好像把訓練集和測試集弄反了，數(shù)目好像測試集比訓練集還多，有點尷尬，不過無傷大雅。結果圖如下：