0%

js | pytorch 和 tensorflow 前端执行

人工智能的计算量非常大,如果放在后端计算,需要的成本是巨大的,所以,应该尽可能的让模型在前端执行。

tensorflow 有官方支持的 tensorflow.js ,但是, pytorch 没有。

所以,使用 tensorflow 构建的模型可以天然在前端运行。

有一个开源项目叫做 onnx.js 可以将 torch 还有 tensorflow 的模型转化为 onnx 格式,然后,使用 onnx.js 进行前端执行。

当然还有其他可能的方案,在这里贴一下「之所以不用,是因为,我尝试过,但是总会出现这样那样的问题」

但是,由于生态混乱,onnx.js 在处理各种模型是总是出现这样那样的混乱,并且,其已经很长时间不更新了。但是,别无选择,这一章,就是使用 onnx.js 来运行 pytorch 模型。

本章使用的模式来自于

我们选用的是

  • face_paint_512_v2.pt

这个模型。

我们主要参考的开源项目是


转化 pytorch 模型到 onnx 格式


参考资料

上面的开源项目,转换 onnx 参考的是

但是这个太麻烦,我们直接参考

建议看原文档

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
import torch
from animegan2_pytorch.model import Generator

device = 'cpu'

torch.set_grad_enabled(False)

model = Generator().eval().to(device)
model.load_state_dict(torch.load(model_file_name))

def convert_to_onnx(net, file_name='output.onnx', input_shape=(512, 512), device='cpu'):
input = torch.randn(1, 3, input_shape[0], input_shape[1]).to(device)

input_layer_names = ['input_image']
output_layer_names = ['output_image']

torch.onnx.export(
net,
input,
file_name,
verbose=True,
opset_version=13,
input_names=input_layer_names,
output_names=output_layer_names,
)

side_length = 512
# ONNXファイル名
output_file_name = model_file_name.replace('512', str(side_length)).replace('.pt', '.onnx')

# ONNX変換
convert_to_onnx(model, output_file_name, (side_length, side_length), device)

python 使用



node 使用


onnxruntime-web

请我喝杯咖啡吧~