Skip to content

BUG: Unknown type name '__torch__.torch.classes.fbgemm.AtomicCounter' #3643

Description

@intfish123

When I load the model using C++, I get an error:

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

c++:

#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;
}

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions