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

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

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

tensorflow.js

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

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

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

本章使用的模式来自于

animegan2-pytorch

我们选用的是

face_paint_512_v2.pt

这个模型。

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

image-transformer

转化 pytorch 模型到 onnx 格式

参考资料

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

pytorch-fast-neural-style-for-web

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

AnimeGANv2-ONNX-Sample
- 转换脚本

建议看原文档

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 使用

python

js | pytorch 和 tensorflow 前端执行

转化 pytorch 模型到 onnx 格式

参考资料

python 使用

node 使用

onnxruntime-web