terminate called after throwing an instance of 'torch::jit::ErrorReport'
what():
Unknown type name '__torch__.torch.classes.fbgemm.AtomicCounter':
File "code/__torch__/fbgemm_gpu/split_table_batched_embeddings_ops_inference.py", line 67
cache_assoc : int
cache_algorithm : __torch__.fbgemm_gpu.split_table_batched_embeddings_ops_common.CacheAlgorithm
timestep_counter : __torch__.torch.classes.fbgemm.AtomicCounter
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ <--- HERE
timestep_prefetch_size : __torch__.torch.classes.fbgemm.AtomicCounter
max_prefetch_depth : int
#include <iostream>
#include "absl/flags/flag.h"
#include "absl/flags/parse.h"
#include "absl/strings/str_format.h"
#include "predictor.pb.h"
#include "predictor.grpc.pb.h"
#include <dlfcn.h>
#include <torch/torch.h>
#include <torch/script.h>
#include <torch/nn/functional/activation.h>
#include <grpcpp/grpcpp.h>
#include "grpcpp/health_check_service_interface.h"
#include "grpcpp/ext/proto_server_reflection_plugin.h"
ABSL_FLAG(u_int16_t, port, 50051, "Server port for the service");
#define NUM_BYTES_FLOAT_FEATURES 4
#define NUM_BYTES_SPARSE_FEATURES 4
class PredictorServiceHandler final : public predictor::Predictor::Service {
public:
PredictorServiceHandler(torch::jit::script::Module &module) : module_(module) {
}
grpc::Status Predict(grpc::ServerContext *context, const predictor::PredictionRequest *request,
predictor::PredictionResponse *response) override {
std::cout << "Predict Called!" << std::endl;
// at (ATen) 全称: A Tensor Library。它是 PyTorch 的核心数学运算库。
// c10 (Core) 全称: Core Tensor Library (也有说法是 Caffe2 + PyTorch 合并时的 Core 库),它是 PyTorch 的底层基础设施库,比 ATen 更底层。它不包含复杂的数学运算,只负责管理最基础的东西
c10::Dict<std::string, at::Tensor> dict;
predictor::FloatFeatures floatFeature = request->float_features();
std::string floatFeatureBlob = floatFeature.values();
auto numFloatFeatures = floatFeature.num_features();
auto batchSize = floatFeatureBlob.size() / (NUM_BYTES_FLOAT_FEATURES * numFloatFeatures);
std::cout << "Size: " << floatFeatureBlob.size() << " Num Features: " << numFloatFeatures << std::endl;
auto floatFeatureTensor = torch::from_blob(floatFeatureBlob.data(), {static_cast<int64_t>(batchSize), numFloatFeatures},torch::kFloat32);
predictor::SparseFeatures idListFeature = request->id_list_features();
auto numIdListFeatures = idListFeature.num_features();
auto lengthsBlob = idListFeature.lengths();
auto valuesBlob = idListFeature.values();
std::cout << "Lengths Size: " << lengthsBlob.size() << " Num Features: " << numIdListFeatures << std::endl;
assert(batchSize ==(lengthsBlob.size() / (NUM_BYTES_SPARSE_FEATURES * numIdListFeatures)));
auto lengthsTensor = torch::from_blob(
lengthsBlob.data(),
{static_cast<long>(lengthsBlob.size()) / NUM_BYTES_SPARSE_FEATURES},
torch::kInt32);
auto valuesTensor = torch::from_blob(
valuesBlob.data(),
{static_cast<long>(valuesBlob.size()) / NUM_BYTES_SPARSE_FEATURES},
torch::kInt32);
dict.insert("float_features", floatFeatureTensor.to(torch::kCUDA));
dict.insert("id_list_features.lengths", lengthsTensor.to(torch::kCUDA));
dict.insert("id_list_features.values", valuesTensor.to(torch::kCUDA));
std::vector<c10::IValue> input;
input.push_back(c10::IValue(dict));
torch::Tensor output = this->module_.forward(input).toGenericDict().at("default").toTensor();
auto predictions = response->mutable_predictions();
predictor::FloatVec fv;
fv.mutable_data()->Add(output.data_ptr<float>(), output.data_ptr<float>() + output.numel());
(*predictions)["default"] = fv;
return grpc::Status::OK;
}
private:
torch::jit::script::Module &module_;
};
void RunServer(uint16_t port, torch::jit::script::Module &module) {
std::string server_address = absl::StrFormat("0.0.0.0:%d", port);
PredictorServiceHandler service(module);
grpc::EnableDefaultHealthCheckService(true);
grpc::reflection::InitProtoReflectionServerBuilderPlugin();
grpc::ServerBuilder builder;
builder.AddListeningPort(server_address, grpc::InsecureServerCredentials());
builder.RegisterService(&service);
std::unique_ptr<grpc::Server> server(builder.BuildAndStart());
std::cout << "Server listening on " << server_address << std::endl;
server->Wait();
}
int main(int argc, char **argv) {
if (argc != 2) {
std::cerr << "Usage: ts-infer <path-to-exported-model>\n";
return -1;
}
std::cout << "Loading model...\n";
torch::jit::script::Module module;
try {
module = torch::jit::load(argv[1]);
} catch (const c10::Error& e) {
std::cerr << "Error loading model: " << argv[1] << ", error: " << e.msg() << std::endl;
return -1;
}
torch::NoGradGuard no_grad;
module.eval();
std::cout << "Sanity Check with dummy inputs" << std::endl;
c10::Dict<std::string, at::Tensor> dict;
dict.insert("float_features", torch::ones({1, 13}, torch::dtype(torch::kFloat32).device(torch::kCUDA, 0)));
dict.insert("id_list_features.lengths", torch::ones({26}, torch::dtype(torch::kLong).device(torch::kCUDA, 0)));
dict.insert("id_list_features.values", torch::ones({26}, torch::dtype(torch::kLong).device(torch::kCUDA, 0)));
std::vector<c10::IValue> input;
input.push_back(c10::IValue(dict));
auto output = module.forward(input).toGenericDict().at("default").toTensor();
std::cout << "Model Forward Completed, Output: "<< output.item<float>() << std::endl;
RunServer(absl::GetFlag(FLAGS_port), module);
return 0;
}
When I load the model using C++, I get an error:
c++: