Understand torch.nn.init.calculate_gain() with Examples – PyTorch Tutorial

By | May 5, 2022

torch.nn.init.calculate_gain() will return a gain value based on nonlinearity. In this tutorial, we will use some examples to show you how to use it.


It is defined as:

torch.nn.init.calculate_gain(nonlinearity, param=None)

It will return a recommended gain value for the given nonlinearity function.

As to nonlinearity, it can be:

torch.nn.init.calculate_gain() nonlinearity value list

As to negative_slope, it can be set by param.

For example:

gain = nn.init.calculate_gain('leaky_relu', 0.2)  # leaky_relu with negative_slope=0.2

How to use torch.nn.init.calculate_gain()?

We can call torch.nn.init.calculate_gain() in some torch init method, such as:

  • torch.nn.init.xavier_normal_(tensor, gain=1.0)
  • torch.nn.init.xavier_uniform_(tensor, gain=1.0)
  • torch.nn.init.orthogonal_(tensor, gain=1)

Here gain parameter can be calculated by torch.nn.init.calculate_gain().

For example:

self.linear_layer = torch.nn.Linear(in_dim, out_dim, bias=bias)

torch.nn.init.calculate_gain(“linear”) means gain = 1.0