Manual Implementation of weight_norm() in C++ API

username: Koro
cooked: <p>Hello Pytorch Community,</p>
<p>since at this point in time, there’s no weight normalization available in Libtorch, but I still need it to port a model over to the C++ API, I tried to manually implement it in C++. I used the source code for the python implementation: <a href="https://pytorch.cn/docs/master/_modules/torch/nn/utils/weight_norm.html" rel="nofollow noopener">https://pytorch.cn/docs/master/_modules/torch/nn/utils/weight_norm.html</a> as inspiration.</p>
<p>But I’m just not sure if I can actually use this in real projects or if this approach is incorrect.</p>
<p>First, in the Constructor of the network, I call init_pre_forward_normalize(*layer) after initialising and registering the layer to my network:</p>
<pre><code class="lang-auto">void init_pre_forward_normalize(torch::nn::Module &amp; module)
{
	//Since we cannot delete the old Paramter, we treat it as our "weight_v"
	//We assume that we only want to apply this normalization to "weight" parameter
	auto old_weight = module.named_parameters().find("weight")-&gt;data();
	torch::Tensor new_weight = _norm(old_weight);
	module.register_parameter("weight_g", new_weight.data());
}
</code></pre>
<p>Using a simplified _norm function:</p>
<pre><code class="lang-auto">torch::Tensor _norm(torch::Tensor &amp;old_weight)
{
	//We assume, that always: dim=0
	torch::Tensor new_weight;
	if (old_weight.dim() == 1)
	{
		new_weight = old_weight.contiguous().view({ old_weight.size(0), -1 }).norm(2, 1)
			.view({ old_weight.size(0) });
	}
	if (old_weight.dim() == 2)
	{
		new_weight = old_weight.contiguous().view({ old_weight.size(0), -1 }).norm(2, 1)
			.view({ old_weight.size(0), 1 });
	}
	if (old_weight.dim() == 3)
	{
		new_weight = old_weight.contiguous().view({ old_weight.size(0), -1 }).norm(2, 1)
			.view({ old_weight.size(0), 1, 1 });
	}
	if (old_weight.dim() == 4)
	{
		new_weight = old_weight.contiguous().view({ old_weight.size(0), -1 }).norm(2, 1)
			.view({ old_weight.size(0), 1, 1, 1 });
	}
	return new_weight;
}
</code></pre>
<p>And then, in the forward Method of my network, I just call pre_forward_normalize(*layer) everytime before using that specific layer’s forward function.</p>
<pre><code class="lang-auto">void pre_forward_normalize(torch::nn::Module &amp; module)
{
	auto v = *(module.named_parameters().find("weight")); //This will be normalized and stay between -1.0 and 1.0
	auto g = *(module.named_parameters().find("weight_g")); //This will stay constant through out Learning process
	module.named_parameters().find("weight")-&gt;data() = v * (g / _norm(v));
}
</code></pre>
<p>I tested it with a simple 2-Linear-Layer Network and the XOR Dataset, and it seems to be working?, since the “weight” parameters of the layers are still changing and the loss is decreasing (though the network’s performance is getting worse, especially when using Normalization on both layers).<br>
I also noticed, that the weight’s values are always between -1.0 and 1.0 for the layers with Normalization, while the “weight_g” values are just staying the same during the training process.</p>
<p>I hope that somebody can look at my code and tell me if there’s something wrong here, because I’m not sure if I can rely on this implementation or not. Any help or suggestions are welcome.</p>
<p>with kind regards, Florian Korotschenko</p>

username: Koro
cooked: <p>So I compared the result to the same network, implemented in Pyton with the same seed and all. Turns out the resulting weights are not the same between this naive C++ implementation and pytorch. I guess there’s more happening inside pytorch than I thought, especially in the backwards pass.</p>
<p>Has somebody figured out how to implement weight_norm? Or is it just best to wait for official support in upcomming Libtorch releases?</p>

