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

Question

I am creating a model that is empty like so:

struct TestNet : torch::nn::Module {

    TestNet() {
    }

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

I then register new modules to the model:

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

Where ConvLayer is a Module with a convolutional layer:

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

I can print out the parameters of TestNet now and see the convolutional layer, however I can not utilize it in a forward pass. What am I missing to do this?

Answer 1

I found a way to do this using torch::nn::Sequential , hope this helps anyone else:

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);