mirror of
https://github.com/pytorch/pytorch.git
synced 2025-10-20 21:14:14 +08:00
e0643fa3fc
Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/66744 Modified loops in files under fbsource/fbcode/caffe2/ from the format `for(TYPE var=x0;var<x_max;x++)` to the format `for(const auto var: irange(xmax))` This was achieved by running r-barnes's loop upgrader script (D28874212) with some modification to exclude all files under /torch/jit and a number of reversions or unused variable suppression warnings added by hand. Test Plan: Sandcastle Reviewed By: ngimel Differential Revision: D31705358 fbshipit-source-id: d6ea350cbaa8f452fc78f238160e5374be637a48
49 lines
1.3 KiB
C++
49 lines
1.3 KiB
C++
// This is a simple predictor binary that loads a TorchScript CV model and runs
|
|
// a forward pass with fixed input `torch::ones({1, 3, 224, 224})`.
|
|
// It's used for end-to-end integration test for custom mobile build.
|
|
|
|
#include <iostream>
|
|
#include <string>
|
|
#include <c10/util/irange.h>
|
|
#include <torch/script.h>
|
|
|
|
using namespace std;
|
|
|
|
namespace {
|
|
|
|
struct MobileCallGuard {
|
|
// Set InferenceMode for inference only use case.
|
|
c10::InferenceMode guard;
|
|
// Disable graph optimizer to ensure list of unused ops are not changed for
|
|
// custom mobile build.
|
|
torch::jit::GraphOptimizerEnabledGuard no_optimizer_guard{false};
|
|
};
|
|
|
|
torch::jit::Module loadModel(const std::string& path) {
|
|
MobileCallGuard guard;
|
|
auto module = torch::jit::load(path);
|
|
module.eval();
|
|
return module;
|
|
}
|
|
|
|
} // namespace
|
|
|
|
int main(int argc, const char* argv[]) {
|
|
if (argc < 2) {
|
|
std::cerr << "Usage: " << argv[0] << " <model_path>\n";
|
|
return 1;
|
|
}
|
|
auto module = loadModel(argv[1]);
|
|
auto input = torch::ones({1, 3, 224, 224});
|
|
auto output = [&]() {
|
|
MobileCallGuard guard;
|
|
return module.forward({input}).toTensor();
|
|
}();
|
|
|
|
std::cout << std::setprecision(3) << std::fixed;
|
|
for (const auto i : c10::irange(5)) {
|
|
std::cout << output.data_ptr<float>()[i] << std::endl;
|
|
}
|
|
return 0;
|
|
}
|