train.py

import os
os.environ["CUDA_VISIBLE_DEVICES"] = "0"
import numpy as np
import tensorflow as tf
import cv2
import cfg

from shufflenetv2_centernet_V2 import ShuffleNetV2_centernet
# from shufflenetv2_centernet_V2_SEB import Shufflenetv2_Centernet_SEB
# from yolov3_centernet_V2 import yolov3_centernet
from create_label import CreatGroundTruth


def parse_color_data(example_proto):
    features = {"img_raw": tf.FixedLenFeature([], tf.string),
                "label": tf.FixedLenFeature([], tf.string),
                "width": tf.FixedLenFeature([], tf.int64),
                "height": tf.FixedLenFeature([], tf.int64)}
    parsed_features = tf.parse_single_example(example_proto, features)
    img = parsed_features["img_raw"]
    img = tf.decode_raw(img, tf.uint8)
    width = parsed_features["width"]
    height = parsed_features["height"]
    img = tf.reshape(img, [height, width, 3])
    img = tf.cast(img, tf.float32) * (1. / 255.) - 0.5
    label = parsed_features["label"]
    label = tf.decode_raw(label, tf.float32)

    return img, label
def erase_invalid_val(sequence):
    label = []
    h, w = sequence.shape
    mask = (sequence != -1.0)
    for i in range(h):
        seq_new = sequence[i][mask[i]]
        label.append(list(seq_new))
    return label

filenames = [cfg.tfrecords_path]
dataset = tf.data.TFRecordDataset(filenames)
dataset = dataset.shuffle(buffer_size=1000)
dataset = dataset.map(parse_color_data)
val1=tf.constant(-0.5,tf.float32)
val2 = tf.constant(-1, tf.float32)
dataset = dataset.padded_batch(cfg.batch_size, padded_shapes=([None, None, 3], [None]), padding_values=(val1, val2))
dataset = dataset.repeat(cfg.epochs)
iterator = dataset.make_one_shot_iterator()
next_element = iterator.get_next()

train_start_time = cv2.getTickCount()


model=ShuffleNetV2_centernet()

sess = tf.Session()
sess.run(tf.global_variables_initializer())
train_writer = tf.summary.FileWriter("shufflenetv2_voc_summary", sess.graph)

saver=tf.train.Saver(max_to_keep=20)

if 0:#reload model
    model_file = tf.train.latest_checkpoint('shufflenetv2_voc/')
    saver.restore(sess, model_file)
    print("reload ckpt from "+model_file)
try:
    while True:
        batch_start_time=cv2.getTickCount()
        img_batch, label_batch = sess.run(next_element)
        label_batch = erase_invalid_val(label_batch)
        cls_gt_batch, size_gt_batch = CreatGroundTruth(label_batch)
        feed = {model.inputs: img_batch,
                model.is_training:True,
                model.size_gt:size_gt_batch,
                model.cls_gt:cls_gt_batch
               }
        fetches = [
                   model.cls_loss,
                   model.size_loss,
                   model.total_loss,
                   model.global_step,
                   model.lr,
                   model.merged_summay,
                   model.train_op,
                   ]
        cls_loss,size_loss,total_loss,global_step, lr, summary, _ = sess.run(fetches, feed)

        train_writer.add_summary(summary, global_step)
        time_elapsed = (cv2.getTickCount()-batch_start_time)/cv2.getTickFrequency()

        if global_step%200==0:
            saver.save(sess,"shufflenetv2_seb_voc/shufflenetv2_seb_voc.ckpt",global_step=global_step)
            # saver.save(sess,"shufflenetv2_face_SEB_summary/shufflenetv2_face_SEB.ckpt",global_step=global_step)
            # saver.save(sess,"shufflenetv2_voc/shufflenetv2_voc.ckpt",global_step=global_step)
            # saver.save(sess,"yolov3_voc/yolov3_voc.ckpt",global_step=global_step)
            # saver.save(sess,"shufflenev2_face_ori/shufflenev2_face.ckpt",global_step=global_step)

        if global_step % 10 == 0:
            print("-------Training {0}th batch-------".format(global_step))
            print("global_step:{0} total_loss:{1:0.3f} cls_loss:{2:0.3f} size_loss:{3:0.3f}".format(global_step,total_loss,cls_loss,size_loss))
            print("learning_rate:{0:0.6f}".format(lr))
            # print("predicts:", predicts)
            print('The batch run total {0:0.5f}s'.format(time_elapsed))

except tf.errors.OutOfRangeError:
    print('Training has completed...')
train_total_time=(cv2.getTickCount()-train_start_time)/cv2.getTickFrequency()
print('Training has stopped...')
hour=train_total_time // 3600
minute=(train_total_time-hour*3600)//60
print('Training runs {:.0f}h {:.0f}m...'.format(hour,minute))
sess.close()