diff --git a/fastNLP/modules/decoder/CRF.py b/fastNLP/modules/decoder/CRF.py index e6327ec0..991927da 100644 --- a/fastNLP/modules/decoder/CRF.py +++ b/fastNLP/modules/decoder/CRF.py @@ -1,6 +1,7 @@ import torch from torch import nn +from fastNLP.modules.utils import initial_parameter def log_sum_exp(x, dim=-1): max_value, _ = x.max(dim=dim, keepdim=True) @@ -19,7 +20,7 @@ def seq_len_to_byte_mask(seq_lens): class ConditionalRandomField(nn.Module): - def __init__(self, tag_size, include_start_end_trans=True): + def __init__(self, tag_size, include_start_end_trans=True ,initial_method = None): """ :param tag_size: int, num of tags :param include_start_end_trans: bool, whether to include start/end tag @@ -35,8 +36,8 @@ class ConditionalRandomField(nn.Module): self.start_scores = nn.Parameter(torch.randn(tag_size)) self.end_scores = nn.Parameter(torch.randn(tag_size)) - self.reset_parameter() - + # self.reset_parameter() + initial_parameter(self, initial_method) def reset_parameter(self): nn.init.xavier_normal_(self.transition_m) if self.include_start_end_trans: diff --git a/fastNLP/modules/decoder/MLP.py b/fastNLP/modules/decoder/MLP.py index c70aa0e9..b8fb95f0 100644 --- a/fastNLP/modules/decoder/MLP.py +++ b/fastNLP/modules/decoder/MLP.py @@ -1,8 +1,8 @@ import torch import torch.nn as nn - +from fastNLP.modules.utils import initial_parameter class MLP(nn.Module): - def __init__(self, size_layer, num_class=2, activation='relu'): + def __init__(self, size_layer, num_class=2, activation='relu' , initial_method = None): """Multilayer Perceptrons as a decoder Args: @@ -36,7 +36,7 @@ class MLP(nn.Module): self.hidden_active = activation else: raise ValueError("should set activation correctly: {}".format(activation)) - + initial_parameter(self, initial_method ) def forward(self, x): for layer in self.hiddens: x = self.hidden_active(layer(x)) diff --git a/fastNLP/modules/encoder/char_embedding.py b/fastNLP/modules/encoder/char_embedding.py index 72680e5b..1da63947 100644 --- a/fastNLP/modules/encoder/char_embedding.py +++ b/fastNLP/modules/encoder/char_embedding.py @@ -1,11 +1,12 @@ import torch import torch.nn.functional as F from torch import nn +# from torch.nn.init import xavier_uniform - +from fastNLP.modules.utils import initial_parameter class ConvCharEmbedding(nn.Module): - def __init__(self, char_emb_size=50, feature_maps=(40, 30, 30), kernels=(3, 4, 5)): + def __init__(self, char_emb_size=50, feature_maps=(40, 30, 30), kernels=(3, 4, 5),initial_method = None): """ Character Level Word Embedding :param char_emb_size: the size of character level embedding. Default: 50 @@ -20,6 +21,8 @@ class ConvCharEmbedding(nn.Module): nn.Conv2d(1, feature_maps[i], kernel_size=(char_emb_size, kernels[i]), bias=True, padding=(0, 4)) for i in range(len(kernels))]) + initial_parameter(self,initial_method) + def forward(self, x): """ :param x: [batch_size * sent_length, word_length, char_emb_size] @@ -53,7 +56,7 @@ class LSTMCharEmbedding(nn.Module): :param hidden_size: int, the number of hidden units. Default: equal to char_emb_size. """ - def __init__(self, char_emb_size=50, hidden_size=None): + def __init__(self, char_emb_size=50, hidden_size=None , initial_method= None): super(LSTMCharEmbedding, self).__init__() self.hidden_size = char_emb_size if hidden_size is None else hidden_size @@ -62,7 +65,7 @@ class LSTMCharEmbedding(nn.Module): num_layers=1, bias=True, batch_first=True) - + initial_parameter(self, initial_method) def forward(self, x): """ :param x:[ n_batch*n_word, word_length, char_emb_size] diff --git a/fastNLP/modules/encoder/conv.py b/fastNLP/modules/encoder/conv.py index 06a31dd8..68536e5d 100644 --- a/fastNLP/modules/encoder/conv.py +++ b/fastNLP/modules/encoder/conv.py @@ -6,6 +6,7 @@ import torch.nn as nn from torch.nn.init import xavier_uniform_ # import torch.nn.functional as F +from fastNLP.modules.utils import initial_parameter class Conv(nn.Module): """ @@ -15,7 +16,7 @@ class Conv(nn.Module): def __init__(self, in_channels, out_channels, kernel_size, stride=1, padding=0, dilation=1, - groups=1, bias=True, activation='relu'): + groups=1, bias=True, activation='relu',initial_method = None ): super(Conv, self).__init__() self.conv = nn.Conv1d( in_channels=in_channels, @@ -26,7 +27,7 @@ class Conv(nn.Module): dilation=dilation, groups=groups, bias=bias) - xavier_uniform_(self.conv.weight) + # xavier_uniform_(self.conv.weight) activations = { 'relu': nn.ReLU(), @@ -37,6 +38,7 @@ class Conv(nn.Module): raise Exception( 'Should choose activation function from: ' + ', '.join([x for x in activations])) + initial_parameter(self, initial_method) def forward(self, x): x = torch.transpose(x, 1, 2) # [N,L,C] -> [N,C,L] diff --git a/fastNLP/modules/encoder/conv_maxpool.py b/fastNLP/modules/encoder/conv_maxpool.py index f666e7f9..7aa897cf 100644 --- a/fastNLP/modules/encoder/conv_maxpool.py +++ b/fastNLP/modules/encoder/conv_maxpool.py @@ -5,7 +5,7 @@ import torch import torch.nn as nn import torch.nn.functional as F from torch.nn.init import xavier_uniform_ - +from fastNLP.modules.utils import initial_parameter class ConvMaxpool(nn.Module): """ @@ -14,7 +14,7 @@ class ConvMaxpool(nn.Module): def __init__(self, in_channels, out_channels, kernel_sizes, stride=1, padding=0, dilation=1, - groups=1, bias=True, activation='relu'): + groups=1, bias=True, activation='relu',initial_method = None ): super(ConvMaxpool, self).__init__() # convolution @@ -47,6 +47,8 @@ class ConvMaxpool(nn.Module): raise Exception( "Undefined activation function: choose from: relu") + initial_parameter(self, initial_method) + def forward(self, x): # [N,L,C] -> [N,C,L] x = torch.transpose(x, 1, 2) diff --git a/fastNLP/modules/encoder/linear.py b/fastNLP/modules/encoder/linear.py index 9582d9f9..a7c5f6c3 100644 --- a/fastNLP/modules/encoder/linear.py +++ b/fastNLP/modules/encoder/linear.py @@ -1,6 +1,6 @@ import torch.nn as nn - +from fastNLP.modules.utils import initial_parameter class Linear(nn.Module): """ Linear module @@ -12,10 +12,10 @@ class Linear(nn.Module): bidirectional : If True, becomes a bidirectional RNN """ - def __init__(self, input_size, output_size, bias=True): + def __init__(self, input_size, output_size, bias=True,initial_method = None ): super(Linear, self).__init__() self.linear = nn.Linear(input_size, output_size, bias) - + initial_parameter(self, initial_method) def forward(self, x): x = self.linear(x) return x diff --git a/fastNLP/modules/encoder/lstm.py b/fastNLP/modules/encoder/lstm.py index bed6c276..5af09f29 100644 --- a/fastNLP/modules/encoder/lstm.py +++ b/fastNLP/modules/encoder/lstm.py @@ -1,6 +1,6 @@ import torch.nn as nn - +from fastNLP.modules.utils import initial_parameter class Lstm(nn.Module): """ LSTM module @@ -13,11 +13,13 @@ class Lstm(nn.Module): bidirectional : If True, becomes a bidirectional RNN. Default: False. """ - def __init__(self, input_size, hidden_size=100, num_layers=1, dropout=0, bidirectional=False): + def __init__(self, input_size, hidden_size=100, num_layers=1, dropout=0, bidirectional=False , initial_method = None): super(Lstm, self).__init__() self.lstm = nn.LSTM(input_size, hidden_size, num_layers, bias=True, batch_first=True, dropout=dropout, bidirectional=bidirectional) - + initial_parameter(self, initial_method) def forward(self, x): x, _ = self.lstm(x) return x +if __name__ == "__main__": + lstm = Lstm(10) diff --git a/fastNLP/modules/encoder/masked_rnn.py b/fastNLP/modules/encoder/masked_rnn.py index 76f828a9..c1ef15d0 100644 --- a/fastNLP/modules/encoder/masked_rnn.py +++ b/fastNLP/modules/encoder/masked_rnn.py @@ -4,7 +4,7 @@ import torch import torch.nn as nn import torch.nn.functional as F - +from fastNLP.modules.utils import initial_parameter def MaskedRecurrent(reverse=False): def forward(input, hidden, cell, mask, train=True, dropout=0): """ @@ -192,7 +192,7 @@ def AutogradMaskedStep(num_layers=1, dropout=0, train=True, lstm=False): class MaskedRNNBase(nn.Module): def __init__(self, Cell, input_size, hidden_size, num_layers=1, bias=True, batch_first=False, - layer_dropout=0, step_dropout=0, bidirectional=False, **kwargs): + layer_dropout=0, step_dropout=0, bidirectional=False, initial_method = None , **kwargs): """ :param Cell: :param input_size: @@ -226,7 +226,7 @@ class MaskedRNNBase(nn.Module): cell = self.Cell(layer_input_size, hidden_size, self.bias, **kwargs) self.all_cells.append(cell) self.add_module('cell%d' % (layer * num_directions + direction), cell) # Max的代码写得真好看 - + initial_parameter(self, initial_method) def reset_parameters(self): for cell in self.all_cells: cell.reset_parameters() diff --git a/fastNLP/modules/encoder/variational_rnn.py b/fastNLP/modules/encoder/variational_rnn.py index b08bdd2d..fb75fabb 100644 --- a/fastNLP/modules/encoder/variational_rnn.py +++ b/fastNLP/modules/encoder/variational_rnn.py @@ -6,6 +6,7 @@ import torch.nn.functional as F from torch.nn._functions.thnn import rnnFusedPointwise as fusedBackend from torch.nn.parameter import Parameter +from fastNLP.modules.utils import initial_parameter def default_initializer(hidden_size): stdv = 1.0 / math.sqrt(hidden_size) @@ -172,7 +173,7 @@ def AutogradVarMaskedStep(num_layers=1, lstm=False): class VarMaskedRNNBase(nn.Module): def __init__(self, Cell, input_size, hidden_size, num_layers=1, bias=True, batch_first=False, - dropout=(0, 0), bidirectional=False, initializer=None, **kwargs): + dropout=(0, 0), bidirectional=False, initializer=None,initial_method = None, **kwargs): super(VarMaskedRNNBase, self).__init__() self.Cell = Cell @@ -193,7 +194,7 @@ class VarMaskedRNNBase(nn.Module): cell = self.Cell(layer_input_size, hidden_size, self.bias, p=dropout, initializer=initializer, **kwargs) self.all_cells.append(cell) self.add_module('cell%d' % (layer * num_directions + direction), cell) - + initial_parameter(self, initial_method) def reset_parameters(self): for cell in self.all_cells: cell.reset_parameters() @@ -284,7 +285,7 @@ class VarFastLSTMCell(VarRNNCellBase): \end{array} """ - def __init__(self, input_size, hidden_size, bias=True, p=(0.5, 0.5), initializer=None): + def __init__(self, input_size, hidden_size, bias=True, p=(0.5, 0.5), initializer=None,initial_method =None): super(VarFastLSTMCell, self).__init__() self.input_size = input_size self.hidden_size = hidden_size @@ -311,7 +312,7 @@ class VarFastLSTMCell(VarRNNCellBase): self.p_hidden = p_hidden self.noise_in = None self.noise_hidden = None - + initial_parameter(self, initial_method) def reset_parameters(self): for weight in self.parameters(): if weight.dim() == 1: diff --git a/fastNLP/modules/utils.py b/fastNLP/modules/utils.py index 442944e7..22139668 100644 --- a/fastNLP/modules/utils.py +++ b/fastNLP/modules/utils.py @@ -2,8 +2,8 @@ from collections import defaultdict import numpy as np import torch - - +import torch.nn.init as init +import torch.nn as nn def mask_softmax(matrix, mask): if mask is None: result = torch.nn.functional.softmax(matrix, dim=-1) @@ -11,6 +11,51 @@ def mask_softmax(matrix, mask): raise NotImplementedError return result +def initial_parameter(net ,initial_method =None): + + if initial_method == 'xavier_uniform': + init_method = init.xavier_uniform_ + elif initial_method=='xavier_normal': + init_method = init.xavier_normal_ + elif initial_method == 'kaiming_normal' or initial_method =='msra': + init_method = init.kaiming_normal + elif initial_method == 'kaiming_uniform': + init_method = init.kaiming_normal + elif initial_method == 'orthogonal': + init_method = init.orthogonal_ + elif initial_method == 'sparse': + init_method = init.sparse_ + elif initial_method =='normal': + init_method = init.normal_ + elif initial_method =='uniform': + initial_method = init.uniform_ + else: + init_method = init.xavier_normal_ + def weights_init(m): + # classname = m.__class__.__name__ + if isinstance(m, nn.Conv2d) or isinstance(m,nn.Conv1d) or isinstance(m,nn.Conv3d): # for all the cnn + if initial_method != None: + init_method(m.weight.data) + else: + init.xavier_normal_(m.weight.data) + init.normal_(m.bias.data) + elif isinstance(m, nn.LSTM): + for w in m.parameters(): + if len(w.data.size())>1: + init_method(w.data) # weight + else: + init.normal_(w.data) # bias + elif hasattr(m, 'weight') and m.weight.requires_grad: + init_method(m.weight.data) + else: + for w in m.parameters() : + if w.requires_grad: + if len(w.data.size())>1: + init_method(w.data) # weight + else: + init.normal_(w.data) # bias + # print("init else") + net.apply(weights_init) def seq_mask(seq_len, max_len): mask = [torch.ge(torch.LongTensor(seq_len), i + 1) for i in range(max_len)]