使用CNN 实现小型Google你画电脑猜
你画电脑猜
Google 的猜画小歌风靡一时,使用 CNN 我们也可以搞一个出来。
数据集
Google AI Lab 已经开源了用来训练的数据集:https://github.com/googlecreativelab/quickdraw-dataset
里边有 354 类共5千万张涂鸦图片,并提供多种格式,从原始的笔迹数据到转换好的位图都有。我们将使用其中的 numpy 格式版本。
如果能够访问 Google Drive 的话,可以从 https://console.cloud.google.com/storage/quickdraw_dataset/full/numpy_bitmap 下载或者有装过 gsutil 的话,可以 gsutil -m rsync gs://quickdraw_dataset/full/numpy_bitmap/ .
由于这个真实项目的数据集规模相当大,有文件小很多的矢量格式可以下载,比如 binary 里就是一种。另外 https://console.cloud.google.com/storage/browser/quickdraw_dataset/sketchrnn 这里提供了一个npz格式的子集,只包含7万5千张图片,numpy.load() 既能加载 .npy 也能加载 .npz 倒是方便。不过所有的矢量格式,下载回来都要预处理转成位图,不然没法喂给咱们的 CNN 网络。这两个文件格式可以参考 binaryparser.ipbn
模型
Google AI Lab 在 https://arxiv.org/abs/1704.03477 这篇论文里解释了他们是怎么干的,并且在 github 上提供了实现,那是一个相对复杂的 RNN 模型。这里我们用一个较为简单的 CNN 模型来解决这个问题。
模型结构如下:
导入
import tensorflow as tf
import numpy as np
import os
import matplotlib.pyplot as plt
from sklearn.model_selection import train_test_split
加载数据集
names = []
# 加载数据集
def load_dataset(dafile):
count = 0
# 每次加载多少张图片
images_per_class = 10000
# 大小为(0,784) 784的来源就是上面的原因
X = np.empty([0, 784])
Y = np.empty([0, 1])
# 加载数据集[butterfly.npy , clock.npy, door.npy, ...]
for file in os.listdir(dafile):
fn = dafile + file
names.append(file)
# 加载npy, 提取10000张图片分析
images = np.load(fn).astype('float32')
# (10000,784)
images = images[0:images_per_class, :]
# 返回给定形状和类型的新数组
# 比如 count = 3时, label = [[3],[3],[3],...], (10000,1)
labels = np.full((images_per_class, 1), count)
# 连接,count = 2时,X(30000, 784), Y(30000, 1)
X = np.concatenate((X, images), axis = 0)
Y = np.concatenate((Y, labels), axis = 0)
count += 1
# 洗牌, 打乱顺序,0 到 Y.shape[0],
order = np.random.permutation(Y.shape[0])
X = X[order, :]
Y = Y[order, :]
# 拆分测试数据集和训练集,1:9
X_train, X_test, Y_train, Y_test = train_test_split(
X, Y, random_state = 0, test_size = 0.1)
return X_train, Y_train, X_test, Y_test, count
下载小型的对应文件夹,有12种类别,放在data的目录下
dataset_files = "./data/"
X_train, Y_train, X_test, Y_test, count = load_dataset(dataset_files)
模型结构:
对应模型代码实现
# 重置图,可以反复运行,
tf.reset_default_graph()
# 所需变量
x = None
y = None
# 三层卷积,三层池化,三层激活层,两层全连接层
# 权重,
w_conv = [None, None, None]
# 偏差值
b_conv = [None, None, None]
# 卷积
r_conv = [None, None, None]
# 激活
h_conv = [None, None, None]
# 池化层
h_pool = [None, None, None]
keep_prob = None
# 全连接层
w_fc = [None, None]
b_fc = [None, None]
h_fc = [None, None]
def build_model():
global x, y
global w_conv, b_conv, r_conv, h_conv
global h_pool, keep_prob
global w_fc, b_fc, h_fc
# 输入层
with tf.variable_scope('input'):
x = tf.placeholder(tf.float32
, shape=[None, 28, 28, 1]
, name = 'x')
# 卷积层 1: conv1 [None, 28, 28, 16]
with tf.variable_scope('conv1'):
# truncated_normal 从截断的正态分布中输出随机值
w_conv[0] = tf.Variable(tf.truncated_normal([3, 3, 1, 16]
, stddev = 0.1))
b_conv[0] = tf.Variable(tf.truncated_normal([16]
, stddev = 0.1))
r_conv[0] = tf.nn.conv2d(x, w_conv[0], strides=[1,1,1,1]
, padding="SAME")
# wx + b
r_conv[0] = r_conv[0] + b_conv[0]
# 激活层 1: h1
with tf.variable_scope('h1'):
h_conv[0] = tf.nn.relu(r_conv[0])
# 池化层 1: pool1 [None, 14, 14, 16]
with tf.variable_scope('pool1'):
h_pool[0] = tf.nn.max_pool(h_conv[0], ksize = [1,2,2,1]
, strides = [1,2,2,1], padding = 'SAME')
# 卷积层 2: conv2
with tf.variable_scope('conv2'):
w_conv[1] = tf.Variable(tf.truncated_normal([3, 3, 16, 32]
, stddev = 0.1))
b_conv[1] = tf.Variable(tf.truncated_normal([32]
, stddev = 0.1))
r_conv[1] = tf.nn.conv2d(h_pool[0], w_conv[1], strides=[1,1,1,1]
, padding = 'SAME')
r_conv[1] = r_conv[1] + b_conv[1]
# 激活层 2: h2
with tf.variable_scope('h2'):
h_conv[1] = tf.nn.relu(r_conv[1])
# 池化层 2: pool2 [None, 7, 7, 32]
with tf.variable_scope('pool2'):
h_pool[1] = tf.nn.max_pool(h_conv[1], ksize=[1,2,2,1]
, strides = [1,2,2,1], padding='SAME')
# 卷积层 3: conv3
with tf.variable_scope('conv3'):
w_conv[2] = tf.Variable(tf.truncated_normal([3, 3, 32, 64]
, stddev = 0.1))
b_conv[2] = tf.Variable(tf.truncated_normal([64]
, stddev = 0.1))
r_conv[2] = tf.nn.conv2d(h_pool[1], w_conv[2], strides=[1,1,1,1]
, padding = 'SAME')
r_conv[2] = r_conv[2] + b_conv[2]
# 激活层 3: h3
with tf.variable_scope('h3'):
h_conv[2] = tf.nn.relu(r_conv[2])
# 池化层 3: pool3 [None, 4, 4, 64]
with tf.variable_scope('pool3'):
h_pool[2] = tf.nn.max_pool(h_conv[2], ksize=[1,2,2,1]
, strides = [1,2,2,1], padding='SAME')
# 全连接层 1 : fc1
with tf.variable_scope('fc1'):
keep_prob = tf.placeholder(tf.float32)
h_pool3_flat = tf.reshape(h_pool[2], [-1, 4*4*64])
w_fc[0] = tf.Variable(tf.truncated_normal([4*4*64, 1024]
, stddev = 0.1))
b_fc[0] = tf.Variable(tf.truncated_normal([1024]
, stddev = 0.1))
h_fc[0] = tf.nn.relu(tf.matmul(h_pool3_flat, w_fc[0]))
h_fc[0] = h_fc[0] + b_fc[0]
# 正则化
h_fc[0] = tf.nn.dropout(h_fc[0], keep_prob)
# 全连接层 2 : fc2
with tf.variable_scope('fc2'):
# 分类 count 12种
w_fc[1] = tf.Variable(tf.truncated_normal([1024, count]
, stddev = 0.1))
b_fc[1] = tf.Variable(tf.truncated_normal([count]
, stddev = 0.1))
h_fc[1] = tf.matmul(h_fc[0], w_fc[1])
h_fc[1] = h_fc[1] + b_fc[1]
# output,多分类softmax函数
with tf.variable_scope('output'):
prediction = tf.nn.softmax(h_fc[1])
return prediction
训练和保存模型
# 训练模型
y = tf.placeholder(tf.float32, shape=[None, count], name = 'y')
batch_size = 512
train_step = None
model = build_model()
reduce_sum = -tf.reduce_sum(y * tf.log(model)
, reduction_indices = [1])
loss = tf.reduce_mean(reduce_sum)
train_step = tf.train.AdamOptimizer(1e-4).minimize(loss)
saver = tf.train.Saver()
with tf.Session() as sess:
init = tf.global_variables_initializer()
sess.run(init)
for i in range(10):
train_batch = zip(range(0, len(X_train), batch_size)
, range(batch_size, len(X_train) + 1, batch_size))
for start, end in train_batch:
sess.run(train_step, feed_dict = {x:X_train[start:end]
,y:Y_train[start:end]
,keep_prob:0.75})
saver.save(sess, './model/mymodel')
加载恢复模型
def predictByModel(img):
with tf.Session() as sess:
saver = tf.train.import_meta_graph('model/mymodel.meta')
saver.restore(sess, tf.train.latest_checkpoint("model/"))
result = sess.run(model, feed_dict={x:img, keep_prob:1})
rIdx = np.argmax(result)
pred = names[rIdx]
return pred
实现预测
from matplotlib import image
mytest = image.imread('./test.png')
mytest_gray = mytest[:, :, 0]
mytest_img = np.reshape(mytest_gray, [1, 28, 28, 1])
print('your draw is ', predictByModel(mytest_img))
plt.imshow(mytest)
your draw is eraser.npy
<matplotlib.image.AxesImage at 0x7f3d4b6e0160>
总结:这里只是使用CNN实现一个小项目,达到使用CNN分类的情况,上面的步骤省去了,测试。直接训练完后,拿来预测分类。
数据集和源代码: GitHub