mirror of
https://github.com/pytorch/pytorch.git
synced 2025-10-21 05:34:18 +08:00
4b68801c69
Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/62521 This diff did the following few things to enable the tests: 1. Exposed IMethod as TORCH_API. 2. Linked torch_deploy to test_api if USE_DEPLOY == 1. Test Plan: ./build/bin/test_api --gtest_filter=IMethodTest.* To be noted, one needs to run `python torch/csrc/deploy/example/generate_examples.py` before the above command. Reviewed By: ezyang Differential Revision: D30055372 Pulled By: alanwaketan fbshipit-source-id: 50eb3689cf84ed0f48be58cd109afcf61ecca508
58 lines
2.0 KiB
C++
58 lines
2.0 KiB
C++
// (c) Facebook, Inc. and its affiliates. Confidential and proprietary.
|
|
|
|
#include <gtest/gtest.h>
|
|
#include <torch/csrc/deploy/deploy.h>
|
|
#include <torch/script.h>
|
|
#include <torch/torch.h>
|
|
|
|
using namespace ::testing;
|
|
using namespace caffe2;
|
|
|
|
const char* simple = "torch/csrc/deploy/example/generated/simple";
|
|
const char* simpleJit = "torch/csrc/deploy/example/generated/simple_jit";
|
|
|
|
// TODO(jwtan): Try unifying cmake and buck for getting the path.
|
|
const char* path(const char* envname, const char* path) {
|
|
const char* env = getenv(envname);
|
|
return env ? env : path;
|
|
}
|
|
|
|
TEST(IMethodTest, CallMethod) {
|
|
auto scriptModel = torch::jit::load(path("SIMPLE_JIT", simpleJit));
|
|
auto scriptMethod = scriptModel.get_method("forward");
|
|
|
|
torch::deploy::InterpreterManager manager(3);
|
|
torch::deploy::Package package = manager.load_package(path("SIMPLE", simple));
|
|
auto pyModel = package.load_pickle("model", "model.pkl");
|
|
torch::deploy::PythonMethodWrapper pyMethod(pyModel, "forward");
|
|
|
|
auto input = torch::ones({10, 20});
|
|
auto outputPy = pyMethod({input});
|
|
auto outputScript = scriptMethod({input});
|
|
EXPECT_TRUE(outputPy.isTensor());
|
|
EXPECT_TRUE(outputScript.isTensor());
|
|
auto outputPyTensor = outputPy.toTensor();
|
|
auto outputScriptTensor = outputScript.toTensor();
|
|
|
|
EXPECT_TRUE(outputPyTensor.equal(outputScriptTensor));
|
|
EXPECT_EQ(outputPyTensor.numel(), 200);
|
|
}
|
|
|
|
TEST(IMethodTest, GetArgumentNames) {
|
|
auto scriptModel = torch::jit::load(path("SIMPLE_JIT", simpleJit));
|
|
auto scriptMethod = scriptModel.get_method("forward");
|
|
|
|
auto& scriptNames = scriptMethod.getArgumentNames();
|
|
EXPECT_EQ(scriptNames.size(), 1);
|
|
EXPECT_STREQ(scriptNames[0].c_str(), "input");
|
|
|
|
torch::deploy::InterpreterManager manager(3);
|
|
torch::deploy::Package package = manager.load_package(path("SIMPLE", simple));
|
|
auto pyModel = package.load_pickle("model", "model.pkl");
|
|
torch::deploy::PythonMethodWrapper pyMethod(pyModel, "forward");
|
|
|
|
auto& pyNames = pyMethod.getArgumentNames();
|
|
EXPECT_EQ(pyNames.size(), 1);
|
|
EXPECT_STREQ(pyNames[0].c_str(), "input");
|
|
}
|