mirror of
https://github.com/zebrajr/pytorch.git
synced 2025-12-06 12:20:52 +01:00
33c03cb61a
Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/65861 First in a series. This PR changes the code in deploy.h/cpp and interpreter_impl.h/cpp to be camel case instead of snake case. Starting with this as it has the most impact on downstream users. Test Plan: Imported from OSS Reviewed By: shannonzhu Differential Revision: D31291183 Pulled By: suo fbshipit-source-id: ba6f74042947c9a08fb9cb3ad7276d8dbb5b2934
87 lines
2.5 KiB
C++
87 lines
2.5 KiB
C++
#include <gtest/gtest.h>
|
|
#include <torch/csrc/deploy/deploy.h>
|
|
#include <torch/cuda.h>
|
|
#include <torch/script.h>
|
|
#include <torch/torch.h>
|
|
#include <future>
|
|
#include <iostream>
|
|
#include <string>
|
|
|
|
int main(int argc, char* argv[]) {
|
|
::testing::InitGoogleTest(&argc, argv);
|
|
int rc = RUN_ALL_TESTS();
|
|
return rc;
|
|
}
|
|
|
|
const char* simple = "torch/csrc/deploy/example/generated/simple";
|
|
const char* simple_jit = "torch/csrc/deploy/example/generated/simple_jit";
|
|
|
|
const char* path(const char* envname, const char* path) {
|
|
const char* e = getenv(envname);
|
|
return e ? e : path;
|
|
}
|
|
|
|
TEST(TorchDeployGPUTest, SimpleModel) {
|
|
if (!torch::cuda::is_available()) {
|
|
GTEST_SKIP();
|
|
}
|
|
const char* model_filename = path("SIMPLE", simple);
|
|
const char* jit_filename = path("SIMPLE_JIT", simple_jit);
|
|
|
|
// Test
|
|
torch::deploy::InterpreterManager m(1);
|
|
torch::deploy::Package p = m.loadPackage(model_filename);
|
|
auto model = p.loadPickle("model", "model.pkl");
|
|
{
|
|
auto M = model.acquireSession();
|
|
M.self.attr("to")({"cuda"});
|
|
}
|
|
std::vector<at::IValue> inputs;
|
|
{
|
|
auto I = p.acquireSession();
|
|
auto eg = I.self.attr("load_pickle")({"model", "example.pkl"}).toIValue();
|
|
inputs = eg.toTuple()->elements();
|
|
inputs[0] = inputs[0].toTensor().to("cuda");
|
|
}
|
|
at::Tensor output = model(inputs).toTensor();
|
|
ASSERT_TRUE(output.device().is_cuda());
|
|
|
|
// Reference
|
|
auto ref_model = torch::jit::load(jit_filename);
|
|
ref_model.to(torch::kCUDA);
|
|
at::Tensor ref_output = ref_model.forward(inputs).toTensor();
|
|
|
|
ASSERT_TRUE(ref_output.allclose(output, 1e-03, 1e-05));
|
|
}
|
|
|
|
TEST(TorchDeployGPUTest, UsesDistributed) {
|
|
const auto model_filename = path(
|
|
"USES_DISTRIBUTED",
|
|
"torch/csrc/deploy/example/generated/uses_distributed");
|
|
torch::deploy::InterpreterManager m(1);
|
|
torch::deploy::Package p = m.loadPackage(model_filename);
|
|
{
|
|
auto I = p.acquireSession();
|
|
I.self.attr("import_module")({"uses_distributed"});
|
|
}
|
|
}
|
|
|
|
TEST(TorchDeployGPUTest, TensorRT) {
|
|
if (!torch::cuda::is_available()) {
|
|
GTEST_SKIP();
|
|
}
|
|
auto packagePath = path(
|
|
"MAKE_TRT_MODULE", "torch/csrc/deploy/example/generated/make_trt_module");
|
|
torch::deploy::InterpreterManager m(1);
|
|
torch::deploy::Package p = m.loadPackage(packagePath);
|
|
auto makeModel = p.loadPickle("make_trt_module", "model.pkl");
|
|
{
|
|
auto I = makeModel.acquireSession();
|
|
auto model = I.self(at::ArrayRef<at::IValue>{});
|
|
auto input = at::ones({1, 2, 3}).cuda();
|
|
auto output = input * 2;
|
|
ASSERT_TRUE(
|
|
output.allclose(model(at::IValue{input}).toIValue().toTensor()));
|
|
}
|
|
}
|