从 0 到 1 实现：AI 版你画我猜小游戏

1分钟看图掌握核心观点👇

引  言

• 云端模型：将模型部署在服务集群上，提供一些API能力供端侧来调用，端侧无需处理计算部分，只需调用API获取计算结果即可。比如OpenAI的官方API就是如此。

  
  

• 端侧模型：直接将模型部署在终端设备上，模型的计算、推理完全依赖终端设备，具有更高的实时性、私密性、安全性，但同时对终端设备等硬件要求也较高。

你画我猜AI版-玩法简介

训练模型

首先，第一步是训练模型。

根据上面的玩法简介，我们知道它本质上是一个基于视觉的图片分类AI模型，而这个模型的功能是：输入图片数据后，模型可以计算出图片的分类置信结果。例如，输入一张小猫的图片，模型的分类计算结果可能为：[猫 90%，狗8%，猪2%]，表示模型认为这张图有90%的概率是只猫，8%的概率是条狗，2%概率是只猪。

这样一来，我们通过将用户手绘的canvas中的图片数据丢给模型，并把模型输出的置信概率最大的分类当作AI的猜测结果，就可以模拟出AI猜词的互动了。

而实现这个模型也很简单，但我们需要了解一些深度学习神经网络的知识以及tensorflow.js的基础用法。如果对这两者不太熟悉，可能需要先自行google一下，做点知识储备。

那么假设大家已经有了一些基础的神经网络、TensorFlow.js基础知识，就可以利用TensorFlow.js轻松搭建一个基于CNN的图片分类模型。

在进入模型训练之前，我们需要先获取数据集。

数据集是训练模型的基础，我们可以自己创建数据集（这很困难、费时），或者寻找一些开源数据集。刚好Google Lab提供了一套完整的开源涂鸦数据集（The Quick Draw Dataset），数据集中包含了345个不同类别的涂鸦数据集合，总共有5000万份涂鸦数据，足够我们挑选使用。

下载完训练数据之后，接下来我们需要搭建模型结构并进行模型训练。

如果我们下载了demo代码，可以看到项目结构如下，主要内容为3个部分：

项目目录/
├── 📁 src/                    
│   ├── 📄 index.ts            # 程序入口文件
│   ├── 📁 data/               # 数据集
│   │   ├── 📄 Apple.npy
│   │   ├── 📄 The Great Wall.npy  
│   │   └── 📄 ...       
│   └── 📁 model/              # 训练模型相关
│       ├── 📄 doodle-data.model.ts  # 数据加载
│       └── 📄 classifier.model.ts   # 模型结构
├── 📄 package.json

import { Classifier } from './model/classifier.model';
import { DoodleData } from './model/doodle-data.model';


async function main(){
  const data = new DoodleData({
    directoryData: 'src/data',
    maxImageClass: 20000
  });


  // 1. 加载训练数据
  data.loadData();
  // 2. 创建模型
  const model = new Classifier(data);
  // 3. 训练模型
  await model.train();
  // 4. 保存模型参数
  await model.save();
}


main();

// 加载data目录下的数据
loadData() {
  this.classes = fs.readdirSync(this.directoryData)
    .filter((x) => x.endsWith('.npy'))
    .map((x) => x.replace('.npy', ''));
}


// 数据生成器，预处理数据为tensor张量
*dataGenerator() {
  // ...
  for (let j = 0; j < bytes.length; j = j + this.IMAGE_SIZE) {
    const singleImage = bytes.slice(j, j + this.IMAGE_SIZE);
    const image = tf
      .tensor(singleImage)
      .reshape([this.IMAGE_WIDTH, this.IMAGE_HEIGHT, 1])
      .toFloat();
    const xs = image.div(offset);
    const ys = tf.tensor(this.classes.map((x) => (x === label ? 1 : 0)));
    yield { xs, ys };
  }
}

构建了一个基于CNN的图像分类模型。通过tf.layers.conv3d()构造了卷积神经网络结构。

提供了train()方法，用于训练模型。这里定义了模型训练的迭代次数（epochs）、训练的批次大小（batchSize），这些参数会影响模型训练的最终结果，就是通常我们所说的“模型调参”，当你觉得模型训练效果不佳时，可以调整这些参数重新训练，直到达成不错的准确率。

提供了save()方法，用于保存模型参数。

import * as tf from "@tensorflow/tfjs-node";
import { DoodleData } from "./doodle-data.model";


exportclassClassifier {
  // ...
  // 定义模型结构
  constructor(data: DoodleData) {
    this.data = data;
    this.model = tf.sequential();
    this.model.add(
      tf.layers.conv2d({
        inputShape: [data.IMAGE_WIDTH, data.IMAGE_HEIGHT, 1],
        kernelSize: 3,
        filters: 16,
        strides: 1,
        activation: "relu",
        kernelInitializer: "varianceScaling",
      })
    );
    this.model.add(
      tf.layers.maxPooling2d({
        poolSize: [2, 2],
        strides: [2, 2],
      })
    );
    this.model.add(
      tf.layers.conv2d({
        kernelSize: 3,
        filters: 32,
        strides: 1,
        activation: "relu",
        kernelInitializer: "varianceScaling",
      })
    );
    this.model.add(
      tf.layers.maxPooling2d({
        poolSize: [2, 2],
        strides: [2, 2],
      })
    );
    this.model.add(tf.layers.flatten());
    this.model.add(
      tf.layers.dense({
        units: this.data.totalClasses,
        kernelInitializer: "varianceScaling",
        activation: "softmax",
      })
    );


    const optimizer = tf.train.adam();
    this.model.compile({
      optimizer,
      loss: "categoricalCrossentropy",
      metrics: ["accuracy"],
    });
  }


  // 模型训练
  async train(){
    const trainingData = tf.data
      .generator(() => this.data.dataGenerator("train"))
      .shuffle(this.data.maxImageClass * this.data.totalClasses)
      .batch(64);


    const testData = tf.data
      .generator(() => this.data.dataGenerator("test"))
      .shuffle(this.data.maxImageClass * this.data.totalClasses)
      .batch(64);


    await this.model.fitDataset(trainingData, {
      epochs: 5,
      validationData: testData,
      callbacks: {
        onEpochEnd: async (epoch, logs) => {
          this.logger.debug(
            `Epoch: ${epoch} - acc: ${logs?.acc.toFixed(
              3
            )} - loss: ${logs?.loss.toFixed(3)}`
          );
        },
        onBatchBegin: async (epoch, logs) => {
          console.log("onBatchBegin" + epoch + JSON.stringify(logs));
        },
      },
    });
  }


  // 保存模型
  async save(){
    fs.mkdirSync("doodle-model", { recursive: true });
    fs.writeFileSync(
      "doodle-model/classes.json",
      JSON.stringify({ classes: this.data.classes })
    );
    await this.model.save("file://./doodle-model");
  }
}

npm run start

onBatchBegin0{"batch":0,"size":512}
onBatchBegin1{"batch":1,"size":512}
onBatchBegin2{"batch":2,"size":512}
onBatchBegin3{"batch":3,"size":512}
onBatchBegin4{"batch":4,"size":192}
...
[Classifier] Epoch: 0 - acc: 0.078 - loss: 2.632
...

项目目录/
├── 📁 doodle-model/           # 训练结果（最终模型）
│   │   ├── 📄 classes.json    # 图片分类     
│   │   ├── 📄 model.json      # 模型描述文件
│   │   └── 📄 weights.bin     # 模型参数

这样，我们的模型就训练完成了。

接下来看看如何在页面中集成模型，实现从绘制canvas图片到模型分类预测的效果。

集成至页面

在页面中的集成模型也非常简单，我们只需要创建一个可以绘图的canvas，每隔一段时间就将当前canvas的图像数据传输给模型，触发一次模型预测即可。

先来看下项目的核心目录结构：

项目目录/
├── 📁 public/assets/doodle-modle/   # 将训练生成的模型放置在public目录下
│   │   ├── 📄 classes.json    # 图片分类     
│   │   ├── 📄 model.json      # 模型描述文件
│   │   └── 📄 weights.bin     # 模型参数
├── 📁 src   
│   ├── 📁 models/               
│   │   └── 📄 DoodleClassifier.js  # 图片分类器
│   ├── 📁 views/               
│   │   └── 📄 DoodleView.vue   # 页面视图（canvas画布）

• loadModel：加载模型，包括model.json、classes.json，在model.json中会自动加载weights.bin

• predictTopN：输入图片数据，调用model.predict() 预测最有可能的TopN个分类结果，并按照置信度排序

// 构造图片分类器
this.model = new DoodleClassifier()


// 加载模型
this.model.loadModel()


// 预处理canvas图片数据
const tensor = tf.browser.fromPixels(imgData, 1);
const resized = tf.image
  .resizeBilinear(tensor, [28, 28])
  .reshape([1, 28, 28, 1]) // Reshape to [1, 28, 28, 1] for batch and single channel
  .toFloat();
const normalized = tf.scalar(1.0).sub(resized.div(tf.scalar(255.0)));


// 预测图片分类
this.model.predictTopN(normalized, 5).then((predictions) => {
  if (predictions) {
    this.predictions = predictions;
  }
});

到这为止，你画我猜-AI版就已经基本搭建完成了。实现起来并不复杂。

如果一切顺利，并且你按照我们提供的demo构建页面，就可以直接在项目中运行：

npm run serve

优化措施

当我们去调整canvas画布大小、画笔粗细后，可能会出现预测结果不准确的情况，此时从canvas获取的图像数据和我们喂给模型的训练数据产生了差异。

这时候我们需要在获取到canvas数据后，额外做一些数据预处理，将数据标准化，例如：

• 将画布的内容区域裁剪为正方形，并居中显示

• 将画布的线条适当变粗，使模型更容易识别

总   结

参   考