import tensorflow as tf


# y=ax+b linear model
def regression(x):
    a = tf.Variable(tf.zeros([784, 10]), name="a")
    b = tf.Variable(tf.zeros([10]), name="b")
    y = tf.nn.softmax(tf.matmul(x, a) + b)

    return y, [a, b]


# 定义卷积模型
def convolutional(x, keep_prob):
    def conv2d(x, w):
        return tf.nn.conv2d(x, w, [1, 1, 1, 1], padding='SAME')

    def max_pool_2x2(x):
        return tf.nn.max_pool(
            x, ksize=[1, 2, 2, 1], strides=[1, 2, 2, 1], padding='SAME')

    def weight_variable(shape):
        initial = tf.truncated_normal(shape, stddev=0.1)
        return tf.Variable(initial)

    def bias_variable(shape):
        initial = tf.constant(0.1, shape=shape)
        return tf.Variable(initial)

    x_image = tf.reshape(x, [-1, 28, 28, 1])
    w_conv1 = weight_variable([5, 5, 1, 32])
    b_conv1 = bias_variable([32])
    h_conv1 = tf.nn.relu(conv2d(x_image, w_conv1) + b_conv1)
    h_pool1 = max_pool_2x2(h_conv1)

    w_conv2 = weight_variable([5, 5, 32, 64])
    b_conv2 = bias_variable([64])
    h_conv2 = tf.nn.relu(conv2d(h_pool1, w_conv2) + b_conv2)
    h_pool2 = max_pool_2x2(h_conv2)

    # 全连接层
    w_fc1 = weight_variable([7 * 7 * 64, 1024])
    b_fc1 = bias_variable([1024])
    h_pool2_flat = tf.reshape(h_pool2, [-1, 7 * 7 * 64])
    h_fc1 = tf.nn.relu(tf.matmul(h_pool2_flat, w_fc1) + b_fc1)
    h_fc1_drop = tf.nn.dropout(h_fc1, keep_prob)
    w_fc2 = weight_variable([1024, 10])
    b_fc2 = bias_variable([10])
    y = tf.nn.softmax(tf.matmul(h_fc1_drop, w_fc2))

    return y, [w_conv1, b_conv1, w_conv2, b_conv2, w_fc1, w_fc2, b_fc2]

[(30, 'tensorflow.reshape', 'tf.reshape', True, 'import tensorflow as tf\n'), (44, 'tensorflow.reshape', 'tf.reshape', True, 'import tensorflow as tf\n'), (46, 'tensorflow.nn.dropout', 'tf.nn.dropout', True, 'import tensorflow as tf\n'), (6, 'tensorflow.zeros', 'tf.zeros', True, 'import tensorflow as tf\n'), (7, 'tensorflow.zeros', 'tf.zeros', True, 'import tensorflow as tf\n'), (16, 'tensorflow.nn.conv2d', 'tf.nn.conv2d', True, 'import tensorflow as tf\n'), (19, 'tensorflow.nn.max_pool', 'tf.nn.max_pool', True, 'import tensorflow as tf\n'), (23, 'tensorflow.truncated_normal', 'tf.truncated_normal', True, 'import tensorflow as tf\n'), (24, 'tensorflow.Variable', 'tf.Variable', True, 'import tensorflow as tf\n'), (27, 'tensorflow.constant', 'tf.constant', True, 'import tensorflow as tf\n'), (28, 'tensorflow.Variable', 'tf.Variable', True, 'import tensorflow as tf\n'), (49, 'tensorflow.matmul', 'tf.matmul', True, 'import tensorflow as tf\n'), (8, 'tensorflow.matmul', 'tf.matmul', True, 'import tensorflow as tf\n'), (45, 'tensorflow.matmul', 'tf.matmul', True, 'import tensorflow as tf\n')]