pytorch_基於cifar創建自己的數據集並訓練

基於cifar創建自己的數據集並訓練

整體流程

• resize圖像，在圖像名稱附上標簽 eg: 1_a.jpg
• 獲取圖像的像素 R G B 合並在一個列表
• 將文件名、數據、標簽和batch_label合並在一個字典中
• 用二進制的方式存儲字典內容,得到和cifar相同的
• 修改 batches.meta 中的 內容，保存並替換原文件

源碼如下

import numpy as np
import operator
from os import listdir
import sys
import pickle
import random

class makecifaData(object):
    def __init__(self):
        self.imgfolder = ''
        self.imgSavePath = ''
        self.DATA_dict = {}
        self.labels = []
        self.filenames = []
        self.data = []
        self.savefilepath = ''
        self.batch_label = ''
        self.img_size = 32

    def img_resize(self,label,_size = 32):
        # image resize to 32 32
        # image name list
        imglist = listdir(self.imgfolder)
        for img_name in imglist:
            currentpath=self.imgfolder+"/"+img_name
            im=Image.open(currentpath)
            #width=im.size[0]
            #height=im.size[1]
            x_s=_size
            y_s=_size
            out = im.resize((x_s,y_s),Image.ANTIALIAS)
            out.save(self.imgSavePath+"/"+str(label)+ '_' +str(img_name))

    def mkcf(self):
        imglist = listdir(self.imgfolder)
        for img_name in imglist:
            temp_R, temp_G, temp_B = [], [], []
            currentpath = self.imgfolder + "/" + img_name
            im = Image.open(currentpath)
            for i in range(0, self.img_size):
                for j in range(0, self.img_size):
                    temp_R.append(im.getpixel((i, j))[0])
                    temp_G.append(im.getpixel((i, j))[1])
                    temp_B.append(im.getpixel((i, j))[2])
            self.data.extend(temp_R)
            self.data.extend(temp_G)
            self.data.extend(temp_B)
            self.filenames.append(img_name.encode('utf-8'))
            self.labels.append(int(img_name.split('_')[0]))
        self.data = np.array(self.data, dtype=np.uint8)
        self.DATA_dict['batch_label'.encode('utf-8')] = self.batch_label.encode('utf-8')
        self.DATA_dict.setdefault('labels'.encode('utf-8'), self.labels)
        self.DATA_dict.setdefault('data'.encode('utf-8'), self.data)
        self.DATA_dict.setdefault('filenames'.encode('utf-8'), self.filenames)
        output = open(self.savefilepath, 'wb')
        pickle.dump(self.DATA_dict, output)
        output.close()
    def unpickle(self,file):
        with open(file, 'rb') as fo:
            dict = pickle.load(fo, encoding='latin-1')
        return dict
def main():
    MD = makecifaData()
    file_trian_val = ['train','val']
    Img_source_filename = ['cover','other','uncover']
    for save_file in file_trian_val:
        for i in range(len(Img_source_filename)):
            MD.imgfolder = 'E:\\2019PythonProject\\Pytorch_step1\\SourceImg\\'+ save_file + '\\' + Img_source_filename[i]
            MD.imgSavePath = 'E:\\2019PythonProject\\Pytorch_step1\\data_shanbu\\'+save_file
            MD.img_resize(i+1,_size=32)
    MD.img_size = 32
    MD.imgfolder = r'E:\2019PythonProject\Pytorch_step1\data_shanbu\train'
    MD.savefilepath = r'E:\2019PythonProject\Pytorch_step1\data_shanbu\data_batch_1.bin'
    MD.batch_label = 'train label 1 of 1'
    MD.mkcf()
    dic = {'num_cases_per_batch': 100, 'label_names': ['airplane', 'automobile', 'bird', 'cat', 'deer', 'dog', 'frog', 'horse', 'ship', 'truck'], 'num_vis': 3072}
    out_file = open('./batches.meta', 'wb')
    pickle.dump(dic, out_file)
    out_file.close()
    pass

開始訓練

• 修改訓練集和測試機list，原始的cifar有五個訓練集，自己得到的只有一個，同時改一下名稱.bin
  
• 遇到的報錯，看源碼無關，就直接注釋了
  

• 又遇到了報錯，是因為新文件中 字典內的key 是b'data'，都改一下即可
  
• 接下來就可以訓練自己的數據集了,訓練完可以保存模型

# 保存網絡模型 保存整個模型
        torch.save(self.net, 'model.pkl')

改變數據集圖像大小

• 之前resize后的圖像大小為3232，有點小，改成128128，再訓練
• 修改網絡參數

class LeNet(nn.Module):
    # 一般在__init__中定義網絡需要的操作算子，比如卷積、全連接算子等等
    def __init__(self):
        super(LeNet, self).__init__()
        # Conv2d的第一個參數是輸入的channel數量，第二個是輸出的channel數量，第三個是kernel size
        self.conv1 = nn.Conv2d(3, 6, 5)
        self.conv2 = nn.Conv2d(6, 16, 5)
        # 由於上一層有16個channel輸出，每個feature map大小為5*5，所以全連接層的輸入是16*5*5
        self.fc1 = nn.Linear(16 * 29 * 29, 120)
        self.fc2 = nn.Linear(120, 84)
        # 最終有10類，所以最后一個全連接層輸出數量是10
        self.fc3 = nn.Linear(84, 10)
        self.pool = nn.MaxPool2d(2, 2)
    # forward這個函數定義了前向傳播的運算，只需要像寫普通的python算數運算那樣就可以了
    def forward(self, x):
        x = F.relu(self.conv1(x))
        x = self.pool(x)
        x = F.relu(self.conv2(x))
        x = self.pool(x)
        # 下面這步把二維特征圖變為一維，這樣全連接層才能處理
        x = x.view(-1, 16*29*29)
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = self.fc3(x)
        return x

• 開始訓練