".github/git@developer.sourcefind.cn:OpenDAS/torchaudio.git" did not exist on "eabf1a1324e476f9b3d54c68c0c475455edb6e3a"
Commit 9af1a2b0 authored by Kai Chen's avatar Kai Chen
Browse files

add weight initialization methods

parent e363982a
from .resnet import ResNet, make_res_layer from .resnet import ResNet, make_res_layer
from .weight_init import xavier_init, normal_init, uniform_init, kaiming_init
__all__ = ['ResNet', 'make_res_layer'] __all__ = [
'ResNet', 'make_res_layer', 'xavier_init', 'normal_init', 'uniform_init',
'kaiming_init'
]
import torch.nn as nn
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'):
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'):
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'):
nn.init.constant_(module.bias, bias)
def kaiming_init(module,
mode='fan_out',
nonlinearity='relu',
bias=0,
distribution='normal'):
assert distribution in ['uniform', 'normal']
if distribution == 'uniform':
nn.init.kaiming_uniform_(
module.weight, mode=mode, nonlinearity=nonlinearity)
else:
nn.init.kaiming_normal_(
module.weight, mode=mode, nonlinearity=nonlinearity)
if hasattr(module, 'bias'):
nn.init.constant_(module.bias, bias)
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment