用python自動生成train,val文件——pytorch數據集

任務——分類

數據集爲5個不同類別的圖片集,每個圖片集大概有3W張圖片。所以要建立一個train訓練的txt文件和一個val驗證的txt文件,裏面放圖片的路徑,因爲只是練手用,所以不放test驗證。

用python自動生成train,val文件——pytorch數據集
最終要的結果是從每個文件裏拿出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

結果用python自動生成train,val文件——pytorch數據集

然後就可以開始寫網絡和訓練模型了