如何将 tensorflow.js 模型权重转换为 pytorch 张量，并返回？

Question

我正在使用 ml5.js，一个围绕 tensorflowjs 的包装器。我想在浏览器中训练一个神经网络，下载权重，在 pyTorch 中将它们处理为张量，然后将它们加载回浏览器的 tensorflowjs 模型中。如何在这些格式之间进行转换 tfjs <-> pytorch?

浏览器模型具有生成三个文件的save() 函数。特定于 ml5.js (json) 的元数据文件、描述模型架构的拓扑文件 (json) 和二进制权重文件 (bin)。

// Browser
model.save()

// HTTP/Download
model_meta.json   (needed by ml5.js)
model.json        (needed by tfjs)
model.weights.bin (needed by tfjs)

# python backend
import json

with open('model.weights.bin', 'rb') as weights_file:
    with open('model.json', 'rb') as model_file:
        weights = weights_file.read()
        model = json.loads(model_file.read())
        ####
        pytorch_tensor = convert2tensor(weights, model) # whats in this function?
        ####
        # Do some processing in pytorch

        ####
        new_weights_bin = convert2bin(pytorch_tensor, model) # and in this?
        ####

Here is sample javascript code 在浏览器中生成并加载 3 个文件。要在对话框中一次加载 select 所有 3 个文件。如果它们是正确的，弹出窗口将显示示例预测。

Answer 1

我找到了一种从 tfjs model.weights.bin 转换为 numpy 的 ndarrays 的方法。从 numpy 数组转换为 pytorch state_dict 是微不足道的，它是张量及其名称的字典。

理论

首先，应该理解模型的tfjs表示。 model.json 描述了模型。在python中，可以当字典来读。它具有以下键：

模型架构被描述为键 modelTopology 下的另一个 json/dictionary。
在键weightsManifest下还有一个json/dictionary，它描述了包裹在相应model.weights.bin文件中的每个权重的type/shape/location。另外，权重清单允许多个 .bin 文件存储权重。

Tensorflow.js 有一个同伴 python 包 tensorflowjs, which comes with utility functions to read and write 权重介于 tf.js 二进制和 numpy 数组格式之间。

每个权重文件都作为一个“组”读取。一个组是一个字典列表，其中的键 name 和 data 指的是权重名称和包含权重的 numpy 数组。还有可选的其他键。

group = [{'name': weight_name, 'data': np.ndarray}, ...]   # 1 *.bin file

申请

安装tensorflowjs。不幸的是，它还会安装tensorflow。

pip install tensorflowjs

使用这些功能。请注意，为了方便起见，我更改了签名。

from typing import Dict, ByteString
import torch
from tensorflowjs.read_weights import decode_weights
from tensorflowjs.write_weights import write_weights

def convert2tensor(weights: ByteString, model: Dict) -> Dict[str, torch.Tensor]:
    manifest = model['weightsManifest']
    # If flatten=False, returns a list of groups equal to the number of .bin files.
    # Use flatten=True to convert to a single group
    group = decode_weights(manifest, weights, flatten=True)
    # Convert dicts in tfjs group format into pytorch's state_dict format:
    # {name: str, data: ndarray} -> {name: tensor}
    state_dict = {d['name']: torch.from_numpy(d['data']) for d in group}
    return state_dict

def convert2bin(state_dict: Dict[str: np.ndarray], model: Dict, directory='./'):
    # convert state_dict to groups (list of 1 group)
    groups = [[{'name': key, 'data': value} for key, value in state_dict.items()]]
    # this library function will write to .bin file[s], but you can read it back
    # or change the function internals my copying them from source
    write_weights(groups, directory, write_manifest=False)

如何将 tensorflow.js 模型权重转换为 pytorch 张量，并返回？

How to convert tensorflow.js model weights to pytorch tensors, and back?

javascript

pytorch

tensorflow.js

ml5.js

理论

申请