How to reimplement regular linear layer using custom module and function in Pytorch C++ (libtorch)

clemsoncoder: <p>I’m trying to implement a custom module that calls a custom function.</p>
<p>I’m trying to just reimplement a simple linear layer in C++ to get it all working like in this python pytorch example:</p>
<p><a href="https://pytorch.org/docs/master/notes/extending.html" class="onebox" target="_blank" rel="noopener nofollow ugc">https://pytorch.org/docs/master/notes/extending.html</a></p>
<p>However I’m having issues declaring things such as wether I should declare<br>
torch::nn::ModuleHolder fc1{nullptr};<br>
or<br>
TORCH_MODULE(LinearModule)</p>
<p>Since the code I modified below compiles and runs but it doesn’t output the right output (outputs nan).</p>
<p>/* Author: Kushashwa, <a href="http://www.github.com/krshrimali" class="inline-onebox" rel="noopener nofollow ugc">krshrimali (Kushashwa Ravi Shrimali) · GitHub</a></p>
<ul>
<li>Reference:<br>
*/</li>
</ul>
<p><span class="hashtag">#include</span> &lt;torch/torch.h&gt;</p>
<p>using namespace torch::autograd;</p>
<p>// Inherit from Function<br>
class LinearFunction : public Function {<br>
public:<br>
// Note that both forward and backward are static functions</p>
<p>// bias is an optional argument<br>
static torch::Tensor forward(<br>
AutogradContext *ctx,<br>
torch::Tensor input,<br>
torch::Tensor weight,<br>
torch::Tensor bias = torch::Tensor()) {</p>
<pre><code>ctx-&gt;save_for_backward({input, weight, bias});

//*
std::cout &lt;&lt; "Input" &lt;&lt; std::endl;
std::cout &lt;&lt; input &lt;&lt; std::endl;
std::cout &lt;&lt; "weight" &lt;&lt; std::endl;
std::cout &lt;&lt; weight &lt;&lt; std::endl;
std::cout &lt;&lt; "DIM input" &lt;&lt; std::endl;
std::cout &lt;&lt; input.dim() &lt;&lt; std::endl;
std::cout &lt;&lt; "DIM weight" &lt;&lt; std::endl;
std::cout &lt;&lt; weight.dim() &lt;&lt; std::endl;

std::cout &lt;&lt; "in dim 0: " &lt;&lt; input.sizes()[0] &lt;&lt; std::endl;
std::cout &lt;&lt; "in dim 1: " &lt;&lt; input.sizes()[1] &lt;&lt; std::endl;

std::cout &lt;&lt; "weight dim 0: " &lt;&lt; weight.sizes()[0] &lt;&lt; std::endl;
std::cout &lt;&lt; "weight dim 1: " &lt;&lt; weight.sizes()[1] &lt;&lt; std::endl;
//*/

getchar();

auto output = input.t().mm(weight.t());
if (bias.defined()) {

    std::cout &lt;&lt; "bias.unsqueeze(0).expand_as(output)" &lt;&lt; std::endl;
    std::cout &lt;&lt; bias.unsqueeze(0).expand_as(output) &lt;&lt; std::endl;
    getchar();

  output += bias.unsqueeze(0).expand_as(output);
}
return output;
</code></pre>
<p>}</p>
<p>static tensor_list backward(AutogradContext *ctx, tensor_list grad_outputs) {<br>
auto saved = ctx-&gt;get_saved_variables();<br>
auto input = saved[0];<br>
auto weight = saved[1];<br>
auto bias = saved[2];</p>
<pre><code>auto grad_output = grad_outputs[0];

//std::cout &lt;&lt; grad_output &lt;&lt; std::endl;
//std::cout &lt;&lt; grad_outputs[0] &lt;&lt; std::endl;
//std::cout &lt;&lt; grad_outputs[0][0] &lt;&lt; std::endl;
//std::cout &lt;&lt; "-------" &lt;&lt; std::endl;

auto grad_input = grad_output.mm(weight);
auto grad_weight = grad_output.t().mm(input);
auto grad_bias = torch::Tensor();
if (bias.defined()) {
  grad_bias = grad_output.sum(0);
}

return {grad_input, grad_weight, grad_bias};
</code></pre>
<p>}<br>
};</p>
<p>/* Sample code for training a FCN on MNIST dataset using PyTorch C++ API */<br>
struct LinearModule: torch::nn::Module {<br>
LinearModule(){}<br>
LinearModule(int64_t in_features, int64_t out_features) {<br>
//weight = register_parameter(“weight”,torch::empty({in_features,out_features}));<br>
weight = register_parameter(“weight”,torch::empty({in_features,out_features}));<br>
bias = register_parameter(“bias”,torch::empty({out_features}));<br>
}<br>
// Implement Algorithm<br>
torch::Tensor forward(torch::Tensor input) {<br>
// std::cout &lt;&lt; x.size(0) &lt;&lt; ", " &lt;&lt; 784 &lt;&lt; std::endl;<br>
auto x = LinearFunction::apply(<br>
input,<br>
weight<br>
//, bias<br>
);<br>
return x;<br>
}</p>
<pre><code>torch::Tensor weight;
torch::Tensor bias;
</code></pre>
<p>};</p>
<p>/* Sample code for training a FCN on MNIST dataset using PyTorch C++ API */<br>
struct Net: torch::nn::Module {</p>
<pre><code>//std::shared_ptr&lt;LinearModule&gt; fc1;
//std::shared_ptr&lt;LinearModule&gt; fc2;

torch::nn::ModuleHolder&lt;LinearModule&gt; fc1{nullptr};
torch::nn::ModuleHolder&lt;LinearModule&gt; fc2{nullptr};

Net() {
	// Initialize CNN
	conv1 = register_module("conv1", torch::nn::Conv2d(torch::nn::Conv2dOptions(1, 10, 5)));
	conv2 = register_module("conv2", torch::nn::Conv2d(torch::nn::Conv2dOptions(10, 20, 5)));
	conv2_drop = register_module("conv2_drop", torch::nn::Dropout());
	//fc1 = register_module("fc1", torch::nn::Linear(320, 50));
	//fc2 = register_module("fc2", torch::nn::Linear(50, 10));

	//fc1 = register_module("fc1", LinearModule(320, 50));
	//fc2 = register_module("fc2", LinearModule(50, 10));
	fc1 = register_module&lt;LinearModule&gt;(
		"fc1", std::make_shared&lt;LinearModule&gt;(320,50));

	fc2 = register_module&lt;LinearModule&gt;(
		"fc2", std::make_shared&lt;LinearModule&gt;(50,10));

	//auto x = torch::randn({320, 50}).requires_grad_();
	//auto weight = torch::randn({320, 50}).requires_grad_();
	//fc1 = LinearFunction::apply(x, weight);
	

	//auto x2 = torch::randn({50,10}).requires_grad_();
	//auto weight2 = torch::randn({50, 10}).requires_grad_();
	//fc2 = LinearFunction::apply(x2, weight2);
}

// Implement Algorithm
torch::Tensor forward(torch::Tensor x) {
	// std::cout &lt;&lt; x.size(0) &lt;&lt; ", " &lt;&lt; 784 &lt;&lt; std::endl;
	x = torch::relu(torch::max_pool2d(conv1-&gt;forward(x), 2));
	x = torch::relu(torch::max_pool2d(conv2_drop-&gt;forward(conv2-&gt;forward(x)), 2));
	x = x.view({-1, 320});
	x = torch::relu(fc1-&gt;forward(x));
	x = torch::dropout(x, 0.5, is_training());
	x = fc2-&gt;forward(x);

	return torch::log_softmax(x, 1);
}
torch::nn::Conv2d conv1{nullptr};
torch::nn::Conv2d conv2{nullptr};
torch::nn::Dropout conv2_drop{nullptr};
//torch::nn::Linear fc1{nullptr}, fc2{nullptr};

//LinearModule fc1():weight(torch::Tensor()),bias(torch::Tensor()){};
//LinearModule fc2():weight(torch::Tensor()),bias(torch::Tensor()){};
//LinearModule fc1():weight(torch::Tensor()),bias(torch::Tensor()){};

//LinearModule fc1{torch::Tensor(),torch::Tensor()};
//LinearModule fc2{torch::Tensor(),torch::Tensor()};
//LinearModule *fc1;
//LinearModule *fc2;
</code></pre>
<p>};</p>
<p>int main() {<br>
if(0){<br>
auto x = torch::randn({2, 3}).requires_grad_();<br>
auto weight = torch::randn({4, 3}).requires_grad_();<br>
auto y = LinearFunction::apply(x, weight);<br>
y.sum().backward();</p>
<pre><code>	std::cout &lt;&lt; x.grad() &lt;&lt; std::endl;
	std::cout &lt;&lt; weight.grad() &lt;&lt; std::endl;


	return 0;
}

auto net = std::make_shared&lt;Net&gt;();

// Create multi-threaded data loader for MNIST data
auto data_loader = torch::data::make_data_loader&lt;torch::data::samplers::SequentialSampler&gt;(
		std::move(torch::data::datasets::MNIST("../data").map(torch::data::transforms::Normalize&lt;&gt;(0.13707, 0.3081)).map(
			torch::data::transforms::Stack&lt;&gt;())), 64);
torch::optim::SGD optimizer(net-&gt;parameters(), 0.01); // Learning Rate 0.01

// net.train();

for(size_t epoch=1; epoch&lt;=10; ++epoch) {
	size_t batch_index = 0;
	// Iterate data loader to yield batches from the dataset
	for (auto&amp; batch: *data_loader) {
		// Reset gradients
		optimizer.zero_grad();
		// Execute the model
		torch::Tensor prediction = net-&gt;forward(batch.data);
		// Compute loss value
		torch::Tensor loss = torch::nll_loss(prediction, batch.target);
		// Compute gradients
		loss.backward();
		// Update the parameters
		optimizer.step();

		// Output the loss and checkpoint every 100 batches
		if (++batch_index % 100 == 0) {
			std::cout &lt;&lt; "Epoch: " &lt;&lt; epoch &lt;&lt; " | Batch: " &lt;&lt; batch_index 
				&lt;&lt; " | Loss: " &lt;&lt; loss.item&lt;float&gt;() &lt;&lt; std::endl;
			torch::save(net, "net.pt");
		}
	}
}
</code></pre>
<p>}</p>