|
import torch.nn as nn
|
|
import torch
|
|
|
|
def constant_init(module, val, bias=0):
|
|
nn.init.constant_(module.weight, val)
|
|
if hasattr(module, 'bias') and module.bias is not None:
|
|
nn.init.constant_(module.bias, bias)
|
|
|
|
def xavier_init(module, gain=1, bias=0, distribution='normal'):
|
|
assert distribution in ['uniform', 'normal']
|
|
if distribution == 'uniform':
|
|
nn.init.xavier_uniform_(module.weight, gain=gain)
|
|
else:
|
|
nn.init.xavier_normal_(module.weight, gain=gain)
|
|
if hasattr(module, 'bias') and module.bias is not None:
|
|
nn.init.constant_(module.bias, bias)
|
|
|
|
|
|
def normal_init(module, mean=0, std=1, bias=0):
|
|
nn.init.normal_(module.weight, mean, std)
|
|
if hasattr(module, 'bias') and module.bias is not None:
|
|
nn.init.constant_(module.bias, bias)
|
|
|
|
|
|
def uniform_init(module, a=0, b=1, bias=0):
|
|
nn.init.uniform_(module.weight, a, b)
|
|
if hasattr(module, 'bias') and module.bias is not None:
|
|
nn.init.constant_(module.bias, bias)
|
|
|
|
|
|
def kaiming_init(module,
|
|
a=0,
|
|
is_rnn=False,
|
|
mode='fan_in',
|
|
nonlinearity='leaky_relu',
|
|
bias=0,
|
|
distribution='normal'):
|
|
assert distribution in ['uniform', 'normal']
|
|
if distribution == 'uniform':
|
|
if is_rnn:
|
|
for name, param in module.named_parameters():
|
|
if 'bias' in name:
|
|
nn.init.constant_(param, bias)
|
|
elif 'weight' in name:
|
|
nn.init.kaiming_uniform_(param,
|
|
a=a,
|
|
mode=mode,
|
|
nonlinearity=nonlinearity)
|
|
else:
|
|
nn.init.kaiming_uniform_(module.weight,
|
|
a=a,
|
|
mode=mode,
|
|
nonlinearity=nonlinearity)
|
|
|
|
else:
|
|
if is_rnn:
|
|
for name, param in module.named_parameters():
|
|
if 'bias' in name:
|
|
nn.init.constant_(param, bias)
|
|
elif 'weight' in name:
|
|
nn.init.kaiming_normal_(param,
|
|
a=a,
|
|
mode=mode,
|
|
nonlinearity=nonlinearity)
|
|
else:
|
|
nn.init.kaiming_normal_(module.weight,
|
|
a=a,
|
|
mode=mode,
|
|
nonlinearity=nonlinearity)
|
|
|
|
if not is_rnn and hasattr(module, 'bias') and module.bias is not None:
|
|
nn.init.constant_(module.bias, bias)
|
|
|
|
|
|
def bilinear_kernel(in_channels, out_channels, kernel_size):
|
|
factor = (kernel_size + 1) // 2
|
|
if kernel_size % 2 == 1:
|
|
center = factor - 1
|
|
else:
|
|
center = factor - 0.5
|
|
og = (torch.arange(kernel_size).reshape(-1, 1),
|
|
torch.arange(kernel_size).reshape(1, -1))
|
|
filt = (1 - torch.abs(og[0] - center) / factor) * \
|
|
(1 - torch.abs(og[1] - center) / factor)
|
|
weight = torch.zeros((in_channels, out_channels,
|
|
kernel_size, kernel_size))
|
|
weight[range(in_channels), range(out_channels), :, :] = filt
|
|
return weight
|
|
|
|
|
|
def init_weights(m):
|
|
|
|
|
|
if isinstance(m, nn.Conv2d):
|
|
kaiming_init(m)
|
|
elif isinstance(m, (nn.BatchNorm2d, nn.GroupNorm)):
|
|
constant_init(m, 1)
|
|
elif isinstance(m, nn.Linear):
|
|
xavier_init(m)
|
|
elif isinstance(m, (nn.LSTM, nn.LSTMCell)):
|
|
kaiming_init(m, is_rnn=True)
|
|
|
|
|
|
|
