#!/usr/bin/env python3
# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
import caffe2.python.predictor.predictor_exporter as pe
import numpy as np
import torch
from caffe2.python import core, workspace
from caffe2.python.onnx import backend as caffe2_backend
from torch.onnx import ExportTypes, OperatorExportTypes
CAFFE2_DB_TYPE = "minidb"
[docs]def convert_caffe2_blob_name(blob_name):
# Caffe2 predictor appends the ":value" suffix to any feature that
# contains a string list
return f"{blob_name}_str:value"
[docs]def pytorch_to_caffe2(
model,
export_input,
external_input_names,
output_names,
export_path,
export_onnx_path=None,
):
num_tensors = 0
for inp in export_input:
num_tensors += len(inp) if isinstance(inp, (tuple, list)) else 1
assert len(external_input_names) == num_tensors
all_input_names = external_input_names[:]
for name, _ in model.named_parameters():
all_input_names.append(name)
# export the pytorch model to ONNX
if export_onnx_path:
print(f"Saving onnx model to: {export_onnx_path}")
else:
export_onnx_path = export_path
model.eval()
with torch.no_grad():
torch.onnx._export(
model,
export_input,
export_onnx_path,
input_names=all_input_names,
output_names=output_names,
export_params=True,
operator_export_type=OperatorExportTypes.ONNX_ATEN_FALLBACK,
export_type=ExportTypes.ZIP_ARCHIVE,
)
# Convert the ONNX model to a caffe2 net
c2_prepared = caffe2_backend.prepare_zip_archive(export_onnx_path)
return c2_prepared
[docs]def create_vocab_index(vocab_list, net, net_workspace, index_name):
vocab_index = net.StringIndexCreate([], index_name)
vocab_blob = net.AddExternalInput(net.NextName())
net_workspace.FeedBlob(str(vocab_blob), vocab_list)
# Populates the index with all the vocab
net.IndexGet([vocab_index, vocab_blob])
# Freeze the index to not add any other words in runtime
net.IndexFreeze([vocab_index], [vocab_index])
return vocab_index
[docs]def create_vocab_indices_map(c2_prepared, init_net, vocab_map):
vocab_indices = {}
for feat_name, vocab in vocab_map.items():
assert len(vocab) > 1
vocab_indices[feat_name] = create_vocab_index(
# Skip index 0 as it is reserved for unkwon tokens
# in Caffe2's index implementation
np.array(vocab[1:], dtype=str),
init_net,
c2_prepared.workspace,
feat_name + "_index",
)
return vocab_indices
[docs]def get_numericalize_net(c2_prepared, vocab_map, input_names):
predict_net = c2_prepared.predict_net
init_net = core.Net(c2_prepared.init_net)
final_input_names = input_names.copy()
with c2_prepared.workspace._ctx:
vocab_indices = create_vocab_indices_map(c2_prepared, init_net, vocab_map)
# Add operators to convert string features to ids based on the vocab
final_predict_net = core.Net(predict_net.name + "_processed")
final_inputs = set(
{
ext_input
for ext_input in input_names
if ext_input not in vocab_map.keys()
}
)
for ext_input in final_inputs:
final_predict_net.AddExternalInput(ext_input)
for feat in vocab_map.keys():
raw_input_blob = final_predict_net.AddExternalInput(
convert_caffe2_blob_name(feat)
)
# IndexGet expects flat tensors, so flatten the batch first then
# Resize it back after the lookup
flattened_input_blob = final_predict_net.FlattenToVec(raw_input_blob)
flattened_ids = final_predict_net.IndexGet(
[vocab_indices[feat], flattened_input_blob]
)
final_predict_net.ResizeLike([flattened_ids, raw_input_blob], [feat])
final_input_names[input_names.index(feat)] = convert_caffe2_blob_name(feat)
return final_predict_net, init_net, final_input_names
[docs]def add_feats_numericalize_ops(c2_prepared, vocab_map, input_names):
predict_net = c2_prepared.predict_net # Protobuf of the predict_net
final_predict_net, init_net, final_input_names = get_numericalize_net(
c2_prepared, vocab_map, input_names
)
with c2_prepared.workspace._ctx:
# Copy over the other list of the ops
final_predict_net.Proto().op.extend(predict_net.op)
# Update predict_net and init_net
c2_prepared.predict_net = final_predict_net.Proto()
c2_prepared.init_net = init_net.Proto()
return c2_prepared, final_input_names
[docs]def export_nets_to_predictor_file(
c2_prepared, input_names, output_names, predictor_path, extra_params=None
):
# netdef external_input includes internally produced blobs
actual_external_inputs = set()
produced = set()
for operator in c2_prepared.predict_net.op:
for blob in operator.input:
if blob not in produced:
actual_external_inputs.add(blob)
for blob in operator.output:
produced.add(blob)
for blob in output_names:
if blob not in produced:
actual_external_inputs.add(blob)
param_names = [
blob
for blob in actual_external_inputs
if blob not in input_names and blob not in output_names
]
if extra_params is not None:
param_names += extra_params
init_net = core.Net(c2_prepared.init_net)
predict_net = core.Net(c2_prepared.predict_net)
# Required because of https://github.com/pytorch/pytorch/pull/6456/files
with c2_prepared.workspace._ctx:
workspace.RunNetOnce(init_net)
predictor_export_meta = pe.PredictorExportMeta(
predict_net=predict_net,
parameters=param_names,
inputs=input_names,
outputs=output_names,
shapes={x: () for x in input_names + output_names},
net_type="simple",
)
pe.save_to_db(
db_type=CAFFE2_DB_TYPE,
db_destination=predictor_path,
predictor_export_meta=predictor_export_meta,
)
[docs]def validate_onnx_export(model):
pass