diff --git a/fastNLP/modules/decoder/MLP.py b/fastNLP/modules/decoder/MLP.py index 766dc225..0470e91b 100644 --- a/fastNLP/modules/decoder/MLP.py +++ b/fastNLP/modules/decoder/MLP.py @@ -4,12 +4,13 @@ from fastNLP.modules.utils import initial_parameter class MLP(nn.Module): - def __init__(self, size_layer, activation='relu', initial_method=None): + def __init__(self, size_layer, activation='relu', initial_method=None, dropout=0.0): """Multilayer Perceptrons as a decoder :param size_layer: list of int, define the size of MLP layers. :param activation: str or function, the activation function for hidden layers. :param initial_method: str, the name of init method. + :param dropout: float, the probability of dropout. .. note:: There is no activation function applying on output layer. @@ -24,6 +25,8 @@ class MLP(nn.Module): else: self.hiddens.append(nn.Linear(size_layer[i-1], size_layer[i])) + self.dropout = nn.Dropout(p=dropout) + actives = { 'relu': nn.ReLU(), 'tanh': nn.Tanh(), @@ -38,8 +41,8 @@ class MLP(nn.Module): def forward(self, x): for layer in self.hiddens: - x = self.hidden_active(layer(x)) - x = self.output(x) + x = self.dropout(self.hidden_active(layer(x))) + x = self.dropout(self.output(x)) return x