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

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.

torch.nn.init.calculate_gain()

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:

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.xavier_uniform_(
            self.linear_layer.weight,
            gain=torch.nn.init.calculate_gain("linear"))

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