用python自動生成train,val文件——pytorch數據集
任務——分類
數據集爲5個不同類別的圖片集,每個圖片集大概有3W張圖片。所以要建立一個train訓練的txt文件和一個val驗證的txt文件,裏面放圖片的路徑,因爲只是練手用,所以不放test驗證。
最終要的結果是從每個文件裏拿出28000個訓練和剩下差不多3000個用來測試。
import os
a=0
while(a<5):
dir = '/home/zyx/data/pic/'+str(a)+'/'
label = a
files = os.listdir(dir)
files.sort()
train = open('/home/zyx/data/train.txt','a')
val = open('/home/zyx/data/val.txt', 'a')
i = 1
for file in files:
if i<29000:
fileType = os.path.split(file)
if fileType[1] == '.txt':
continue
name = str(dir) + file + ' ' + str(int(label)) +'\n'
train.write(name)
i = i+1
print(i)
else:
fileType = os.path.split(file)
if fileType[1] == '.txt':
continue
name = str(dir) +file + ' ' + str(int(label)) +'\n'
val.write(name)
i = i+1
print(i)
val.close()
train.close()
print(a)
a = a + 1
結果
然後就可以開始寫網絡和訓練模型了