Skip to content
Permalink
Browse files

Preserve output PH names in C2 loader (#3798)

Summary:
It's important to preserve the name of each external output in outside models, but since they must be unique with Node names we'll need to claim them before naming any nodes. This fixes the C2 Model Loader to do that.

Documentation: Fixes #3774
Pull Request resolved: #3798

Test Plan: ninja test, image-classifier with -dump-ir to make sure weights were named right.

Differential Revision: D18630019

Pulled By: nickgg

fbshipit-source-id: 0b0e771c0e6489384f7e40a772178bf45ff64b0e
  • Loading branch information
nickgg authored and facebook-github-bot committed Dec 2, 2019
1 parent 5beebe1 commit a4f2540aafe40a29aaa63c03745df98c687ab222
@@ -222,7 +222,7 @@ uint8_t activations[LENET_MNIST_ACTIVATIONS_MEM_SIZE];
uint8_t *inputAddr = GLOW_GET_ADDR(mutableWeight, LENET_MNIST_data);

/// Bundle output data absolute address.
uint8_t *outputAddr = GLOW_GET_ADDR(mutableWeight, LENET_MNIST_softmax__1);
uint8_t *outputAddr = GLOW_GET_ADDR(mutableWeight, LENET_MNIST_softmax);

/// Copy the pre-processed images into the mutable region of the bundle.
static void initInputImages() {
@@ -300,7 +300,7 @@ static uint8_t *allocateMutableWeightVars(const BundleConfig &config) {
static void dumpInferenceResults(const BundleConfig &config,
uint8_t *mutableWeightVars) {
const SymbolTableEntry &outputWeights =
getMutableWeightVar(config, "gpu_0_softmax__1");
getMutableWeightVar(config, "gpu_0_softmax");
int maxIdx = 0;
float maxValue = 0;
float *results = (float *)(mutableWeightVars + outputWeights.offset);
@@ -101,7 +101,7 @@ void dispatchClassify(unsigned int id, HostManager *hostManager,
EXIT_ON_ERR(std::move(err));
auto *bindings = context->getPlaceholderBindings();
size_t maxIdx =
bindings->get(bindings->getPlaceholderByName("gpu_0_softmax__1"))
bindings->get(bindings->getPlaceholderByName("gpu_0_softmax"))
->getHandle()
.minMaxArg()
.second;
@@ -109,6 +109,12 @@ class Module final {
usedNodeNames_.insert(name);
}

/// Registers a name as used by a Storage node (Constant or Placeholder) in
/// this module.
void registerStorageName(llvm::StringRef name) {
usedStorageNames_.insert(name);
}

/// Return a pointer to a uniqued type \p T.
TypeRef uniqueType(const Type &T);

@@ -1444,6 +1444,12 @@ Error Caffe2ModelLoader::loadInputs(const caffe2::NetDef &net,
}

Error Caffe2ModelLoader::loadNetwork(caffe2::NetDef &net) {
// Make a claim on the unique name of all output Placeholders.
for (int i = 0; i < net.external_output_size(); i++) {
auto &outputName = net.external_output(i);
G_.getParent()->registerStorageName(legalizeName(outputName));
}

/// Load the network operators:
for (int i = 0; i < net.op_size(); i++) {
auto &op = net.op(i);
@@ -1464,8 +1470,13 @@ Error Caffe2ModelLoader::loadNetwork(caffe2::NetDef &net) {
auto &outputName = net.external_output(i);
NodeValue r;
ASSIGN_VALUE_OR_RETURN_ERR(r, getNodeValueByName(outputName));
auto *PH =
G_.getParent()->createPlaceholder(r.getType(), outputName, false);

PlaceholderList &PHList = G_.getParent()->getPlaceholders();
// Create a Placeholder with the previously claimed name.
auto *PH = new Placeholder(legalizeName(outputName),
G_.getParent()->uniqueType(*r.getType()), false,
ANY_LAYOUT);
PHList.push_back(PH);
auto *SN = G_.createSave("save_" + outputName, r, PH);
outputVarsByName_[outputName] = SN->getPlaceholder();
}
@@ -2551,7 +2551,8 @@ TEST(caffe2, importNames) {
Tensor input(ElemKind::FloatTy, {6});
Caffe2ModelLoader caffe2LD(NetDescFilename, NetWeightFilename,
{"sigmoid_test_input"}, {&input.getType()}, *F);
EXPECT_TRUE(F->getNodeByName("sigmoid_test_output"));
EXPECT_TRUE(mod.getPlaceholderByName("sigmoid_test_output"));
EXPECT_TRUE(F->getNodeByName("sigmoid_test_output__1"));
}

TEST(caffe2, importSqr) {

0 comments on commit a4f2540

Please sign in to comment.
You can’t perform that action at this time.