Dogs_vs_Cats/train.py at master · DeepVisionTeam/Dogs_vs_Cats · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
import glob
import os

import tflearn
from sklearn.model_selection import train_test_split

from image_utils import preprocess_image

pjoin = os.path.join
TRAIN_DATA = pjoin(os.path.dirname(__file__), 'images')
MODEL_PATH = pjoin(os.path.dirname(__file__), 'model')
if not os.path.exists(MODEL_PATH):
    os.makedirs(MODEL_PATH, mode=0755)
MODEL_NAME = 'resnet_dogs_vs_cats.model'


def read_data():
    X = []
    Y = []
    for f in glob.glob(TRAIN_DATA + '/*.jpg'):
        fname = os.path.basename(f)
        # 0 for cat, 1 for dog
        label = 0 if fname.startswith('cat') else 1
        image = preprocess_image(f, [256, 256, 3])
        X.append(image)
        Y.append(label)
        # split training data and validation set data
    X, X_test, y, y_test = train_test_split(X, Y,
                                            test_size=0.2,
                                            random_state=42)
    return (X, y), (X_test, y_test)


def resnet():
    # Residual blocks
    n = 5

    # Building Residual Network
    net = tflearn.input_data(shape=[None, 256, 256, 3])
    net = tflearn.conv_2d(net, 16, 3, regularizer='L2', weight_decay=0.0001)
    net = tflearn.residual_block(net, n, 16)
    net = tflearn.residual_block(net, 1, 32, downsample=True)
    net = tflearn.residual_block(net, n - 1, 32)
    net = tflearn.residual_block(net, 1, 64, downsample=True)
    net = tflearn.residual_block(net, n - 1, 64)
    net = tflearn.batch_normalization(net)
    net = tflearn.activation(net, 'relu')
    net = tflearn.global_avg_pool(net)
    # Regression
    net = tflearn.fully_connected(net, 2, activation='softmax')
    net = tflearn.regression(net, optimizer='adam',
                             loss='categorical_crossentropy')
    return net


def train():
    (X, Y), (X_test, Y_test) = read_data()
    Y = tflearn.data_utils.to_categorical(Y, 2)
    Y_test = tflearn.data_utils.to_categorical(Y_test, 2)

    # Training
    model = tflearn.DNN(resnet(), checkpoint_path='model_resnet',
                        max_checkpoints=10, tensorboard_verbose=0,
                        clip_gradients=0.)

    if os.path.exists(pjoin(MODEL_PATH, MODEL_NAME)):
        model.load(pjoin(MODEL_PATH, MODEL_NAME))

    model.fit(X, Y, n_epoch=200, validation_set=(X_test, Y_test),
              snapshot_epoch=False, snapshot_step=500,
              show_metric=True, batch_size=16, shuffle=True,
              run_id='resnet_cat_dog')

    model.save(pjoin(MODEL_PATH, MODEL_NAME))


if __name__ == '__main__':
    train()