对Pytorch神经网络初始化kaiming分布详解

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

fan_in和fan_out

pytorch计算fan_in和fan_out的源码


def _calculate_fan_in_and_fan_out(tensor):
 dimensions = tensor.ndimension()
 if dimensions < 2:
  raise ValueError("Fan in and fan out can not be computed 
  for tensor with fewer than 2 dimensions")

 if dimensions == 2: # Linear
  fan_in = tensor.size(1)
  fan_out = tensor.size(0)
 else:
  num_input_fmaps = tensor.size(1)
  num_output_fmaps = tensor.size(0)
  receptive_field_size = 1
  if tensor.dim() > 2:
   receptive_field_size = tensor[0][0].numel()
  fan_in = num_input_fmaps * receptive_field_size
  fan_out = num_output_fmaps * receptive_field_size

 return fan_in, fan_out

torch.nn.init.xavier_uniform_(tensor, gain=1)

>>> w = torch.empty(3, 5)
>>> nn.init.xavier_uniform_(w, gain=nn.init.calculate_gain('relu'))

torch.nn.init.xavier_normal_(tensor, gain=1)

torch.nn.init.kaiming_uniform_
 (tensor, a=0, mode='fan_in', nonlinearity='leaky_relu')

>>> w = torch.empty(3, 5)
>>> nn.init.kaiming_uniform_(w, mode='fan_in', nonlinearity='relu')

torch.nn.init.kaiming_normal_
 (tensor, a=0, mode='fan_in', nonlinearity='leaky_relu')

 >>> w = torch.empty(3, 5)
>>> nn.init.kaiming_normal_(w, mode='fan_out', nonlinearity='relu')