mirror of
https://github.com/zebrajr/pytorch.git
synced 2025-12-07 00:21:07 +01:00
a25e6370e5
Summary: Expose IMethod interface, which provides a unified interface to either script or python methods backed by torchscript or torchdeploy. IMethod provides a way to depend on a torch method without depending on a particular runtime implementation such as torchscript or python/deploy. Test Plan: add unit tests. Reviewed By: suo Differential Revision: D29463455 fbshipit-source-id: 903391d9af9fbdd8fcdb096c1a136ec6ac153b7c
45 lines
1.6 KiB
C++
45 lines
1.6 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;
|
|
|
|
TEST(IMethodTest, CallMethod) {
|
|
auto script_model = torch::jit::load(getenv("SIMPLE_JIT"));
|
|
auto script_method = script_model.get_method("forward");
|
|
|
|
torch::deploy::InterpreterManager manager(3);
|
|
torch::deploy::Package p = manager.load_package(getenv("SIMPLE"));
|
|
auto py_model = p.load_pickle("model", "model.pkl");
|
|
torch::deploy::PythonMethodWrapper py_method(py_model, "forward");
|
|
|
|
auto input = torch::ones({10, 20});
|
|
auto output_py = py_method({input});
|
|
auto output_script = script_method({input});
|
|
EXPECT_TRUE(output_py.isTensor());
|
|
EXPECT_TRUE(output_script.isTensor());
|
|
auto output_py_tensor = output_py.toTensor();
|
|
auto output_script_tensor = output_script.toTensor();
|
|
|
|
EXPECT_TRUE(output_py_tensor.equal(output_script_tensor));
|
|
EXPECT_EQ(output_py_tensor.numel(), 200);
|
|
}
|
|
|
|
TEST(IMethodTest, GetArgumentNames) {
|
|
auto script_model = torch::jit::load(getenv("SIMPLE_JIT"));
|
|
auto script_method = script_model.get_method("forward");
|
|
|
|
torch::deploy::InterpreterManager manager(3);
|
|
torch::deploy::Package p = manager.load_package(getenv("SIMPLE"));
|
|
auto py_model = p.load_pickle("model", "model.pkl");
|
|
torch::deploy::PythonMethodWrapper py_method(py_model, "forward");
|
|
|
|
// TODO(whc) implement and test these
|
|
EXPECT_THROW(script_method.getArgumentNames(), std::runtime_error);
|
|
EXPECT_THROW(py_method.getArgumentNames(), std::runtime_error);
|
|
}
|