# Author : Hellcat

# Time   : 2017/12/8

import os

import time

import numpy as np

import tensorflow as tf

import cifar10

# import models.tutorials.image.cifar10.cifar10 as cifar10

import cifar10_input

# import models.tutorials.image.cifar10.cifar10_input as cifar10_input

data_dir = './cifar-10/'

cifar10.maybe_download_and_extract()

max_steps = 3000

batch_size = 128

IMAGE_SIZE = 24

NUM_CLASSES = 10

NUM_EXAMPLES_PER_EPOCH_FOR_TRAIN = 50000

NUM_EXAMPLES_PER_EPOCH_FOR_EVAL = 10000

def variable_with_weight_loss(shape, stddev, wl):

    '''

    权参数初始化，会使用L2正则化

    :param shape: 权重尺寸

    :param stddev: 标准差

    :param wl: L2项稀疏

    :return: 权重变量

    '''

    var = tf.Variable(tf.truncated_normal(shape, stddev=stddev))

    if wl is not None:

        weight_loss = tf.multiply(tf.nn.l2_loss(var),wl,name='weight_loss')

        tf.add_to_collection('losses', weight_loss)

    return var

# 读取图片并预处理

images_train, labels_train = cifar10_input.distorted_inputs(data_dir=data_dir,

                                                            batch_size=batch_size)

images_test, labels_test = cifar10_input.inputs(eval_data=True,

                                                data_dir=data_dir,

                                                batch_size=batch_size)

# 输入：24*24的RGB三色通道图片

image_holder = tf.placeholder(tf.float32, [batch_size,24,24,3])

label_holder = tf.placeholder(tf.int32, [batch_size])

# 卷积->relu激活->最大池化->标准化

weight1 = variable_with_weight_loss(shape=[5,5,3,64],stddev=5e-2,wl=0.)

bias1 = tf.Variable(tf.constant(0.,shape=[64]))

kernel1 = tf.nn.conv2d(image_holder,weight1,[1,1,1,1],padding='SAME')

conv1 = tf.nn.relu(tf.nn.bias_add(kernel1,bias1))

pool1 = tf.nn.max_pool(conv1, ksize=[1,3,3,1], strides=[1,2,2,1],padding='SAME')

norm1 = tf.nn.lrn(pool1,4,bias=1.,alpha=0.001/9.,beta=0.75)

# 卷积->relu激活->标准化->最大池化

weight2 = variable_with_weight_loss(shape=[5,5,64,64],stddev=5e-2,wl=0.)

bias2 = tf.Variable(tf.constant(0.,shape=[64]))

kernel2 = tf.nn.conv2d(norm1,weight2,[1,1,1,1],padding='SAME')

conv2 = tf.nn.relu(tf.nn.bias_add(kernel2,bias2))

norm2 = tf.nn.lrn(conv2,4,bias=1.,alpha=0.001/9.,beta=0.75)

pool2 = tf.nn.max_pool(norm2, ksize=[1,3,3,1], strides=[1,2,2,1], padding='SAME')

# 全连接：reshape尺寸->384

reshape = tf.reshape(pool2,[batch_size,-1])

dim = reshape.get_shape()[1].value # <-----动态获取tensor大小的方法

weight3 = variable_with_weight_loss(shape=[dim, 384], stddev=0.04, wl=0.004)

bias3 = tf.Variable(tf.constant(0.1,shape=[384]))

local3 = tf.nn.relu(tf.matmul(reshape,weight3)+bias3)

print('reshape.get_shape()[1]:',

      reshape.get_shape()[1],

      type(reshape.get_shape()[1])) # <-----问题

print('reshape.get_shape()[1].value:',

      reshape.get_shape()[1].value,

      type(reshape.get_shape()[1].value)) # <-----问题

print('tf.shape(reshape):',tf.shape(reshape)) # <-----问题

# 全连接：384->192

weight4 = variable_with_weight_loss(shape=[384,192],stddev=0.04,wl=0.004)

bias4 = tf.Variable(tf.constant(0.1,shape=[192]))

# tf.nn.bias_add 是 tf.add 的一个特例

# 二者均支持 broadcasting（广播机制），也即两个操作数最后一个维度保持一致。

# 除了支持最后一个维度保持一致的两个操作数相加外，tf.add 还支持第二个操作数是一维的情况

local4 = tf.nn.relu(tf.nn.bias_add(tf.matmul(local3,weight4), bias4))

# 全连接：192->10

weight5 = variable_with_weight_loss(shape=[192,10],stddev=1/192.,wl=0.)

bias5 = tf.Variable(tf.constant(0.,shape=[10]))

logits = tf.add(tf.matmul(local4,weight5),bias5)

def loss(logits, labels):

    '''

    loss函数计算

    :param logits: 网络输出结果

    :param labels: 真实标签

    :return:

    '''

    labels = tf.cast(labels,tf.int64)

    # 使用SoftMax交叉熵函数，loss计算自带softmax层

    # 对比下面的print可以得知输出的是128张图片各自的交叉熵

    cross_entropy = tf.nn.sparse_softmax_cross_entropy_with_logits(logits=logits,labels=labels,

                                                                   name='cross_entropy_per_example')

    print('交叉熵：',cross_entropy.get_shape()) # (128,)

    cross_entropy_mean = tf.reduce_mean(cross_entropy,name='cross_entropy')

    tf.add_to_collection('losses', cross_entropy_mean)

    # tf.add_n():多项连加

    return tf.add_n(tf.get_collection('losses'),name='total_loss')

loss = loss(logits, label_holder)

train_op = tf.train.AdamOptimizer(1e-3).minimize(loss)

# 输出结果top_k准确率，默认为1

top_k_op = tf.nn.in_top_k(logits, label_holder, 1)

sess = tf.InteractiveSession()

tf.global_variables_initializer().run()

# 训练部分

# 启动数据增强队列

tf.train.start_queue_runners()

for step in range(max_steps):

    start_time = time.time()

    image_batch, label_batch = sess.run([images_train, labels_train])

    _, loss_value = sess.run([train_op, loss],

                             feed_dict={image_holder:image_batch, label_holder:label_batch})

    duration = time.time() - start_time

    if step % 10 == 0:

        examples_per_sec = batch_size / duration

        sec_per_batch = float(duration)

        format_str = ('step %d, loss=%.2f (%.1f examples/sec; %.3f sec/batch)')

        print(format_str % (step, loss_value, examples_per_sec, sec_per_batch))

# 测试部分

num_examples = 10000

import math

num_iter = int(math.ceil(num_examples / batch_size))

true_count = 0

total_sample_count = num_iter * batch_size

step = 0

while step < num_iter:

    image_batch, label_batch = sess.run([images_test, labels_test])

    predictions = sess.run(top_k_op, feed_dict={image_holder:image_batch,

                                                label_holder:label_batch})

    true_count += np.sum(predictions)

    step += 1

prediction = predictions / total_sample_count

print('precision @ 1 = %.3f' % prediction)

print('reshape.get_shape()[1]:',

      reshape.get_shape()[1],

      type(reshape.get_shape()[1]))

print('reshape.get_shape()[1].value:',

      reshape.get_shape()[1].value,

      type(reshape.get_shape()[1].value)) # 动态获取tensor shape的方式，必须调用.value

print('tf.shape(reshape):',tf.shape(reshape))