随着深度学习技术的飞速发展，AI的魔力逐渐渗透到了我们日常生活的每一个角落，开启了无限的可能。特别是在前端开发领域，AI正如一股清流般，给传统的用户交互带来了全新的体验。想象一下，通过手势就能浏览网页，一笑便能播放你喜欢的音乐，或是仅凭一句语音命令，就能操控整个应用——这一切，都不再是科幻电影里的画面，而是即将成为我们日常生活的一部分。
不仅如此，AI的神奇力量还激发了前端产品创新的灵感源泉。无论是文本到语音转换（TTS）、图片分类，还是物体检测，这些看似高不可攀的技术，如今都触手可及。而且，随着技术边界的不断扩展，社区如Kaggle、Hugging Face等涌现出的高能模型更是让人目不暇接。如果作为前端开发者对这些技术视而不见，无疑是对自身能力的一大浪费。
此外，对于那些已经涉足或打算引入AI技术的公司来说，服务器成本无疑是一块巨大的负担。如果能将这些功能渗透到前端，那将大大降低成本，同时提升用户体验。
正因如此，作为一名前端开发者，掌握如何在浏览器中运用机器学习技术已经成为必修课。在本篇文章中，我将通过一系列代码示例，带大家了解如何在前端应用AI的三种方式：借助现成的库、使用TensorFlow.js以及OnnxRuntime-web来点亮你的前端产品，让AI成为你的超级伙伴。
以下内容请结合代码食用,github.com/ymrdf/web-a…
1 利用现成库
使用这种方法，开发者只需通过调用API接口即可将强大的AI能力嵌入到应用中。下面是一些好用的库:
@tensorflow-models/face-detection：专注于面部识别功能，可以识别出图片中的人脸及其关键特征点。@tensorflow-models/mobilenet：一种轻量级的模型，适用于图像分类和其他视觉任务，尤其优化了移动和低功耗设备的性能。@tensorflow-models/pose-detection：用于检测图像或视频中人体的姿势和关键点，支持多种姿势检测算法。
举个具体的例子，我们可以通过@tensorflow-models/face-detection库来实现面部检测功能。以下是一个简单的示例代码，供大家参考。更详细的实现方式和代码可以参见这个GitHub项目
async createDetector() {
try {
this.detector = await faceDetection.createDetector(
faceDetection.SupportedModels.MediaPipeFaceDetector,
{
runtime: "mediapipe",
modelType: "short",
maxFaces: 1,
solutionPath: this.options.solutionPath
? this.options.solutionPath
: `https://cdn.jsdelivr.net/npm/@mediapipe/face_detection@${mpFaceDetection.VERSION}`,
}
);
return;
} catch (e) {
console.warn(e);
this.setStatus(EUserDetectorStatus.faceDetectorCreateError);
}
}
使用这种方式的优点是易于上手, 对于不熟悉深度学习原理和模型训练的开发者来说，这种方法大大降低了技术门槛，使得只需少量的代码就能快速集成AI功能。但是目前能用的现成库很少，很难找到满足需求的库, 到npm和github上搜索到的大部分库都是在node端运行的库， 搜索到的库也质量不高。上面列举的三个库大部分是google给前端封装的tensorflow预训练的模型，质量可以保证。 其它可用的模型可以参考，github.com/tensorflow/…。另外transformersjs也不错, 但因为其模型都是用的transformer并且目前不能使用webGL, 在前端运行速度比较慢:www.npmjs.com/package/@xe…
2 使用tensorflowjs运行模型
TensorFlow.js j是google开发的用于在浏览器或 Node.js 上进行机器学习的库。它允许开发者直接在客户端中训练和部署机器学习模型，无需借助服务器端的计算资源，从而实现高效的数据处理和实时交互。
此外，TensorFlow.js 支持使用已有的 TensorFlow 模型，并将其转换为适用于 Web 的格式。开发者可以利用现有的预训练模型，也可以从零开始创建和训练新的模型，让机器学习变得更加便捷和可扩展。
TensorFlow.js 支持多种模型格式，包括 Keras HDF5、tf.keras SavedModel，以及 Kaggle 中的 TensorFlow Hub 模块。
2.1 用tensorflowjs运行tf.keras上已经预训练好的模型
这种方式可以支持的模型包 Keras HDF5, tf.keras SavedModel 另外包括 kaggle中的TensorFlow Hub module。现在可用的模型参考一下这里:keras.io/api/applica…
主要步聚是:
- 保存模型： 保存现有模型
- 转换模型： 使用TensorFlow.js的转换器将模型文件转换为TensorFlow.js支持的格式。
- 加载模型： 在网页应用中使用TensorFlow.js加载和运行该模型
现在我们看一下怎样运行tf.keras中的预训练过的模型,我从里面选了我比较喜欢的一个图片分类模型InceptionV3:
2.1.1 保存模型
首先，保存 tf.keras 中预训练好的模型, 保存为 Keras HDF5 格式：
import tensorflow as tf
from tensorflow.keras.preprocessing import image
from tensorflow.keras.applications.inception_v3 import preprocess_input, decode_predictions
from tensorflow.keras.layers import Input
input_tensor = Input(shape=(224, 224, 3))
model = tf.keras.applications.InceptionV3(input_tensor=input_tensor,weights='imagenet')
# 保存模型，保存为HDF5文件
model.save('inceptionv3.h5')
2.1.2 转换模型
2.1.2.1 安装并使用TensorFlow.js的转换器, 参考:
注意, 一定要新建一个python环境, 安装python 3.6.8, 然后执行:
pip install tensorflowjs[wizard]
如果出现如下错误,请升级pip:
Could not find a version that satisfies the requirement tensorflow<3,>=2.13.0 (from tensorflowjs[wizard]) (from versions: 0.12.1, 1.0.0, 1.1.0, 1.2.0, 1.2.1, 1.3.0, 1.4.0, 1.4.1, 1.5.0, 1.5.1, 1.6.0, 1.7.0, 1.7.1, 1.8.0, 1.9.0, 1.10.0, 1.10.1, 1.11.0, 1.12.0, 1.12.2, 1.12.3, 1.13.1, 1.13.2, 1.14.0, 1.15.0, 1.15.2, 1.15.3, 1.15.4, 1.15.5, 2.0.0, 2.0.1, 2.0.2, 2.0.3, 2.0.4, 2.1.0, 2.1.1, 2.1.2, 2.1.3, 2.1.4, 2.2.0, 2.2.1, 2.2.2, 2.2.3, 2.3.0, 2.3.1, 2.3.2, 2.3.3, 2.3.4, 2.4.0, 2.4.1, 2.4.2, 2.4.3, 2.4.4, 2.5.0, 2.5.1, 2.5.2, 2.6.0rc0, 2.6.0rc1, 2.6.0rc2, 2.6.0, 2.6.1, 2.6.2)
No matching distribution found for tensorflow<3,>=2.13.0 (from tensorflowjs[wizard])
如果出现证书问题:
Collecting tensorflowjs
Could not fetch URL https://pypi.python.org/simple/tensorflowjs/: There was a problem confirming the ssl certificate: [SSL: CERTIFICATE_VERIFY_FAILED] certificate verify failed (_ssl.c:852) - skipping
Could not find a version that satisfies the requirement tensorflowjs (from versions: )
No matching distribution found for tensorflowjs
请在pip命令后面加--trusted-host pypi.python.org --trusted-host files.pythonhosted.org --trusted-host pypi.org
2.1.2.2 转换模型：
执行:tensorflowjs_wizard根据提示将 .h5 文件转换为 TensorFlow.js 支持的格式。
2.1.3 加载模型：
在前端项目中，用如下代码加载并使用转换后的模型
async function loadModel() {
// 使用tf.loadLayersModel 或loadGraphModel加载转换后的模型
const model = await tf.loadGraphModel('/inceptionv3/model.json');
const result = model.predict(tf.zeros([1, IMAGE_SIZE, IMAGE_SIZE, 3])) as tf.Tensor
result.dispose();
return model;
}
2.1.4 使用模型：
用如下代码用模型进行预测, 首先，通过调用loadModel()函数异步加载预训练的模型，并将某一路径下的图片转化为张量格式，即imageTensor，以便进行模型预测。
随后，利用加载的模型对图像张量进行预测，拿到预测结果后，通过tf.squeeze()去除张量中的单一维度，然后应用tf.softmax()获取预测结果的概率分布。通过tf.topk()方法，提取出概率最高的5个分类及其对应的概率值，最终，使用imagenetClassesTopK()根据索引获取相对应的类别名称。
(代码参考:github.com/ymrdf/web-a…)
const model = await loadModel();
const imageTensor = await getImageTfTensorFromPath(path);
const predictions = model.predict(imageTensor) as tf.Tensor
const squeezed_tensor = tf.squeeze(predictions)
const outputSoftmax = tf.softmax(squeezed_tensor);
const top5 = tf.topk(outputSoftmax, 5);
const top5Indices = top5.indices.dataSync();
const top5Values = top5.values.dataSync();
const results = imagenetClassesTopK(top5Indices, top5Values);
return [results, 0.5];
下面是效果:
2.2 用tensorflowjs运行自己用tensorflow框架训练的模型
2.2.1 训练手写数字识别模型
我们先训练一个简单的手写数字识别模型。 因为这个任务是一个很简单的图片识别任务, 我们就选一个简单的卷积神经网络模型改一下就可以了。 我选LeNet(ieeexplore.ieee.org/document/72…), 它的结构差不多是这样的:
经过不断调整,把平均汇聚层改成最大汇聚层; 然后全连接层的激活函数改成relu函数, 卷积层和全连接层后面都加上批量规范化层,最后模型是这样的:
model = models.Sequential([
Input(shape=(28, 28, 1)), # 使用Input层明确指定输入形状
layers.Conv2D(6, kernel_size=(5, 5), padding="same", activation="sigmoid"),
layers.BatchNormalization(),
layers.MaxPooling2D(pool_size=(2, 2)),
layers.Conv2D(16, kernel_size=(5, 5), activation="sigmoid"),
layers.BatchNormalization(),
layers.MaxPooling2D(pool_size=(2, 2)),
layers.Flatten(),
layers.Dense(120, activation="relu"),
layers.BatchNormalization(),
layers.Dense(84, activation="relu"),
layers.BatchNormalization(),
layers.Dense(10, activation="softmax")
])
后用以下代码训练模型:
import tensorflow as tf
from tensorflow.keras import datasets, layers, models, Input
from tensorflow.keras.utils import to_categorical
# 1. 加载数据集
(training_images, training_labels), (test_images, test_labels) = datasets.mnist.load_data()
# 数据预处理
training_images = training_images.reshape((60000, 28, 28, 1)).astype("float32") / 255
test_images = test_images.reshape((10000, 28, 28, 1)).astype("float32") / 255
# 将标签转换为one-hot编码
training_labels = to_categorical(training_labels)
test_labels = to_categorical(test_labels)
# 2. 定义模型
model = models.Sequential([
Input(shape=(28, 28, 1)),
layers.Conv2D(6, kernel_size=(5, 5), padding="same", activation="sigmoid"),
layers.BatchNormalization(),
layers.MaxPooling2D(pool_size=(2, 2)),
layers.Conv2D(16, kernel_size=(5, 5), activation="sigmoid"),
layers.BatchNormalization(),
layers.MaxPooling2D(pool_size=(2, 2)),
layers.Flatten(),
layers.Dense(120, activation="relu"),
layers.BatchNormalization(),
layers.Dense(84, activation="relu"),
layers.BatchNormalization(),
layers.Dense(10, activation="softmax")
])
model.compile(optimizer=tf.keras.optimizers.SGD(learning_rate=0.05),
loss='categorical_crossentropy',
metrics=['accuracy'])
# 3. 训练模型
model.fit(training_images, training_labels, epochs=5, batch_size=64, verbose=2)
# 4. 评估模型
test_loss, test_acc = model.evaluate(test_images, test_labels, verbose=2)
print(f"\nTest accuracy: {test_acc*100:.2f}%, Test loss: {test_loss:.4f}")
调整超参数学习率到0.05, 在测试集上的预测准确率可达到98%; 在上面代码中增加如下代码,把模型保存为kera文件:
model.save("./numberRecog.h5")
2.2.2 转换模型
直接根据2.1.2节的办法转换模型, 执行tensorflowjs_wizard根据提示生成模型文件
2.2.3 加载模型
使用tf.loadLayersModel 或tf.loadGraphModel加载转换后的模型, 用哪个方法是根据转化模型时选择的格式来的,一般是用loadGraphModel
(代码参考: github.com/ymrdf/web-a…)
async function loadModel() {
const model = await tf.loadGraphModel('/numberRecogV1/model.json');
const result = model.predict(tf.zeros([1, IMAGE_SIZE, IMAGE_SIZE, 1])) as tf.Tensor
result.dispose();
return model;
}
2.2.4 使用模型
export async function inference(path:string) {
const model = await loadModel();
const imageTensor = await getImageTfTensorFromPath(path);
const predictions = model.predict(imageTensor) as tf.Tensor
const squeezed_tensor = tf.squeeze(predictions)
const outputSoftmax = tf.softmax(squeezed_tensor);
const top5 = tf.topk(outputSoftmax, 5);
const top5Indices = top5.indices.dataSync();
return [top5Indices, 0.5];
}
自己训练的模型果然差强人意,效果如下:
上面的方法只能加载TensorFlow训练的模型，不能运行pytorch等其它框架训练的模型。但是因为pytorch的强势，好多书和论文都是以pytorch做为事例的，网上能下载到的模型也是pytorch的远多于TensorFlow的。有没有办法能运行任意框架训练的模型呢？(也能把其它模型训练的参数转化成tensorflow模型参数, 再用tensorflowjs构建模型, 加载这些参数。 但这种方式门槛超高, 非常容易出错, 一般人弄不了)
3 使用onnxruntime-web运行模型
ONNX（Open Neural Network Exchange）是一种开放源代码的深度学习模型交换格式，由微软和Facebook共同开发。它旨在促进不同深度学习框架之间的互操作性，使模型可以在不同平台和工具之间无缝转换和运行。ONNX支持多种主流深度学习框架，如PyTorch、TensorFlow和Caffe2，简化了模型的部署流程，提高了开发效率。通过ONNX，开发者可以轻松地在不同环境中共享和复用深度学习模型，增强了人工智能项目的灵活性和可移植性。
ONNX Runtime Web 是个可以在浏览器里跑 ONNX 模型的工具。它让你在网页前端就能做深度学习推理，不需要靠后端服务器。通过使用 WebAssembly 和 WebGL，它在各种浏览器里都能高效运行。onnxruntime-web是在端上运行AI模型的通用方法, 只要这学会这个方法就能解决所有问题。
在 onnxruntime-web上运行模型的 主要步聚是:
- 保存模型： 保存现有模型为onnx文件
- 加载模型： 在网页应用中使用onnxruntime-web加载和运行该模型
- 运行模型
3.1 运行pytorch框架自带训练好的模型
3.1.1 导出模型
pytorch自带模型有若干，我选择resnet18模型演示, 其它模型可以参考:pytorch.org/vision/stab…
将PyTorch模型转换为ONNX格式。这可以通过使用PyTorch的torch.onnx.export函数来实现。主要过程是,加载预训练的resnet18模型
import torch
from torchvision import models, transforms
from PIL import Image
# 加载预训练的resnet18模型,设置模型为评估模式
resnet18 = models.resnet18(pretrained=True).eval()
transform = transforms.Compose([
transforms.CenterCrop(224),
transforms.ToTensor(),
transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
])
image = Image.open("./1.jpeg")
# 生成输入张量, 其它可以直接新建一个张量,这里用图片只是为了实验模型好不好用
image = transform(image).unsqueeze(0)
with torch.no_grad():
outputs = resnet18(image)
input_names = [ "actual_input_1" ]
output_names = [ "output1" ]
# 导出模型
torch.onnx.export(resnet18, image, "resnet18.onnx", verbose=False, input_names=input_names, output_names=output_names, opset_version=7)
3.1.2 用onnxruntime-web加载模型
onnxruntime-web的InferenceSession.create()是一个异步方法，用来加载一个ONNX模型文件，该文件在这里是以第一个参数的路径指定的。
在.create()方法中，传入了一个配置对象，该对象有两个字段：
executionProviders: ['webgl']：字段指定了模型运行时的后端执行器。在这个例子中，它设置为使用WebGL来加速计算，这意味着模型的推理将利用webGL, 也可选择cpu, wasm.graphOptimizationLevel: 'all'：字段指定了图优化级别。设置为'all'意味着会启用所有可用的图形优化，以提高模型执行的效率和性能。
(代码参考: github.com/ymrdf/web-a…)
const session = await ort.InferenceSession
.create('/resnet18.onnx',
{ executionProviders: ['webgl'], graphOptimizationLevel: 'all' });
3.1.3 用onnxruntime-web运行模型
首先创建了一个空的输入数据对象feeds，并根据模型的输入名称将预处理过的数据preprocessedData作为输入。通过异步执行session.run()方法运行模型. 接下来，从输出数据中获取模型的预测结果，应用softmax函数处理这些结果以获得概率分布，然后使用tf.topk()方法找出概率最高的五个预测结果及其索引。最后，通过imagenetClassesTopK方法将这些索引转换为具体的类别名称，并将这些信息以及推理时间一起返回。
const start = new Date();
const feeds: Record<string, ort.Tensor> = {};
feeds[session.inputNames[0]] = preprocessedData;
const outputData = await session.run(feeds);
const end = new Date();
const inferenceTime = (end.getTime() - start.getTime())/1000;
const output = outputData[session.outputNames[0]];
const outputSoftmax = tf.softmax(tf.tensor(Array.prototype.slice.call(output.data)));
const top5 = tf.topk(outputSoftmax, 5);
const top5Indices = top5.indices.dataSync();
const top5Values = top5.values.dataSync();
const results = imagenetClassesTopK(top5Indices, top5Values);
return [results, inferenceTime];
下面是效果:
3.2 到huggingface和kaggle等平台找训练好的模型
实际上，上述介绍的所有方法虽然实用，但由于框架自带或者用 TensorFlow 预训练的模型相对较少，自行训练高质量模型的难度又相对较大，其实际应用场景可能有限。然而，如果我们能够运行 Hugging Face 和 Kaggle 等平台上的任意模型，应用的空间将会大大拓展。
例如，我随便在 Hugging Face 上找了一个模型：huggingface.co/briaai/RMBG… 这个模型非常强大，能够精准地分割出图片中的物体。
让我们直观地看下其效果：
现在我们就要在浏览器上实现这个功能,是不是感觉有点小兴奋???
3.2.1 准备
首先根据huggingface.co/briaai/RMBG…上的介绍看看怎么加载这个模型:
from transformers import AutoModelForImageSegmentation
model = AutoModelForImageSegmentation.from_pretrained("briaai/RMBG-1.4",trust_remote_code=True)
然后写代码试试这个模型好不好用:
from transformers import AutoModelForImageSegmentation
from torchvision.transforms.functional import normalize
import torch.nn.functional as F
import numpy as np
import torch
from skimage import io
from PIL import Image
model = AutoModelForImageSegmentation.from_pretrained("briaai/RMBG-1.4",trust_remote_code=True)
def preprocess_image(im: np.ndarray, model_input_size: list) -> torch.Tensor:
if len(im.shape) < 3:
im = im[:, :, np.newaxis]
# orig_im_size=im.shape[0:2]
im_tensor = torch.tensor(im, dtype=torch.float32).permute(2,0,1)
im_tensor = F.interpolate(torch.unsqueeze(im_tensor,0), size=model_input_size, mode='bilinear')
image = torch.divide(im_tensor,255.0)
image = normalize(image,[0.5,0.5,0.5],[1.0,1.0,1.0])
return image
def postprocess_image(result: torch.Tensor, im_size: list)-> np.ndarray:
result = torch.squeeze(F.interpolate(result, size=im_size, mode='bilinear') ,0)
ma = torch.max(result)
mi = torch.min(result)
result = (result-mi)/(ma-mi)
im_array = (result*255).permute(1,2,0).cpu().data.numpy().astype(np.uint8)
im_array = np.squeeze(im_array)
return im_array
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
model_input_size = [1024,1024]
image_path = "./profile.jpg"
orig_im = io.imread(image_path)
orig_im_size = orig_im.shape[0:2]
image = preprocess_image(orig_im, model_input_size).to(device)
model.eval()
result=model(image)
result_image = postprocess_image(result[0][0], orig_im_size)
pil_im = Image.fromarray(result_image)
no_bg_image = Image.new("RGBA", pil_im.size, (0,0,0,0))
orig_image = Image.open(image_path)
no_bg_image.paste(orig_image, mask=pil_im)
no_bg_image.show()
运行后能成功分割出图片中的物体。
3.2.2 导出模型onnx文件
然后给上面代码加上这么三行代码再执行,生成onnx (这个模型用了较新的算子,导不出7版本的onnx, 只能用最新版本的onnx 到时候用cpu推理了:
input_names = [ "actual_input_1" ]
output_names = [ "output1" ]
torch.onnx.export(model, image, "rmbg.onnx", verbose=False, input_names=input_names, output_names=output_names)
3.2.3 用onnxruntime-web加载模型
(代码参考: github.com/ymrdf/web-a…; 运行代码前请先解压public/rmbg.onnx文件)
const session = await ort.InferenceSession
.create('/rmbg.onnx',
{ executionProviders: ['cpu'], graphOptimizationLevel: 'all' });
3.2.4 推理
async function runInference(session: ort.InferenceSession, preprocessedData: any): Promise<any> {
const feeds: Record<string, ort.Tensor> = {};
feeds[session.inputNames[0]] = preprocessedData;
const outputData = await session.run(feeds);
const output = outputData[session.outputNames[0]];
return output
}
3.2.5 数据预处理
看了上面代码后你可能吃惊于用模型推理的代码竟如此简单,其实在实际用模型推理的时候, 最麻烦的是数据的预处理, 比如这个模型需要的输入形如[1,3,1024,1024]的张量。 所以要获取一张图片并把它处理成[1,3,1024,1024]的张量。但是因为onnxruntime-web提供的张量运算方法很少,我一般是用tensorflowjs提供的张量进行数据处理,最后把tensorflow张量转化成onnxruntime-web张量:
这里就这个例子说明数据的处理流程, 步骤如下7个步骤:
export async function getImageTfTensorFromPath(path: string ): Promise<tf.Tensor> {
return new Promise((r) => {
const src = path
const $image = new Image();
$image.crossOrigin = 'Anonymous';
$image.onload = function() {
// 1. 将图片元素转换为Tensor
const tensor = tf.browser.fromPixels($image)
.resizeBilinear([1024,1024]) // 2. 更改图片大小
.toFloat() // 3. 转化成浮点数
.div(tf.scalar(255.0)) // 4. 归一化
.transpose([2, 0, 1]) // 5. 把[1024,1024,3]的张量转成[3,1024,1024]
.expandDims(); // 6. 增加一维,使其变成[1,3,1024,1024]:
// 7. 标准化张量
const mean = tf.tensor([0.5,0.5,0.5]);
const std = tf.tensor([1.0,1.0,1.0]);
const normalizedTensor = tensor.sub(mean.reshape([1,3,1,1])).div(std.reshape([1,3,1,1]));
normalizedTensor.print()
r(normalizedTensor);
};
$image.src = src;
})
}
最后是把tensorflow张量转化成onnxruntime-web张量。 转化函数在源码中,大家可以自取就不列了。
3.2.6 数据后处理
模型生成的数据是一个形如[1,1,1024,1024]张量, 每个数据是0到1的浮点数, 代表对应图片位置上的像素点是否保留。
怎么把这个数据变成一张图片呢,方法如下:
// 1. 将图片元素转换为Tensor
const tensor = tf.browser.fromPixels($image).resizeBilinear([1024,1024])
// 2. 将输出数据转换为tf.Tensor
const alpha4 = convertOnnxTensorToTfTensor(alphaExpanded)
// 3. 去掉一维
let alpha3 = tf.squeeze(alpha4, [1]);
// [1,1024,1024] => [1024,1024, 1]
alpha3 = tf.reshape(alpha3, [1024, 1024, 1]);
// 乘255
alpha3 = tf.mul(alpha3, 255)
combine(tensor, alpha3)
function combine(imageTensor:tf.Tensor, alphaExpanded:tf.Tensor){
// 1.沿最后一个维度合并
const combinedTensor = tf.concat([imageTensor, alphaExpanded], -1);
// 2.将tensor转换为Uint8ClampedArray
combinedTensor.data().then(data => {
const clampedArray = new Uint8ClampedArray(data);
// 3. 创建ImageData对象
const imageData = new ImageData(clampedArray, 1024, 1024);
// 4. 绘制到canvas
const canvas = document.querySelector("#test");
const ctx = canvas.getContext('2d');
ctx?.putImageData(imageData, 0, 0);
});
}
把下面图片扔进去:
运行结果如下,效果还不错:
3.3 运行自己训练的模型
3.3.1 用pytorch训练手写数字识别模型
虽然上面介绍了怎样用tensorflowjs运行自己训练的模型，但前面说过了大部分书都是以pytorch框架为教学框架的，对前端开发工程师来说，一般先学习的都是pytorch。所以我们还是再学一遍怎么把用pytorch框架训练的模型加载到web环境下。 这次我们还是训练手写数字识别模型，这次我们用pytorch实现上面tensorflow的模型。
import torch
from torch import nn
from torch.utils.data import DataLoader
from torchvision import datasets
from torchvision.transforms import ToTensor
training_data = datasets.MNIST(
root="data",
train=True,
download=True,
transform=ToTensor(),
)
test_data = datasets.MNIST(
root="data",
train=False,
download=True,
transform=ToTensor(),
)
batch_size = 64
train_dataloader = DataLoader(training_data, batch_size=batch_size)
test_dataloader = DataLoader(test_data, batch_size=batch_size)
for X, y in test_dataloader:
print(f"Shape of X [N, C, H, W]: {X.shape}")
print(f"Shape of y: {y.shape} {y.dtype}")
break
device = (
"cuda"
if torch.cuda.is_available()
else "mps"
if torch.backends.mps.is_available()
else "cpu"
)
class NeuralNetwork(nn.Module):
def __init__(self):
super().__init__()
self.flatten = nn.Flatten()
self.linear_relu_stack = nn.Sequential(
nn.Conv2d(1, 6, kernel_size=5, padding=2),nn.BatchNorm2d(6), nn.Sigmoid(),
nn.MaxPool2d(kernel_size=2, stride=2),
nn.Conv2d(6, 16, kernel_size=5),nn.BatchNorm2d(16), nn.Sigmoid(),
nn.MaxPool2d(kernel_size=2, stride=2),
nn.Flatten(),
nn.Linear(16 * 5 * 5, 120),nn.BatchNorm1d(120), nn.ReLU(),
nn.Linear(120, 84),nn.BatchNorm1d(84), nn.ReLU(),
nn.Linear(84, 10)
)
def forward(self, x):
logits = self.linear_relu_stack(x)
return logits
model = NeuralNetwork().to(device)
print(model)
def init_weights(m):
if type(m) == nn.Linear or type(m) == nn.Conv2d:
nn.init.xavier_uniform_(m.weight)
loss_fn = nn.CrossEntropyLoss()
optimizer = torch.optim.SGD(model.parameters(), lr=0.05) # 0.1-97.7 0.05-96.5 0.075-97.0
def train(dataloader, model, loss_fn, optimizer):
size = len(dataloader.dataset)
model.train()
for batch, (X, y) in enumerate(dataloader):
X, y = X.to(device), y.to(device)
pred = model(X)
loss = loss_fn(pred, y)
loss.backward()
optimizer.step()
optimizer.zero_grad()
if batch % 100 == 0:
loss, current = loss.item(), (batch + 1) * len(X)
print(f"loss: {loss:>7f} [{current:>5d}/{size:>5d}]")
def test(dataloader, model, loss_fn):
size = len(dataloader.dataset)
num_batches = len(dataloader)
model.eval()
test_loss, correct = 0, 0
with torch.no_grad():
for X, y in dataloader:
X, y = X.to(device), y.to(device)
pred = model(X)
test_loss += loss_fn(pred, y).item()
correct += (pred.argmax(1) == y).type(torch.float).sum().item()
test_loss /= num_batches
correct /= size
print(f"Test Error: \n Accuracy: {(100*correct):>0.1f}%, Avg loss: {test_loss:>8f} \n")
epochs = 5
model.linear_relu_stack.apply(init_weights)
for t in range(epochs):
print(f"Epoch {t+1}\n-------------------------------")
train(train_dataloader, model, loss_fn, optimizer)
test(test_dataloader, model, loss_fn)
3.3.2 导出模型为onnx文件
调参数满意后,运行以下代码保存模型，生成onnx文件:
input = torch.rand(size=(1,1, 28, 28), dtype=torch.float32)
torch.onnx.export(model, input , "numberRecog.onnx", verbose=False, opset_version=7)
3.3.3 加载模型,并运行模型
(代码参考: github.com/ymrdf/web-a…)
export async function inference(path:string):Promise<[Uint8Array,number]> {
const imageTensor = await getImageTfTensorFromPath(path);
const preprocessedData = await convertTfTensorToOnnxTensor(imageTensor);
const session = await ort.InferenceSession
.create('/numberRecog.onnx',
{ executionProviders: ['cpu'], graphOptimizationLevel: 'all' });
const start = new Date();
const feeds: Record<string, ort.Tensor> = {};
feeds[session.inputNames[0]] = preprocessedData;
const outputData = await session.run(feeds);
const end = new Date();
const inferenceTime = (end.getTime() - start.getTime())/1000;
const output = outputData[session.outputNames[0]];
const predictions = convertOnnxTensorToTfTensor(output);
const squeezed_tensor = tf.squeeze(predictions)
const outputSoftmax = tf.softmax(squeezed_tensor);
const top5 = tf.topk(outputSoftmax, 5);
const top5Indices = top5.indices.dataSync() as Uint8Array;
return [top5Indices, inferenceTime];
}
效果和tensorflow训练的模型一样如下:
onnxruntime-web也有一些局限性，主要是onnxruntime-web在使用webgl时支持的onnx版本比较老，有些算子可能无法支持。使用assembly可以支持所有算子,但速度相对较慢。 另外, onnx文件不能方便的被分割。
总结
在本文中，我们详细介绍了如何在前端开发中运行机器学习模型，包括使用TensorFlow.js和OnnxRuntime-web这两种主要方法。首先，我们需要将模型保存或转换成相应的格式，然后使用对应的工具进行加载和执行。
TensorFlow.js能够直接在浏览器中运行TensorFlow框架的模型，为开发者提供了一种便捷的模型训练和部署方式。然而，现有的TensorFlow.js资源相对有限，且其转换工具tensorflowjs_wizard使用上具有一定难度。
相比之下，OnnxRuntime-web支持加载使用多种框架训练的ONNX模型，提供了更大的灵活性。然而，由于WebGL的限制，有些算子可能无法支持，需要依赖Assembly来确保模型的兼容性，但会带来执行速度上的折扣。
对于数据处理部分，onnxRuntime-web的张量计算方法相对较少。我们可以采用@tensorflow/tfjs的方法来进行数据处理，然后将处理后的结果转换为onnxRuntime-web所支持的张量对象进行进一步应用。
通过学习和掌握这些技术，我们可以将AI技术无缝集成到前端开发中，将网页应用的用户体验提升到一个全新的高度。无论是作为前端开发者还是AI爱好者，继续探索和学习AI相关技术，无疑是打开未来技术大门的一把钥匙。让我们一起拥抱AI技术，勇敢探索未知的前沿领域，为用户构建更智能、更人性化的应用！
希望这篇文章能为你提供有益的帮助，愿我们在前端开发与AI技术的结合之路上共同成长，创造出更多可能性！
