MaskDetection/train.py at main · saswat0/MaskDetection · 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
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
import os
import sys
import time

import tensorflow as tf
from absl import flags, logging, app
from absl.flags import FLAGS

from utils import config
from utils.lr_scheduler import MultiStepWarmUpLR
from utils.prior_box import priors_box
from utils.utils import set_memory_growth
from dataset.preprocess import load_dataset
from network.loss import MultiBoxLoss
from network.net import SSDModel

flags.DEFINE_string('gpu', '0', 'which gpu to use')


def main(_):
    global load_t1
    os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'
    os.environ['CUDA_VISIBLE_DEVICES'] = FLAGS.gpu
    logger = tf.get_logger()
    logger.disabled = True
    logger.setLevel(logging.FATAL)
    set_memory_growth()

    weights_dir = 'checkpoints/'
    if not os.path.exists(weights_dir):
        os.mkdir(weights_dir)

    logging.info("Load configuration...")
    cfg = config.cfg
    label_classes = cfg['labels_list']
    logging.info(f"Total image sample:{cfg['dataset_len']},Total classes number:"
                 f"{len(label_classes)},classes list:{label_classes}")

    logging.info("Compute priors boxes...")
    priors, num_cell = priors_box(cfg)
    logging.info(f"Prior boxes number:{len(priors)},default anchor box number per feature map cell:{num_cell}")

    logging.info("Loading dataset...")
    train_dataset = load_dataset(cfg, priors, shuffle=True, train=True)
    val_dataset = load_dataset(cfg, priors, shuffle=False, train=False)

    logging.info("Create Model...")
    try:
        model = SSDModel(cfg=cfg, num_cell=num_cell, training=True)
        model.summary()
        tf.keras.utils.plot_model(model, to_file=os.path.join(os.getcwd(), 'model.png'),
                                  show_shapes=True, show_layer_names=True)
    except Exception as e:
        logging.error(e)
        logging.info("Create network failed.")
        sys.exit()

    init_epoch = -1

    steps_per_epoch = cfg['dataset_len'] // cfg['batch_size']
    val_steps_per_epoch = cfg['val_len'] // cfg['batch_size']

    logging.info(f"steps_per_epoch:{steps_per_epoch}")

    logging.info("Define optimizer and loss computation and so on...")

    learning_rate = MultiStepWarmUpLR(
        initial_learning_rate=cfg['init_lr'],
        lr_steps=[e * steps_per_epoch for e in cfg['lr_decay_epoch']],
        lr_rate=cfg['lr_rate'],
        warmup_steps=cfg['warmup_epoch'] * steps_per_epoch,
        min_lr=cfg['min_lr'])

    optimizer = tf.keras.optimizers.SGD(learning_rate=learning_rate, momentum=cfg['momentum'], nesterov=True)

    multi_loss = MultiBoxLoss(num_class=len(label_classes), neg_pos_ratio=3)

    train_log_dir = 'logs/train'
    train_summary_writer = tf.summary.create_file_writer(train_log_dir)

    @tf.function
    def train_step(inputs, labels):
        with tf.GradientTape() as tape:
            predictions = model(inputs, training=True)
            losses = {}
            losses['reg'] = tf.reduce_sum(model.losses)
            losses['loc'], losses['class'] = multi_loss(labels, predictions)
            total_loss = tf.add_n([l for l in losses.values()])

        grads = tape.gradient(total_loss, model.trainable_variables)
        optimizer.apply_gradients(zip(grads, model.trainable_variables))

        return total_loss, losses

    for epoch in range(init_epoch + 1, cfg['epoch']):
        try:
            start = time.time()
            avg_loss = 0.0
            for step, (inputs, labels) in enumerate(train_dataset.take(steps_per_epoch)):

                load_t0 = time.time()
                total_loss, losses = train_step(inputs, labels)
                avg_loss = (avg_loss * step + total_loss.numpy()) / (step + 1)
                load_t1 = time.time()
                batch_time = load_t1 - load_t0

                steps = steps_per_epoch * epoch + step
                with train_summary_writer.as_default():
                    tf.summary.scalar('loss/total_loss', total_loss, step=steps)
                    for k, l in losses.items():
                        tf.summary.scalar('loss/{}'.format(k), l, step=steps)
                    tf.summary.scalar('learning_rate', optimizer.lr(steps), step=steps)

                print(
                    f"\rEpoch: {epoch + 1}/{cfg['epoch']} | Batch {step + 1}/{steps_per_epoch} | Batch time {batch_time:.3f} || Loss: {total_loss:.6f} | loc loss:{losses['loc']:.6f} | class loss:{losses['class']:.6f} ",
                    end='', flush=True)

            print(
                f"\nEpoch: {epoch + 1}/{cfg['epoch']}  | Epoch time {(load_t1 - start):.3f} || Average Loss: {avg_loss:.6f}")

            with train_summary_writer.as_default():
                tf.summary.scalar('loss/avg_loss', avg_loss, step=epoch)

            if (epoch + 1) % cfg['save_freq'] == 0:
                filepath = os.path.join(weights_dir, f'weights_epoch_{(epoch + 1):03d}.h5')
                model.save_weights(filepath)
                if os.path.exists(filepath):
                    print(f">>>>>>>>>>Save weights file at {filepath}<<<<<<<<<<")

        except KeyboardInterrupt:
            print('interrupted')
            exit(0)


if __name__ == '__main__':
    app.run(main)