Tensorflow.js项目实战(7):手写数字识别1

import * as tf from"@tensorflow/tfjs";

// 定义超参数
// 学习率：模型在训练过程中更新权重的速度
const LEARNING_RATE = 0.1;
// 批大小：每次训练使用的样本数
const BATCH_SIZE = 64;
// 总训练步骤数
const TRAIN_STEPS = 400;

// 数据相关的常量
// 图像尺寸，这里假设是28x28像素
const IMAGE_SIZE = 28;
// 类别数目，对于MNIST数据集是10个数字
const LABELS_SIZE = 10;
// 初始化优化器，采用随机梯度下降法(SGD)，学习率为0.1
const optimizer = tf.train.sgd(LEARNING_RATE);

// 创建一个顺序模型（Sequential model）
// 这是一种可以按顺序添加多层的模型结构
const model = tf.sequential();

// 添加第一个卷积层
// 卷积层参数说明：
// - inputShape: 输入数据形状，这里是[28, 28, 1]，即28x28像素的单通道图像
// - kernelSize: 卷积核大小为5x5
// - filters: 输出通道数为8
// - strides: 卷积步长为1，意味着卷积核每次在图像上滑动时，水平和垂直方向都移动1个像素
// - activation: 激活函数选用ReLU
// - kernelInitializer: 初始化权重的方法为变方差缩放初始化
model.add(tf.layers.conv2d({
  inputShape: [IMAGE_SIZE, IMAGE_SIZE, 1],
  kernelSize: 5,
  filters: 8,
  strides: 1,//解释一下strides的作用，就是卷积核的移动步长，步长为1的话，卷积核每次只移动一个像素点，步长为2的话，卷积核每次移动两个像素点，以此类推。
  activation: 'relu',
  kernelInitializer: 'varianceScaling'
}));
// 添加第一个最大池化层
// 参数poolSize指定池化窗口大小为2x2
// strides也为2，这意味着每次池化后图像尺寸会减半
model.add(tf.layers.maxPooling2d({ poolSize: 2, strides: 2 }));

// 添加第二个卷积层，结构类似第一个卷积层但通道数更多
model.add(tf.layers.conv2d({
  kernelSize: 5,
  filters: 16,
  strides: 1,
  activation: 'relu',
  kernelInitializer: 'varianceScaling'
}));
// 添加第二个最大池化层，同样使图像尺寸减半
model.add(tf.layers.maxPooling2d({ poolSize: 2, strides: 2 }));
// 将多维的卷积特征展平成一维向量
model.add(tf.layers.flatten());

// 添加全连接层（Dense layer），输出节点数等于类别数
// 使用softmax作为激活函数，使得输出可以表示每个类别的概率分布
model.add(tf.layers.dense({
  units: LABELS_SIZE,
  kernelInitializer: 'varianceScaling',
  activation: 'softmax'
}));

// 训练模型
exportasyncfunction train(data) {
   // 编译模型，设置优化器和损失函数
  model.compile({
    optimizer: optimizer,
    loss: 'categoricalCrossentropy'
  });

   // 开始训练循环
  for (let i = 0; i < TRAIN_STEPS; i++) {
    const batch = data.nextTrainBatch(BATCH_SIZE);
    //reshape -1,28,28,1  第一个参数-1 表示不确定
    const history = await model.fit(batch.xs.reshape([-1, IMAGE_SIZE, IMAGE_SIZE, 1]), batch.labels, {
      batchSize: BATCH_SIZE,
      epochs: 1
    });
    // 打印当前批次的平均损失值以及训练步数
    console.log(history.history.loss[0], i);
    await tf.nextFrame();
  }
  //保存模型
  model.save('localstorage://mnist-model');
}

// 预测模型
exportfunction predict(x) {
  const predictions = model.predict(x.reshape([-1, IMAGE_SIZE, IMAGE_SIZE, 1]));
  returnArray.from(predictions.argMax(-1).dataSync());
}

// 辅助函数，用于从标签向量中获取类别索引
exportfunction classesFromLabel(y) {
  returnArray.from(y.argMax(-1).dataSync());
}