How to get model attribute using c++ api


(Sixer Wang) #1

I want to predict my model using pytorch C++ api, and I want to access one model attribute. Model definition is like:

class MyModule(torch.jit.ScriptModule):
    __constants__ = ['mods']
    def __init__(self, N, M):
        super(MyModule, self).__init__()
        self.linear = nn.Linear(N, M)
        self.mods = N

    @torch.jit.script_method
    def forward(self, inputs):
        output = self.linear(inputs)
        return output

And I want to get “mods” in c++. However, when I use the following code, It doesn’t work:

std::shared_ptr<torch::jit::script::Module> module = torch::jit::load(argv[1]);

int mods = module->mods;

So, what could I do to access the model attribute using c++ api?


(Martin Huber) #2

I guess its supposed to be done with

module->get_attribute("mods");

but somehow it segfaults for me


(Sixer Wang) #3

Thanks for ur advice. I refer the doc(https://github.com/pytorch/pytorch/blob/v1.0.1/torch/csrc/jit/script/module.h), and find the method “run_method”. I get the value according to this method. And my pytorch version is v1.0.1.

My model is like:

class MyModule(torch.jit.ScriptModule):
    __constants__ = ['mods']
    def __init__(self, N, M):
        super(MyModule, self).__init__()
        self.linear = nn.Linear(N, M)
        self.mods = N

    @torch.jit.script_method
    def get_mods(self):
    	return self.mods

    @torch.jit.script_method
    def forward(self, inputs):
        output = self.linear(inputs)
        return output

"get_mods"method is used to get the attribute “mods”.

My cpp code is like:

auto mods = module1->run_method("get_mods");
auto value = mods.toScalar();
int mods_value = value.to<int>();

It works well for me.