tensorflow简洁实现MNIST识别

一、MyData.py

• 数据读取，进行了简单的归一化，即除以255.0
• dataset_path是保存图片路径和标签的txt的路径，每行的格式是/home/lwd/data/mnist/image/59999.png 8
• MNIST转化成图片看这里

import os
import cv2
import random
import numpy as np
import tensorflow as tf

class Dataset(object):
    def __init__(self, dataset_path, train, batch_size=1):
        self.all = []
        for line in open(dataset_path):
            self.all.append(line)
        if train == True:
            random.shuffle(self.all)
        self.bs = batch_size
        self.batch_num = len(self.all) // self.bs
        self.bid = 0
    def __iter__(self):
        return self
    def __len__(self):
        return self.batch_num
    def __next__(self):
        batch_image = np.zeros((self.bs, 28, 28, 1))
        batch_label = np.zeros((self.bs, 10))
        if self.bid < self.batch_num:
            for j in range(self.bs):
                sp=self.all[self.bid*self.bs+j].split()
                im=cv2.imread(sp[0], 0).astype(np.float32) / 255.0
                im = np.expand_dims(im, axis = 2)
                batch_image[j, :, :, :] = im
                lab=int(sp[1])
                gt=np.zeros((10))
                gt[lab]=1
                batch_label[j,:]=gt
            self.bid += 1
            return batch_image, batch_label
        else:
            self.bid = 0
            random.shuffle(self.all)
            raise StopIteration

二、mnist.py

• 两种方式实现了两种结构的网络
• 参考了这个github项目

import tensorflow as tf
import numpy as np
import random
import cv2,sys
import MyData


data=tf.placeholder(tf.float32, [None, 28, 28, 1],name='data')
label=tf.placeholder(tf.float32, [None, 10], name='label')
# model
x_image = tf.reshape(data, [-1, 28, 28, 1])
with tf.variable_scope('conv1'): # output is 28x28
    weight = tf.get_variable(name='weight', dtype=tf.float32, trainable=True, shape=[5,5,1,6], initializer=tf.random_normal_initializer(stddev=0.01))
    conv1 = tf.nn.conv2d(x_image, weight, [1,1,1,1], 'SAME')
    bias = tf.get_variable(name='bias', shape=6, trainable=True, dtype=tf.float32, initializer=tf.constant_initializer(0.0))
    conv1 = tf.nn.bias_add(conv1, bias)
with tf.variable_scope('pool1'):
    pool1 = tf.nn.max_pool(conv1, ksize=[1, 2, 2, 1], strides=[1, 2, 2, 1], padding='VALID')
with tf.variable_scope('conv2'): # output is 10x10
    weight2 = tf.get_variable(name='weight', dtype=tf.float32, trainable=True, shape=[5,5,6,16], initializer=tf.random_normal_initializer(stddev=0.01))
    conv2 = tf.nn.conv2d(pool1, weight2, [1,1,1,1], 'VALID')
    bias2 = tf.get_variable(name='bias', shape=16, trainable=True, dtype=tf.float32, initializer=tf.constant_initializer(0.0))
    conv2 = tf.nn.bias_add(conv2, bias2)
with tf.variable_scope('pool2'):
    pool2 = tf.nn.max_pool(conv2, ksize=[1, 2, 2, 1], strides=[1, 2, 2, 1], padding='VALID')
with tf.variable_scope('dense1'):
    flat = tf.reshape(pool2, [-1, 5*5*16])
    dense1 = tf.layers.dense(inputs=flat, units=200, activation=tf.nn.relu, use_bias=True)
with tf.variable_scope('dense2'):
    dense2 = tf.layers.dense(inputs=dense1, units=10, activation=None, use_bias=True)
y = tf.nn.softmax(dense2)
# loss
cross_entropy = tf.nn.softmax_cross_entropy_with_logits(labels=label, logits=dense2)
loss=tf.reduce_sum(cross_entropy)
train_step = tf.train.GradientDescentOptimizer(1e-4).minimize(loss)
# accuracy
correct_prediction = tf.equal(tf.argmax(y, 1), tf.argmax(label, 1))
accuracy = tf.reduce_mean(tf.cast(correct_prediction, tf.float32))

'''
data = tf.placeholder(tf.float32, [None, 28, 28, 1])
label = tf.placeholder(tf.float32, [None, 10])
# model
K = 4  
L = 8  
M = 12  
N = 200
W1 = tf.Variable(tf.truncated_normal([5, 5, 1, K], stddev=0.1))  
B1 = tf.Variable(tf.ones([K])/10)
W2 = tf.Variable(tf.truncated_normal([5, 5, K, L], stddev=0.1))
B2 = tf.Variable(tf.ones([L])/10)
W3 = tf.Variable(tf.truncated_normal([4, 4, L, M], stddev=0.1))
B3 = tf.Variable(tf.ones([M])/10)
W4 = tf.Variable(tf.truncated_normal([7 * 7 * M, N], stddev=0.1))
B4 = tf.Variable(tf.ones([N])/10)
W5 = tf.Variable(tf.truncated_normal([N, 10], stddev=0.1))
B5 = tf.Variable(tf.ones([10])/10)
stride = 1  # output is 28x28
Y1 = tf.nn.relu(tf.nn.conv2d(data, W1, strides=[1, stride, stride, 1], padding='SAME') + B1)
stride = 2  # output is 14x14
Y2 = tf.nn.relu(tf.nn.conv2d(Y1, W2, strides=[1, stride, stride, 1], padding='SAME') + B2)
stride = 2  # output is 7x7
Y3 = tf.nn.relu(tf.nn.conv2d(Y2, W3, strides=[1, stride, stride, 1], padding='SAME') + B3)
YY = tf.reshape(Y3, shape=[-1, 7 * 7 * M])
Y4 = tf.nn.relu(tf.matmul(YY, W4) + B4)
Ylogits = tf.matmul(Y4, W5) + B5
y = tf.nn.softmax(Ylogits)
# loss
cross_entropy = tf.nn.softmax_cross_entropy_with_logits(logits=Ylogits, labels=label)
loss = tf.reduce_mean(cross_entropy)*100
correct_prediction = tf.equal(tf.argmax(y, 1), tf.argmax(label, 1))
accuracy = tf.reduce_mean(tf.cast(correct_prediction, tf.float32))
#lr = 0.0001 +  tf.train.exponential_decay(0.003, step, 2000, 1/math.e)
train_step = tf.train.AdamOptimizer(1e-4).minimize(loss)
'''


gpu_options = tf.GPUOptions(per_process_gpu_memory_fraction=0.1)
sess = tf.Session(config=tf.ConfigProto(gpu_options=gpu_options))
sess.run(tf.global_variables_initializer())
input_data = MyData.Dataset('/home/lwd/data/mnist/train.txt', True, 32)
test_data = MyData.Dataset('/home/lwd/data/mnist/test.txt', False, 32)
# save ckpt model
saver = tf.train.Saver()
# for tensorboard
summary_writer = tf.summary.FileWriter('./log/', sess.graph)

# load
# saver = tf.train.import_meta_graph('./checkpoint/my.meta')
# saver.restore(sess,tf.train.latest_checkpoint('./checkpoint'))
# xh = 1
# acc = 0
# for item in test_data:
#     yy = sess.run(y, feed_dict={data:item[0], label:item[1]})
#     for k in range(yy.shape[0]):
#         # print(xh, np.argmax(yy[k]))
#         if(np.argmax(yy[k]) == np.argmax(item[1][k])) : acc += 1
#         xh += 1
#     #summary_writer.add_summary(summary, xh)
# print(acc * 1.0 / (xh - 1))
# sys.exit(0)

for i in range(100):
    total = 0
    cnt = 0
    tl = 0
    for item in input_data:
        _, acc, lo = sess.run([train_step, accuracy, loss], feed_dict={data:item[0], label:item[1]})
        total += acc
        cnt += 1.0
        tl += lo
    print(total/cnt, tl / cnt)
    if total/cnt > 0.98:
        xh = 1
        for item in test_data:
            yy = sess.run(y, feed_dict={data:item[0]})
            for k in range(yy.shape[0]):
                print(xh, np.where(yy[k]==np.max(yy[k])))
                xh += 1
        # save ckpt model
        saver.save(sess, './checkpoint/my')
        sys.exit(0)

三、ckpt转pb

import tensorflow as tf
from tensorflow.python.framework import graph_util

with tf.Graph().as_default() as graph_old:
    isess = tf.InteractiveSession()
    ckpt_filename = './checkpoint/my'
    isess.run(tf.global_variables_initializer())
    saver = tf.train.import_meta_graph(ckpt_filename+'.meta', clear_devices=True)
    saver.restore(isess, ckpt_filename)
    
    constant_graph = graph_util.convert_variables_to_constants(isess, isess.graph_def, ["Softmax"])
    constant_graph = graph_util.remove_training_nodes(constant_graph)
    with tf.gfile.GFile('./pb_model/model.pb', mode='wb') as f:
        f.write(constant_graph.SerializeToString())