username: Koro
cooked: <p>Wrapping the desired Layer into a custom module solves the issue here. Using _norm() from above, you can implement a module with, for example, a Linear Layer + Weight Norm this way:</p>
<pre><code class="lang-auto">struct Linear_WN_Impl : torch::nn::Module
{
	torch::Tensor v;
	torch::Tensor g;
	torch::Tensor bias;
	torch::nn::LinearOptions _options;	//Hold Addition Information like (Stride, Padding, ...)

	Linear_WN_Impl(torch::nn::LinearOptions options) : _options(options)
	{
		torch::manual_seed(187); //Only for repeatability and comparison
		//Initialize v, g and bias
		auto default_layer = torch::nn::Linear(_options);
		v = default_layer-&gt;weight.data();
		g = _norm(v);
		bias = default_layer-&gt;bias.data();
		//Register as Parameters, make trainable
		v = this-&gt;register_parameter("v", v);
		g = this-&gt;register_parameter("g", g);
		bias = this-&gt;register_parameter("bias", bias);
	}
	torch::Tensor forward(torch::Tensor &amp; input)
	{
		//Forward Pass, use v and g to form normalized weight
		return torch::linear(input, v * (g / _norm(v)), bias);
	}
};
TORCH_MODULE(Linear_WN_);
</code></pre>
<p>then, you can utilze the module like here for example:</p>
<pre><code class="lang-auto">struct Net : torch::nn::Module
{
	Linear_WN_ fc1{ nullptr };
	torch::nn::Linear fc2{ nullptr };

	Net()
	{
                fc1 = register_module("fc1", Linear_WN_(torch::nn::LinearOptions(2, 4)));
		fc2 = register_module("fc2", torch::nn::Linear(4, 1));
	}

	torch::Tensor forward(torch::Tensor x) 
	{
		auto y = torch::sigmoid(fc1-&gt;forward(x));
		y = torch::sigmoid(fc2-&gt;forward(y));
		return y;
	}
};
</code></pre>
<p>I tested and compared it to equivalent Pytorch Code with the built-in weight_norm on the first layer, here are the results of both programs:<br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/1/9/1938b4bdceee0b17291e2fe5a7f763cf32f51539.png" data-download-href="https://discuss.pytorch.org/uploads/default/1938b4bdceee0b17291e2fe5a7f763cf32f51539" title="grafik"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/1/9/1938b4bdceee0b17291e2fe5a7f763cf32f51539_2_690x282.png" alt="grafik" data-base62-sha1="3B7qAviNFDGWT71Lw3JatvxF4Nz" width="690" height="282" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/1/9/1938b4bdceee0b17291e2fe5a7f763cf32f51539_2_690x282.png, https://discuss.pytorch.org/uploads/default/optimized/3X/1/9/1938b4bdceee0b17291e2fe5a7f763cf32f51539_2_1035x423.png 1.5x, https://discuss.pytorch.org/uploads/default/original/3X/1/9/1938b4bdceee0b17291e2fe5a7f763cf32f51539.png 2x" data-dominant-color="272722"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">grafik</span><span class="informations">1122×459 70.7 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div><br>
and after Training Loop:<br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/2/0/207f55559828ac8b28ec7886183c6a3bf56a3e9d.png" data-download-href="https://discuss.pytorch.org/uploads/default/207f55559828ac8b28ec7886183c6a3bf56a3e9d" title="grafik"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/2/0/207f55559828ac8b28ec7886183c6a3bf56a3e9d_2_690x439.png" alt="grafik" data-base62-sha1="4Du5xVxcBKAwl92388fbPj4wOzX" width="690" height="439" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/2/0/207f55559828ac8b28ec7886183c6a3bf56a3e9d_2_690x439.png, https://discuss.pytorch.org/uploads/default/optimized/3X/2/0/207f55559828ac8b28ec7886183c6a3bf56a3e9d_2_1035x658.png 1.5x, https://discuss.pytorch.org/uploads/default/original/3X/2/0/207f55559828ac8b28ec7886183c6a3bf56a3e9d.png 2x" data-dominant-color="292724"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">grafik</span><span class="informations">1151×733 141 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div><br>
So as you can see, the resulting weights are the same.</p>