Как преобразовать веса модели tensorflow.js в тензоры pytorch и обратно?

Я использую 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?
        ####

Вот пример кода javascript для создания и загрузки 3 файлов в браузере. Для загрузки выберите в диалоговом окне сразу все 3 файла. Если они верны, во всплывающем окне будет показан пример прогноза.


person hazrmard    schedule 18.12.2020    source источник
comment
привет, возможно, это может вас заинтересовать datascience.stackexchange.com/questions/40496/ и эти преобразователи моделей   -  person IronMan    schedule 18.12.2020


Ответы (1)


Мне удалось найти способ конвертировать tfjs model.weights.bin в numpy ndarrays. Преобразовать массивы numpy в pytorch state_dict, который представляет собой словарь тензоров и их имен, тривиально.

Теория

Во-первых, следует понять представление модели в формате tfjs. model.json описывает модель. В питоне это можно читать как словарь. Он имеет следующие ключи:

  1. Архитектура модели описывается как еще один json / словарь под ключом modelTopology.

  2. У него также есть json / словарь под ключом weightsManifest, который описывает тип / форму / расположение каждого веса, заключенного в соответствующий model.weights.bin файл. Кроме того, манифест весов позволяет хранить веса в нескольких .bin файлах.

Tensorflow.js имеет сопутствующий пакет python с служебными функциями для и nofolights между двоичным форматом 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)
person hazrmard    schedule 19.12.2020