無論是訓練機器學習或是深度學習，第一步當然是先划分數據集啦，今天小白整理了一些划分數據集的方法，希望大佬們多多指教啊，嘻嘻~
首先看一下數據集的樣子，flower_data文件夾下有四個文件夾，每個文件夾表示一種花的類別
划分數據集的主要步驟：
1. 定義一個空字典，用來存放各個類別的訓練集、測試集和驗證集，字典的key是類別，value也是一個字典，存放該類別的訓練集、測試集和驗證集；
2.使用python獲取所有的類別文件夾；
3.對每個類別划分訓練集、測試集和驗證集：（1）把該類別的所有有效圖片放入一個列表中；（2）設置一個隨機數對列表進行划分。
具體的代碼實現如下所示
import glob import os.path import random import numpy as np # 圖片數據文件夾 INPUT_DATA = './flower_data' # 這個函數從數據文件夾中讀取所有的圖片列表並按訓練、驗證、測試數據分開 # testing_percentage和validation_percentage指定了測試數據集和驗證數據集的大小 def create_image_lists(testing_percentage,validation_percentage): # 得到的所有圖片都存在result這個字典里，key為類別的名稱，value值也是一個字典，存放的是該類別的 # 文件名、訓練集、測試集和驗證集 result = {} # 獲取當前目錄下所有的子目錄,這里x 是一個三元組(root,dirs,files)，第一個元素表示INPUT_DATA當前目錄， # 第二個元素表示當前目錄下的所有子目錄,第三個元素表示當前目錄下的所有的文件 sub_dirs = [x[0] for x in os.walk(INPUT_DATA)] # sub_dirs = ['./flower_data','./flower_data\\daisy','./flower_data\\dandelion', # './flower_data\\roses','./flower_data\\sunflowers','./flower_data\\tulips'] # 每個子目錄表示一類花，現在對每類花划分訓練集、測試集和驗證集 # sub_dirs[0]表示當前文件夾本身的地址，不予考慮，只考慮他的子目錄（各個類別的花） for sub_dir in sub_dirs[1:]: # 獲取當前目錄下所有的有效圖片文件 extensions = ['jpg','jpeg'] # 把圖片存放在file_list列表里 file_list = [] # os.path.basename(sub_dir)返回sub_sir最后的文件名 # 如os.path.basename('./flower_data/daisy')返回daisy dir_name = os.path.basename(sub_dir) for extension in extensions: file_glob = os.path.join(INPUT_DATA,dir_name,'*.'+extension) # glob.glob(file_glob)獲取指定目錄下的所有圖片，存放在file_list中 file_list.extend(glob.glob(file_glob)) if not file_list: continue # 通過目錄名獲取類別的名稱,返回將字符串中所有大寫字符轉換為小寫后生成的字符串 label_name = dir_name.lower() # 初始化當前類別的訓練數據集、測試數據集和驗證數據集 training_images = [] testing_images = [] validation_images = [] for file_name in file_list: base_name = os.path.basename(file_name) # 隨機將數據分到訓練數據集、測試數據集和驗證數據集 # 產生一個隨機數，最大值為100 chance = np.random.randint(100) if chance < validation_percentage: validation_images.append(base_name) elif chance < (testing_percentage+validation_percentage): testing_images.append(base_name) else: training_images.append(base_name) # 將當前類別是數據放入結果字典 result[label_name]={'dir':dir_name, 'training':training_images, 'testing':testing_images, 'validation':validation_images} # 返回整理好的所有數據 return result result = create_image_lists(10,30) print(result)
運行結果：
可以看出字典result中有五個key,表示五個類別。
下圖是各個類別的划分情況：
