代码拉取完成，页面将自动刷新
#!/usr/bin/env python3
# -*- coding: utf-8 -*-
# @Time : 18-5-18 下午7:31
# @Author : Luo Yao
# @Site : http://icode.baidu.com/repos/baidu/personal-code/Luoyao
# @File : train_lanenet.py
# @IDE: PyCharm Community Edition
"""
训练lanenet模型
"""
import argparse
import os
import os.path as ops
import time
import glog as log
import numpy as np
import tensorflow as tf
import sys
from config import global_config
from lanenet_model import lanenet_merge_model
from data_provider import lanenet_data_processor
CFG = global_config.cfg
def init_args():
"""
:return:
"""
parser = argparse.ArgumentParser()
parser.add_argument('--dataset_dir', type=str, help='The training dataset dir path')
parser.add_argument('--net', type=str, help='Which base net work to use', default='vgg')
parser.add_argument('--weights_path', type=str, help='The pretrained weights path')
return parser.parse_args()
def average_gradients(tower_grads):
"""Calculate the average gradient for each shared variable across all towers.
Note that this function provides a synchronization point across all towers.
Args:
tower_grads: List of lists of (gradient, variable) tuples. The outer list
is over individual gradients. The inner list is over the gradient
calculation for each tower.
Returns:
List of pairs of (gradient, variable) where the gradient has been averaged
across all towers.
"""
average_grads = []
for grad_and_vars in zip(*tower_grads):
# Note that each grad_and_vars looks like the following:
# ((grad0_gpu0, var0_gpu0), ... , (grad0_gpuN, var0_gpuN))
grads = []
for g, _ in grad_and_vars:
# Add 0 dimension to the gradients to represent the tower.
expanded_g = tf.expand_dims(g, 0)
# Append on a 'tower' dimension which we will average over below.
grads.append(expanded_g)
# Average over the 'tower' dimension.
grad = tf.concat(grads, 0)
grad = tf.reduce_mean(grad, 0)
# Keep in mind that the Variables are redundant because they are shared
# across towers. So .. we will just return the first tower's pointer to
# the Variable.
v = grad_and_vars[0][1]
grad_and_var = (grad, v)
average_grads.append(grad_and_var)
return average_grads
def forward(batch_queue, net, phase, scope, optimizer=None):
img_batch, label_instance_batch, label_existence_batch = batch_queue.dequeue()
inference = net.inference(img_batch, phase, 'lanenet_loss')
_ = net.loss(inference, label_instance_batch, label_existence_batch, 'lanenet_loss')
total_loss = tf.add_n(tf.get_collection('total_loss', scope))
instance_loss = tf.add_n(tf.get_collection('instance_seg_loss', scope))
existence_loss = tf.add_n(tf.get_collection('existence_pre_loss', scope))
out_logits = tf.add_n(tf.get_collection('instance_seg_logits', scope))
# calculate the accuracy
out_logits = tf.nn.softmax(logits=out_logits)
out_logits_out = tf.argmax(out_logits, axis=-1)
pred_0 = tf.count_nonzero(tf.multiply(tf.cast(tf.equal(label_instance_batch, 0), tf.int32),
tf.cast(tf.equal(out_logits_out, 0), tf.int32)),
dtype=tf.int32)
pred_1 = tf.count_nonzero(tf.multiply(tf.cast(tf.equal(label_instance_batch, 1), tf.int32),
tf.cast(tf.equal(out_logits_out, 1), tf.int32)),
dtype=tf.int32)
pred_2 = tf.count_nonzero(tf.multiply(tf.cast(tf.equal(label_instance_batch, 2), tf.int32),
tf.cast(tf.equal(out_logits_out, 2), tf.int32)),
dtype=tf.int32)
pred_3 = tf.count_nonzero(tf.multiply(tf.cast(tf.equal(label_instance_batch, 3), tf.int32),
tf.cast(tf.equal(out_logits_out, 3), tf.int32)),
dtype=tf.int32)
pred_4 = tf.count_nonzero(tf.multiply(tf.cast(tf.equal(label_instance_batch, 4), tf.int32),
tf.cast(tf.equal(out_logits_out, 4), tf.int32)),
dtype=tf.int32)
gt_all = tf.count_nonzero(tf.cast(tf.greater(label_instance_batch, 0), tf.int32), dtype=tf.int32)
gt_back = tf.count_nonzero(tf.cast(tf.equal(label_instance_batch, 0), tf.int32), dtype=tf.int32)
pred_all = tf.add(tf.add(tf.add(pred_1, pred_2), pred_3), pred_4)
accuracy = tf.divide(tf.cast(pred_all, tf.float32), tf.cast(gt_all, tf.float32))
accuracy_back = tf.divide(tf.cast(pred_0, tf.float32), tf.cast(gt_back, tf.float32))
# Compute mIoU of Lanes
overlap_1 = pred_1
union_1 = tf.add(tf.count_nonzero(tf.cast(tf.equal(label_instance_batch, 1),
tf.int32), dtype=tf.int32),
tf.count_nonzero(tf.cast(tf.equal(out_logits_out, 1),
tf.int32), dtype=tf.int32))
union_1 = tf.subtract(union_1, overlap_1)
IoU_1 = tf.divide(tf.cast(overlap_1, tf.float32), tf.cast(union_1, tf.float32))
overlap_2 = pred_2
union_2 = tf.add(tf.count_nonzero(tf.cast(tf.equal(label_instance_batch, 2),
tf.int32), dtype=tf.int32),
tf.count_nonzero(tf.cast(tf.equal(out_logits_out, 2),
tf.int32), dtype=tf.int32))
union_2 = tf.subtract(union_2, overlap_2)
IoU_2 = tf.divide(tf.cast(overlap_2, tf.float32), tf.cast(union_2, tf.float32))
overlap_3 = pred_3
union_3 = tf.add(tf.count_nonzero(tf.cast(tf.equal(label_instance_batch, 3),
tf.int32), dtype=tf.int32),
tf.count_nonzero(tf.cast(tf.equal(out_logits_out, 3),
tf.int32), dtype=tf.int32))
union_3 = tf.subtract(union_3, overlap_3)
IoU_3 = tf.divide(tf.cast(overlap_3, tf.float32), tf.cast(union_3, tf.float32))
overlap_4 = pred_4
union_4 = tf.add(tf.count_nonzero(tf.cast(tf.equal(label_instance_batch, 4),
tf.int64), dtype=tf.int32),
tf.count_nonzero(tf.cast(tf.equal(out_logits_out, 4),
tf.int64), dtype=tf.int32))
union_4 = tf.subtract(union_4, overlap_4)
IoU_4 = tf.divide(tf.cast(overlap_4, tf.float32), tf.cast(union_4, tf.float32))
IoU = tf.reduce_mean(tf.stack([IoU_1, IoU_2, IoU_3, IoU_4]))
tf.get_variable_scope().reuse_variables()
if optimizer is not None:
grads = optimizer.compute_gradients(total_loss)
else:
grads = None
return total_loss, instance_loss, existence_loss, accuracy, accuracy_back, IoU, out_logits_out, grads
def train_net(dataset_dir, weights_path=None, net_flag='vgg'):
train_dataset_file = ops.join(dataset_dir, 'train_gt.txt')
val_dataset_file = ops.join(dataset_dir, 'val_gt.txt')
assert ops.exists(train_dataset_file)
phase = tf.placeholder(dtype=tf.string, shape=None, name='net_phase')
train_dataset = lanenet_data_processor.DataSet(train_dataset_file)
val_dataset = lanenet_data_processor.DataSet(val_dataset_file)
net = lanenet_merge_model.LaneNet()
tower_grads = []
global_step = tf.Variable(0, trainable=False)
learning_rate = tf.train.polynomial_decay(CFG.TRAIN.LEARNING_RATE, global_step,
CFG.TRAIN.EPOCHS, power=0.9)
optimizer = tf.train.MomentumOptimizer(learning_rate=learning_rate, momentum=0.9)
img, label_instance, label_existence = train_dataset.next_batch(CFG.TRAIN.BATCH_SIZE)
batch_queue = tf.contrib.slim.prefetch_queue.prefetch_queue(
[img, label_instance, label_existence], capacity=2 * CFG.TRAIN.GPU_NUM, num_threads=CFG.TRAIN.CPU_NUM)
val_img, val_label_instance, val_label_existence = val_dataset.next_batch(CFG.TRAIN.VAL_BATCH_SIZE)
val_batch_queue = tf.contrib.slim.prefetch_queue.prefetch_queue(
[val_img, val_label_instance, val_label_existence], capacity=2 * CFG.TRAIN.GPU_NUM,
num_threads=CFG.TRAIN.CPU_NUM)
with tf.variable_scope(tf.get_variable_scope()):
for i in range(CFG.TRAIN.GPU_NUM):
with tf.device('/gpu:%d' % i):
with tf.name_scope('tower_%d' % i) as scope:
total_loss, instance_loss, existence_loss, accuracy, accuracy_back, _, out_logits_out, \
grad = forward(batch_queue, net, phase, scope, optimizer)
tower_grads.append(grad)
with tf.name_scope('test_%d' % i) as scope:
val_op_total_loss, val_op_instance_loss, val_op_existence_loss, val_op_accuracy, \
val_op_accuracy_back, val_op_IoU, _, _ = forward(val_batch_queue, net, phase, scope)
grads = average_gradients(tower_grads)
train_op = optimizer.apply_gradients(grads, global_step=global_step)
train_cost_time_mean = []
train_instance_loss_mean = []
train_existence_loss_mean = []
train_accuracy_mean = []
train_accuracy_back_mean = []
saver = tf.train.Saver()
model_save_dir = 'model/culane_lanenet/culane_scnn'
if not ops.exists(model_save_dir):
os.makedirs(model_save_dir)
train_start_time = time.strftime('%Y-%m-%d-%H-%M-%S', time.localtime(time.time()))
model_name = 'culane_lanenet_{:s}_{:s}.ckpt'.format(net_flag, str(train_start_time))
model_save_path = ops.join(model_save_dir, model_name)
sess_config = tf.ConfigProto(device_count={'GPU': CFG.TRAIN.GPU_NUM}, allow_soft_placement=True)
sess_config.gpu_options.per_process_gpu_memory_fraction = CFG.TRAIN.GPU_MEMORY_FRACTION
sess_config.gpu_options.allow_growth = CFG.TRAIN.TF_ALLOW_GROWTH
sess_config.gpu_options.allocator_type = 'BFC'
with tf.Session(config=sess_config) as sess:
with sess.as_default():
if weights_path is None:
log.info('Training from scratch')
init = tf.global_variables_initializer()
sess.run(init)
else:
log.info('Restore model from last model checkpoint {:s}'.format(weights_path))
saver.restore(sess=sess, save_path=weights_path)
# 加载预训练参数
if net_flag == 'vgg' and weights_path is None:
pretrained_weights = np.load(
'./data/vgg16.npy',
encoding='latin1').item()
for vv in tf.trainable_variables():
weights = vv.name.split('/')
if len(weights) >= 3 and weights[-3] in pretrained_weights:
try:
weights_key = weights[-3]
weights = pretrained_weights[weights_key][0]
_op = tf.assign(vv, weights)
sess.run(_op)
except Exception as e:
continue
tf.train.start_queue_runners(sess=sess)
for epoch in range(CFG.TRAIN.EPOCHS):
t_start = time.time()
_, c, train_accuracy, train_accuracy_back, train_instance_loss, train_existence_loss, _ = \
sess.run([train_op, total_loss, accuracy, accuracy_back, instance_loss, existence_loss, out_logits_out],
feed_dict={phase: 'train'})
cost_time = time.time() - t_start
train_cost_time_mean.append(cost_time)
train_instance_loss_mean.append(train_instance_loss)
train_existence_loss_mean.append(train_existence_loss)
train_accuracy_mean.append(train_accuracy)
train_accuracy_back_mean.append(train_accuracy_back)
if epoch % CFG.TRAIN.DISPLAY_STEP == 0:
print(
'Epoch: {:d} loss_ins= {:6f} ({:6f}) loss_ext= {:6f} ({:6f}) accuracy= {:6f} ({:6f}) '
'accuracy_back= {:6f} ({:6f}) mean_time= {:5f}s '.format(epoch + 1, train_instance_loss,
np.mean(train_instance_loss_mean),
train_existence_loss,
np.mean(train_existence_loss_mean),
train_accuracy,
np.mean(train_accuracy_mean),
train_accuracy_back,
np.mean(train_accuracy_back_mean),
np.mean(train_cost_time_mean)))
if epoch % 500 == 0:
train_cost_time_mean.clear()
train_instance_loss_mean.clear()
train_existence_loss_mean.clear()
train_accuracy_mean.clear()
train_accuracy_back_mean.clear()
if epoch % 1000 == 0:
saver.save(sess=sess, save_path=model_save_path, global_step=epoch)
if epoch % 10000 != 0 or epoch == 0:
continue
val_cost_time_mean = []
val_instance_loss_mean = []
val_existence_loss_mean = []
val_accuracy_mean = []
val_accuracy_back_mean = []
val_IoU_mean = []
for epoch_val in range(int(len(val_dataset) / CFG.TRAIN.VAL_BATCH_SIZE / CFG.TRAIN.GPU_NUM)):
t_start_val = time.time()
c_val, val_accuracy, val_accuracy_back, val_IoU, val_instance_loss, val_existence_loss = \
sess.run(
[val_op_total_loss, val_op_accuracy, val_op_accuracy_back,
val_op_IoU, val_op_instance_loss, val_op_existence_loss],
feed_dict={phase: 'test'})
cost_time_val = time.time() - t_start_val
val_cost_time_mean.append(cost_time_val)
val_instance_loss_mean.append(val_instance_loss)
val_existence_loss_mean.append(val_existence_loss)
val_accuracy_mean.append(val_accuracy)
val_accuracy_back_mean.append(val_accuracy_back)
val_IoU_mean.append(val_IoU)
if epoch_val % 1 == 0:
print('Epoch_Val: {:d} loss_ins= {:6f} ({:6f}) '
'loss_ext= {:6f} ({:6f}) accuracy= {:6f} ({:6f}) accuracy_back= {:6f} ({:6f}) '
'mIoU= {:6f} ({:6f}) mean_time= {:5f}s '.
format(epoch_val + 1, val_instance_loss, np.mean(val_instance_loss_mean), val_existence_loss,
np.mean(val_existence_loss_mean), val_accuracy, np.mean(val_accuracy_mean),
val_accuracy_back, np.mean(val_accuracy_back_mean), val_IoU, np.mean(val_IoU_mean),
np.mean(val_cost_time_mean)))
val_cost_time_mean.clear()
val_instance_loss_mean.clear()
val_existence_loss_mean.clear()
val_accuracy_mean.clear()
val_accuracy_back_mean.clear()
val_IoU_mean.clear()
if __name__ == '__main__':
# init args
args = init_args()
# train lanenet
train_net(args.dataset_dir, args.weights_path, net_flag=args.net)
此处可能存在不合适展示的内容，页面不予展示。您可通过相关编辑功能自查并修改。
如您确认内容无涉及 不当用语 / 纯广告导流 / 暴力 / 低俗色情 / 侵权 / 盗版 / 虚假 / 无价值内容或违法国家有关法律法规的内容，可点击提交进行申诉，我们将尽快为您处理。
马建仓 AI 助手