ishigenの技術ブログ

卓球×テクノロジー、機械学習

yolov3を自前データで学習(失敗編)

今回はyolov3を自前のデータで学習させて、動かしてみます。学習させるのは動画のデータです。

環境

ubuntu 16.04
python 3.6
keras
Geforce GTX 1050

こちらのkeras実装を使わせていただきました。正直なところ、VOTTを利用して教師データを作成しYOLOで学習させるのであれば、別の実装を使ったほうがいいとは思います。

github.com

参考記事

qiita.com


また、アノテーションにはVOTTを使いました。

ishigentech.hatenadiary.jp

このモデルの学習用のデータフォーマットと出力されるファイル形式が異なるので、モデルのフォーマットに合わせたテキストデータを作成します。このモデルの学習用のフォーマットは

画像のパス min_x,min_y,max_x,max_y,クラス番号 min_x,min_y,max_x,max_y,クラス番号
画像のパス ~

上記のように一つのテキストファイルに全ての画像のパス、ボックスの位置、クラス番号を記載します。訓練用と検証用のデータは分けずに記載して構いません。train.py実行時に分けてくれるみたいです。

以下のソースでtrain.txtという名前の訓練用データを作成しました。

WIDTH = 640
HEIGHT = 360
txtpath = '...../data/obj/'
imgpath = '...../data/obj/'
w = open('train.txt','w')
w.close()
#ファイル数に合わせてループさせる
for i in range(1, 103):
    txtfilename = '....._frame_'+str(i)+'.txt'
    imgfilename = '....._frame_'+str(i)+'.jpg'
    txt = imgpath+imgfilename
    f = open(txtpath+txtfilename)
    boxes = f.readlines()
    f.close()
    for box in boxes:
        boxElementList = box.split()
        boxElementList = [float(n) for n in boxElementList]
        for i, boxElement in enumerate(boxElementList):
            if i == 0: continue
            if i == 1:
                x_mean = boxElement * WIDTH
            if i == 2:
                y_mean = boxElement * HEIGHT
            if i == 3:
                x_width = boxElement * WIDTH / 2
            if i == 4:
                y_height = boxElement * HEIGHT / 2
            #print(boxElement)
        x_min = round(x_mean - x_width)
        y_min = round(y_mean - y_height)
        x_max = round(x_mean + x_width)
        y_max = round(y_mean + y_height)
        txt += ' '+str(x_min)+','+str(y_min)+','+str(x_max)+','+str(y_max)+','
        txt += str(int(boxElementList[0]))
    txt += '\n'
    #print(txt)
    w = open('train.txt','a')
    w.write(txt)
    w.close()

batch_sizeはとりあえず1でやってみます。

python train.py

これで学習できると思っていると、私のマシンに搭載されているGeforce GTX 1050(2G)だと落ちます。
はじめの50エポックまではいけることもありますが、50エポックから学習する層の数が増えるので、間違いなく落ちます。

f:id:ishigentech:20181006211133p:plain

遅いかもしれないけど、学習はできると思っていたので、けっこうショックです。。。
AWSでやってみるか、1080Tiとかを頑張って買うかしないようです。

環境を変えてまた再チャレンジします。