programmers_contest/train.py at master · Soonmok/programmers_contest · 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
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
import numpy as np
import pandas as pd
import tensorflow as tf

from config import get_config_from_json
from data_loader import get_train_data_loader, get_test_data_loader, get_paths, get_dataset, get_data
from fixresnet import resnet50
from model import get_se_model, get_resnet_model, get_sphere_model


def train():
    config = get_config_from_json("config.json")

    config.training = True
    if config.model == "baseline":
        model = get_resnet_model(config)
    elif config.model == "senet":
        model = get_se_model(config)
    elif config.model == "resnet":
        model = resnet50()
        model.build(input_shape=(128, 128, 3))
    elif config.model == "sphere":
        model = get_sphere_model(config)
    else:
        model = None
    print(model.summary())

    # define losses and optimizer
    loss_object = tf.keras.losses.SparseCategoricalCrossentropy()
    optimizer = tf.keras.optimizers.Adam()

    train_loss = tf.keras.metrics.Mean(name='train_loss')
    train_accuracy = tf.keras.metrics.SparseCategoricalAccuracy(name='train_accuracy')

    test_loss = tf.keras.metrics.Mean(name='test_loss')
    test_accuracy = tf.keras.metrics.SparseCategoricalAccuracy(name='test_accuracy')

    @tf.function
    def train_step(x_train, y_train):
        with tf.GradientTape() as senet_tape:
            y_pred = model(x_train, training=True)
            loss = loss_object(y_train, y_pred)
        # tf.print(tf.argmax(y_pred, axis=-1))
        gradient_of_model = senet_tape.gradient(loss, model.trainable_variables)
        optimizer.apply_gradients(zip(gradient_of_model, model.trainable_variables))
        train_loss(loss)
        train_accuracy(y_train, y_pred)

    @tf.function
    def dev_step(x_dev, y_dev):
        y_pred = model(x_dev, training=False)
        # tf.print(y_pred)
        t_loss = loss_object(y_dev, y_pred)

        test_loss(t_loss)
        test_accuracy(y_dev, y_pred)

    def test_step(x_test):
        y_pred = model(x_test, training=False)
        return y_pred

    default_dev_score = 85
    train_steps = 5500 // config.data.batch_size
    dev_steps = 500 // config.data.batch_size

    train_img_paths, dev_img_paths, train_labels, dev_labels = get_paths(config, "train_vision.csv")
    test_data_loader = get_test_data_loader(config, "test_vision.csv")

    train_data_loader = get_data(config, train_img_paths, train_labels)
    train_data_loader = train_data_loader.__iter__()
    dev_data_loader = get_data(config, dev_img_paths, dev_labels)
    dev_data_loader = dev_data_loader.__iter__()

    for epoch in range(config.trainer.epochs):
        config.training = True
        for idx in range(train_steps):
            x_train, y_train = next(train_data_loader)
            train_step(x_train, y_train)
            template = 'Epoch {}, Steps : {}/{}, Loss: {}, Accuracy: {}'
            # print(template.format(epoch+1, idx+1, train_steps, train_loss.result(), train_accuracy.result() * 100))

        config.training = False
        for idx in range(dev_steps):
            x_dev, y_dev = next(dev_data_loader)
            dev_step(x_dev, y_dev)
        template = 'Epoch {}, Loss: {}, Accuracy: {}, Dev Loss: {}, Dev Accuracy: {}'
        print(template.format(epoch + 1,
                              train_loss.result(),
                              train_accuracy.result() * 100,
                              test_loss.result(),
                              test_accuracy.result() * 100))

        if test_accuracy.result() * 100 > default_dev_score:
            labels = []
            for idx, x_test in enumerate(test_data_loader):
                y_pred = test_step(x_test)
                labels.append(y_pred.numpy())
            labels = np.concatenate(labels, axis=0)
            labels = np.argmax(labels, axis=1)
            df = pd.read_csv("test_vision.csv")
            df['label'] = labels + 1
            df.drop(df.columns[0], axis=1)
            del df['filename']
            df.to_csv(f"test_result_{epoch + 1}.csv", mode='w')

        train_loss.reset_states()
        train_accuracy.reset_states()
        test_loss.reset_states()
        test_accuracy.reset_states()


if __name__ == "__main__":
    train()