前段时间风靡盆友圈的【猜画小歌】大家应该都玩过,画一幅灵魂画作由AI猜。国外有盆友通过Google Colab训练模型,并基于TensorFlow.js实现了浏览器版的你画我猜,下面是对原文的翻译
代码和Demo
代码:https://github.com/zaidalyafeai/zaidalyafeai.github.io/tree/master/sketcher
Demo:https://zaidalyafeai.github.io/sketcher/
Google Colab notebook:https://colab.research.google.com/github/zaidalyafeai/zaidalyafeai.github.io/blob/master/sketcher/Sketcher.ipynb
数据集
我们使用卷积神经网络CNN来识别画出的图像的类型,在Quick Draw数据集上进行模型训练,这个数据集有345个类别的大约5千万张手绘图像
流程
我们将使用 Keras 框架在Google Colab 免费提供的 GPU 上训练模型,然后在浏览器里基于TensorFlow.js运行模型,我之前写过TensorFlow.js的教程,在继续阅读下问之前,可以先看一下这个教程。下图为这个项目的实现流程
在Colab上训练
Google提供了免费的GPU,你可以在这个教程查看如何创建一个notebook和进行GPU编程
导入
我们使用Keras,以tensorflow作为后端运行
import os
import glob
import numpy as np
from tensorflow.keras import layers
from tensorflow import keras
import tensorflow as tf
加载数据
由于内存有限,所以我们不会训练所有的分类,只选了其中的100个分类作为数据集。每个分类的数据都可以在Google Cloud上以numpy数组的形式找到,数组的大小为 [N, 784],N为这个分类下图片的数量。我们先下载这个数据集
import urllib.request
def download():
base = 'https://storage.googleapis.com/quickdraw_dataset/full/numpy_bitmap/'
for c in classes:
cls_url = c.replace('_', '%20')
path = base+cls_url+'.npy'
print(path)
urllib.request.urlretrieve(path, 'data/'+c+'.npy')
由于内存有限,我们只加载每个分类里的5000张图片,并且预览其中的20%作为测试数据
def load_data(root, vfold_ratio=0.2, max_items_per_class= 5000 ):
all_files = glob.glob(os.path.join(root, '*.npy'))
#initialize variables
x = np.empty([0, 784])
y = np.empty([0])
class_names = []
#load a subset of the data to memory
for idx, file in enumerate(all_files):
data = np.load(file)
data = data[0: max_items_per_class, :]
labels = np.full(data.shape[0], idx)
x = np.concatenate((x, data), axis=0)
y = np.append(y, labels)
class_name, ext = os.path.splitext(os.path.basename(file))
class_names.append(class_name)
data = None
labels = None
#separate into training and testing
permutation = np.random.permutation(y.shape[0])
x = x[permutation, :]
y = y[permutation]
vfold_size = int(x.shape[0]/100*(vfold_ratio*100))
x_test = x[0:vfold_size, :]
y_test = y[0:vfold_size]
x_train = x[vfold_size:x.shape[0], :]
y_train = y[vfold_size:y.shape[0]]
return x_train, y_train, x_test, y_test, class_names
数据预处理
在训练模型之前,需要对数据进行预处理。该模型将使用规模为 [N, 28, 28, 1] 的批处理,并且输出规模为 [N, 100] 的概率。
# Reshape and normalize
x_train = x_train.reshape(x_train.shape[0], image_size, image_size, 1).astype('float32')
x_test = x_test.reshape(x_test.shape[0], image_size, image_size, 1).astype('float32')
x_train /= 255.0
x_test /= 255.0
# Convert class vectors to class matrices
y_train = keras.utils.to_categorical(y_train, num_classes)
y_test = keras.utils.to_categorical(y_test, num_classes)
创建模型
我们将创建一个简单的卷积神经网络,模型越简单、参数越少越好。因为我们将在浏览器里运行模型,并且希望能快速得到预测结果。因此我们的模型只包含 3 个卷积层和 2 个全连接层:
# Define model
model = keras.Sequential()
model.add(layers.Convolution2D(16, (3, 3),
padding='same',
input_shape=x_train.shape[1:], activation='relu'))
model.add(layers.MaxPooling2D(pool_size=(2, 2)))
model.add(layers.Convolution2D(32, (3, 3), padding='same', activation= 'relu'))
model.add(layers.MaxPooling2D(pool_size=(2, 2)))
model.add(layers.Convolution2D(64, (3, 3), padding='same', activation= 'relu'))
model.add(layers.MaxPooling2D(pool_size =(2,2)))
model.add(layers.Flatten())
model.add(layers.Dense(128, activation='relu'))
model.add(layers.Dense(100, activation='softmax'))
# Train model
adam = tf.train.AdamOptimizer()
model.compile(loss='categorical_crossentropy',
optimizer=adam,
metrics=['top_k_categorical_accuracy'])
print(model.summary())
拟合、验证及测试
接着我们对模型进行了 5 轮训练,将训练数据分成了 256 批输入模型,并且分离出 10% 作为验证集。
#fit the model
model.fit(x = x_train, y = y_train, validation_split=0.1, batch_size = 256, verbose=2, epochs=5)
#evaluate on unseen data
score = model.evaluate(x_test, y_test, verbose=0)
print('Test accuarcy: {:0.2f}%'.format(score[1] * 100))
下面是训练结果
测试结果达到了 92.20% 的 top 5 准确率
准备web格式的模型
得到准确率符合预期的模型之后,我们将模型保存下来
model.save('keras.h5')
安装tensorflow.js
!pip install tensorflowjs
转换模型
!mkdir model
!tensorflowjs_converter --input_format keras keras.h5 model/
转换之后生成了一些权重文件和包含模型架构的 json 文件
通过 zip 将模型进行压缩,以便将其下载到本地机器上:
!zip -r model.zip model
下载模型
from google.colab import files
files.download('model.zip')
在浏览器进行预测
在这部分,我们将介绍如何在浏览器加载模型,并进行预测。我们将创建一个300*300的画布,这里我将不会介绍画布相关实现,而是将重点放在TensorFlow.js这部分
加载模型
为了使用 TensorFlow.js,我们首先需要加载对应脚本
<script src="https://cdn.jsdelivr.net/npm/@tensorflow/tfjs@latest"> </script>
你的本地机器上需要有一台运行中的服务器来托管权重文件,你可以使用apache 服务器或者和我一样托管在GitHub上
model = await tf.loadModel('model/model.json')
预处理
在进行预测前,需要把得到的数据进行一定的处理,首先从canvas上获取图片数据
//the minimum boudning box around the current drawing
const mbb = getMinBox()
//cacluate the dpi of the current window
const dpi = window.devicePixelRatio
//extract the image data
const imgData = canvas.contextContainer.getImageData(mbb.min.x * dpi, mbb.min.y * dpi,
(mbb.max.x - mbb.min.x) * dpi, (mbb.max.y - mbb.min.y) * dpi);
文章稍后将介绍 getMinBox(),dpi 变量用来根据屏幕像素的密度对裁剪出的画布进行拉伸
我们将画布当前的图像数据转化为一个张量,调整大小并进行归一化处理:
function preprocess(imgData)
{
return tf.tidy(()=>{
//convert the image data to a tensor
let tensor = tf.fromPixels(imgData, numChannels= 1)
//resize to 28 x 28
const resized = tf.image.resizeBilinear(tensor, [28, 28]).toFloat()
// Normalize the image
const offset = tf.scalar(255.0);
const normalized = tf.scalar(1.0).sub(resized.div(offset));
//We add a dimension to get a batch shape
const batched = normalized.expandDims(0)
return batched
})
}
我们使用 model.predict 进行预测,这将返回一个规模为「N, 100」的概率。
const pred = model.predict(preprocess(imgData)).dataSync()
我们可以排序得到Top5的概率
提升准确率
我们的模型接受的输入数据是规模为 [N, 28, 28, 1] 的张量,我们绘图画布的尺寸为 300*300,但是用户可能在这个画布上画了个很小的图像,因此我们需要剪裁出包含手绘图像的那部分,通过找到手绘内容最左上方和右下方的点即可实现
//record the current drawing coordinates
function recordCoor(event)
{
//get current mouse coordinate
var pointer = canvas.getPointer(event.e);
var posX = pointer.x;
var posY = pointer.y;
//record the point if withing the canvas and the mouse is pressed
if(posX >=0 && posY >= 0 && mousePressed)
{
coords.push(pointer)
}
}
//get the best bounding box by finding the top left and bottom right cornders
function getMinBox(){
var coorX = coords.map(function(p) {return p.x});
var coorY = coords.map(function(p) {return p.y});
//find top left corner
var min_coords = {
x : Math.min.apply(null, coorX),
y : Math.min.apply(null, coorY)
}
//find right bottom corner
var max_coords = {
x : Math.max.apply(null, coorX),
y : Math.max.apply(null, coorY)
}
return {
min : min_coords,
max : max_coords
}
}
手绘测试
下图显示了一些第一次绘制的图像以及准确率最高的类别,所有的手绘图像都是我用鼠标画的,用笔绘制的话应该会得到更高的准确率
原文地址:https://medium.com/tensorflow/train-on-google-colab-and-run-on-the-browser-a-case-study-8a45f9b1474e