mirror of
https://github.com/saymrwulf/pytorch.git
synced 2026-05-14 20:57:59 +00:00
8286ce1e3a
Summary: Closes https://github.com/caffe2/caffe2/pull/1260 Differential Revision: D5906739 Pulled By: Yangqing fbshipit-source-id: e482ba9ba60b5337d9165f28f7ec68d4518a0902
87 lines
3.2 KiB
Python
87 lines
3.2 KiB
Python
# Copyright (c) 2016-present, Facebook, Inc.
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
##############################################################################
|
|
|
|
## @package mobile_exporter
|
|
# Module caffe2.python.mobile_exporter
|
|
|
|
from __future__ import absolute_import
|
|
from __future__ import division
|
|
from __future__ import print_function
|
|
from __future__ import unicode_literals
|
|
from caffe2.python import core, utils
|
|
from caffe2.proto import caffe2_pb2
|
|
|
|
|
|
def Export(workspace, net, params):
|
|
"""Returns init_net and predict_net suitable for writing to disk
|
|
and loading into a Predictor"""
|
|
proto = net if isinstance(net, caffe2_pb2.NetDef) else net.Proto()
|
|
predict_net = caffe2_pb2.NetDef()
|
|
predict_net.CopyFrom(proto)
|
|
init_net = caffe2_pb2.NetDef()
|
|
# Populate the init_net.
|
|
ssa, blob_versions = core.get_ssa(net)
|
|
inputs = []
|
|
for versioned_inputs, _ in ssa:
|
|
inputs += [name for name, _ in versioned_inputs]
|
|
|
|
input_blobs = [blob_name for blob_name, version in
|
|
blob_versions.items()
|
|
if version == 0 and blob_name not in params]
|
|
# Blobs that are never used as an input to another layer,
|
|
# i.e. strictly output blobs.
|
|
output_blobs = [blob_name for blob_name, version in
|
|
blob_versions.items()
|
|
if version != 0 and blob_name not in inputs]
|
|
|
|
for blob_ref in params:
|
|
blob_name = str(blob_ref)
|
|
blob = workspace.FetchBlob(blob_name)
|
|
init_net.op.extend(
|
|
[
|
|
core.CreateOperator(
|
|
"GivenTensorFill", [], [blob_name],
|
|
arg=[
|
|
utils.MakeArgument("shape", blob.shape),
|
|
utils.MakeArgument("values", blob)
|
|
]
|
|
)
|
|
]
|
|
)
|
|
# We have to make sure the blob exists in the namespace
|
|
# and we can do so with fake data. (Which is immediately overwritten
|
|
# by any typical usage)
|
|
for blob_name in input_blobs:
|
|
init_net.op.extend(
|
|
[
|
|
core.CreateOperator(
|
|
"GivenTensorFill", [], [blob_name],
|
|
arg=[
|
|
utils.MakeArgument("shape", [1, 1]),
|
|
utils.MakeArgument("values", [0.0])
|
|
]
|
|
)
|
|
]
|
|
)
|
|
|
|
# Now we make input/output_blobs line up with what Predictor expects.
|
|
del predict_net.external_input[:]
|
|
predict_net.external_input.extend(input_blobs)
|
|
# For populating weights
|
|
predict_net.external_input.extend(proto.external_input)
|
|
# Ensure the output is also consistent with what we want
|
|
del predict_net.external_output[:]
|
|
predict_net.external_output.extend(output_blobs)
|
|
return init_net, predict_net
|