Search code examples
c++pytorchforwardlibtorch

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


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?


Solution

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