pytorc/libtorch C++中的自定义子模块

Custom submodules in pytorch / libtorch C++

本文关键字：自定义 模块 libtorch C++ pytorc      更新时间：2024-04-28

完全披露，大约几天前我在PyTorch论坛上问了同样的问题，但没有得到回复，所以从技术上讲，这是一个转发，但我相信这仍然是一个好问题，因为我在网上找不到答案。如下：

您能展示一个将register_module与自定义模块一起使用的示例吗我在网上找到的唯一例子是将线性层或卷积层注册为子模块。

我试着写我自己的模块，并将其注册到另一个模块，但我无法让它工作。我的IDE告诉我no instance of overloaded function "MyModel::register_module" matches the argument list -- argument types are: (const char [14], TreeEmbedding)

(TreeEmbedding是我制作的另一个结构的名称，它扩展了torch:：nn:：Module。(

我是不是错过了什么？举个例子会很有帮助。



编辑：下面是其他上下文。

我有一个头文件"model.h"，其中包含以下内容：

struct TreeEmbedding : torch::nn::Module {
TreeEmbedding();
torch::Tensor forward(Graph tree);
};
struct MyModel : torch::nn::Module{
size_t embeddingSize;
TreeEmbedding treeEmbedding;
MyModel(size_t embeddingSize=10);
torch::Tensor forward(std::vector<Graph> clauses, std::vector<Graph> contexts);
};

我还有一个cpp文件"model.cpp"，其中包含以下内容：

MyModel::MyModel(size_t embeddingSize) :
embeddingSize(embeddingSize)
{
treeEmbedding = register_module("treeEmbedding", TreeEmbedding{});
}

此设置仍然存在与上面相同的错误。文档中的代码确实有效(使用线性层等内置组件(，但使用自定义模块则无效。在追踪到torch:：nn:：Linear之后，它看起来就像是一个ModuleHolder(不管是什么…(

谢谢，Jack

如果有人能提供更多细节，我会接受更好的答案，但为了防止有人想知道，我想我会提供我能找到的一些信息：

register_module接收一个字符串作为它的第一个参数，它的第二个参数可以是ModuleHolder(我不知道这是什么…(，也可以是模块的shared_ptr。下面是我的例子：

treeEmbedding = register_module<TreeEmbedding>("treeEmbedding", make_shared<TreeEmbedding>());

到目前为止，这似乎对我有效。

相关文章：