AlexNet系列I:基于Tensorflow 模型代码

最新推荐文章于 2024-07-19 18:17:55 发布

Elffer

最新推荐文章于 2024-07-19 18:17:55 发布

阅读量323

点赞数

分类专栏：深度学习

本文链接：https://blog.csdn.net/bmwlwg/article/details/100738572

版权

本文档介绍了如何使用Tensorflow实现AlexNet模型，应用于Udacity的交通灯检测项目。包含4个部分：train.py负责训练，test.py进行测试，model.py定义模型，images.py处理图像。

摘要由CSDN通过智能技术生成

AlexNet

以下代码是针对 Udacity Traffic light detection的应用:

1. train.py

import tensorflow as tf

import model
import utils
import images


def get_loss(result, y):
    cross_entropy = tf.nn.sigmoid_cross_entropy_with_logits(logits=result, labels=y)
    return tf.reduce_mean(cross_entropy)


def get_optimizer(loss):
    train_variables = tf.trainable_variables()
    optimizer = tf.train.AdamOptimizer(learning_rate=utils.lr).minimize(loss, var_list=train_variables)
    return optimizer


def train():
    x_train, y_train, _ = images.read_traffic_light(True)
    x_test, y_test, _ = images.read_traffic_light(False)

    train_batches = x_train.shape[0]

    x, y, one_hot, result = model.get_model(is_train=True)
    loss = get_loss(result, one_hot)
    optimizer = get_optimizer(loss)

    saver = tf.train.Saver()
    with tf.Session() as sess:
        sess.run(tf.global_variables_initializer())

        for epoch in range(utils.epochs):
            for batch in range(train_batches // utils.batch_size):
                start = batch * utils.batch_size
                next_x = x_train[start:start + utils.batch_size]
                next_y = y_train[start:start + utils.batch_size]

                sess.run(optimizer, feed_dict={x: next_x, y: next_y})

            loss_result = sess.run(loss, feed_dict={x: x_test, y: y_test})
            print("epoch: {}, loss: {}".format(epoch, loss_result))

        saver.save(sess, "./result/result.ckpt")


if __name__ == '__main__':
    train()