PyTorch C++ 前端：注册新模块并在转发期间使用它们

Question

我正在创建一个空模型，如下所示：

struct TestNet : torch::nn::Module {

    TestNet() {
    }

    torch::Tensor Forward(torch::Tensor x)
    {
        return x;
    }
};

然后我将新模块注册到模型中：

auto net = std::make_shared<TestNet>();
torch::nn::ModuleHolder<ConvLayer> conv(1, 1, 3, 1, 1);
net->register_module("conv1", conv);

其中 ConvLayer 是一个带有卷积层的模块：

struct ConvLayer : torch::nn::Module {
    ConvLayer() {}
    ConvLayer(int in_ch, int out_ch, int kernel, int pad, int stride)
        : conv1(torch::nn::Conv2dOptions(in_ch, out_ch, kernel)
            .stride(stride)
            .padding(pad)
            .bias(false))
    {
        register_module("Conv1", conv1);
    }
    torch::Tensor Forward(torch::Tensor x)
    {
        return conv1(x);
    }

    torch::nn::Conv2d conv1{ nullptr };
};

我现在可以打印出 TestNet 的参数并查看卷积层，但是我不能在前向传播中使用它。我缺少什么才能做到这一点？

Answer 1

我找到了一种使用 torch::nn::Sequential 来做到这一点的方法，希望这对其他人有帮助：

struct TestNet2 : torch::nn::Module {

    TestNet2() {
        layers = register_module("layers", torch::nn::Sequential());
    }
    template <typename T>
    void sequentialLayer(T Layer)
    {
        layers->push_back(Layer);
    }
    torch::Tensor Forward(torch::Tensor x)
    {
        return layers->forward(x);
    }
    torch::nn::Sequential layers;
};

...

auto net = std::unique_ptr<TestNet2>();
auto convLayer = torch::nn::Conv2d(torch::nn::Conv2dOptions(1, 1, 3)
        .stride(1)
        .padding(1)
        .bias(false));

net->sequentialLayer(convLayer);

PyTorch C++ 前端：注册新模块并在转发期间使用它们

PyTorch C++ Frontend: Registering New Modules and using them during Forward

c++

pytorch

forward

libtorch