Add regularization parameter to Models (#3120)

allennlp/models/basic_classifier.py

@@ -1,12 +1,12 @@
-from typing import Dict
+from typing import Dict, Optional
 
 from overrides import overrides
 import torch
 
 from allennlp.data import Vocabulary
 from allennlp.models.model import Model
 from allennlp.modules import Seq2SeqEncoder, Seq2VecEncoder, TextFieldEmbedder
-from allennlp.nn import InitializerApplicator
+from allennlp.nn import InitializerApplicator, RegularizerApplicator
 from allennlp.nn.util import get_text_field_mask
 from allennlp.training.metrics import CategoricalAccuracy
 
@@ -41,6 +41,8 @@ class BasicClassifier(Model):
         Vocabulary namespace corresponding to labels. By default, we use the "labels" namespace.
     initializer : ``InitializerApplicator``, optional (default=``InitializerApplicator()``)
         If provided, will be used to initialize the model parameters.
+    regularizer : ``RegularizerApplicator``, optional (default=``None``)
+        If provided, will be used to calculate the regularization penalty during training.
     """
     def __init__(self,
                  vocab: Vocabulary,
@@ -50,9 +52,10 @@ def __init__(self,
                  dropout: float = None,
                  num_labels: int = None,
                  label_namespace: str = "labels",
-                 initializer: InitializerApplicator = InitializerApplicator()) -> None:
+                 initializer: InitializerApplicator = InitializerApplicator(),
+                 regularizer: Optional[RegularizerApplicator] = None) -> None:
 
-        super().__init__(vocab)
+        super().__init__(vocab, regularizer)
         self._text_field_embedder = text_field_embedder
 
         if seq2seq_encoder:

allennlp/models/bert_for_classification.py

@@ -1,4 +1,4 @@
-from typing import Dict, Union
+from typing import Dict, Union, Optional
 
 from overrides import overrides
 import torch
@@ -8,6 +8,7 @@
 from allennlp.models.model import Model
 from allennlp.modules.token_embedders.bert_token_embedder import PretrainedBertModel
 from allennlp.nn.initializers import InitializerApplicator
+from allennlp.nn import RegularizerApplicator
 from allennlp.training.metrics import CategoricalAccuracy
 
 
@@ -42,6 +43,8 @@ class BertForClassification(Model):
         Otherwise, they will be frozen and only the final linear layer will be trained.
     initializer : ``InitializerApplicator``, optional
         If provided, will be used to initialize the final linear layer *only*.
+    regularizer : ``RegularizerApplicator``, optional (default=``None``)
+        If provided, will be used to calculate the regularization penalty during training.
     """
     def __init__(self,
                  vocab: Vocabulary,
@@ -51,8 +54,9 @@ def __init__(self,
                  index: str = "bert",
                  label_namespace: str = "labels",
                  trainable: bool = True,
-                 initializer: InitializerApplicator = InitializerApplicator()) -> None:
-        super().__init__(vocab)
+                 initializer: InitializerApplicator = InitializerApplicator(),
+                 regularizer: Optional[RegularizerApplicator] = None,) -> None:
+        super().__init__(vocab, regularizer)
 
         if isinstance(bert_model, str):
             self.bert_model = PretrainedBertModel.load(bert_model)

allennlp/models/bidirectional_lm.py

@@ -1,9 +1,11 @@
+from typing import Optional
+
 from allennlp.data.vocabulary import Vocabulary
 from allennlp.models.language_model import LanguageModel
 from allennlp.models.model import Model
 from allennlp.modules.text_field_embedders import TextFieldEmbedder
 from allennlp.modules.seq2seq_encoders import Seq2SeqEncoder
-from allennlp.nn import InitializerApplicator
+from allennlp.nn import InitializerApplicator, RegularizerApplicator
 
 
 @Model.register('bidirectional-language-model')
@@ -37,6 +39,8 @@ class BidirectionalLanguageModel(LanguageModel):
         the full ``_SoftmaxLoss`` defined above.
     sparse_embeddings: ``bool``, optional (default: False)
         Passed on to ``SampledSoftmaxLoss`` if True.
+    regularizer : ``RegularizerApplicator``, optional (default=``None``)
+        If provided, will be used to calculate the regularization penalty during training.
     """
     def __init__(self,
                  vocab: Vocabulary,
@@ -45,12 +49,14 @@ def __init__(self,
                  dropout: float = None,
                  num_samples: int = None,
                  sparse_embeddings: bool = False,
-                 initializer: InitializerApplicator = None) -> None:
+                 initializer: InitializerApplicator = None,
+                 regularizer: Optional[RegularizerApplicator] = None) -> None:
         super().__init__(vocab=vocab,
                          text_field_embedder=text_field_embedder,
                          contextualizer=contextualizer,
                          dropout=dropout,
                          num_samples=num_samples,
                          sparse_embeddings=sparse_embeddings,
                          bidirectional=True,
-                         initializer=initializer)
+                         initializer=initializer,
+                         regularizer=regularizer)

allennlp/models/event2mind.py

@@ -17,6 +17,7 @@
 from allennlp.models.model import Model
 from allennlp.nn.beam_search import BeamSearch
 from allennlp.nn.util import get_text_field_mask, sequence_cross_entropy_with_logits
+from allennlp.nn import RegularizerApplicator
 from allennlp.training.metrics import UnigramRecall
 
 
@@ -55,6 +56,8 @@ class Event2Mind(Model):
     target_embedding_dim : int, optional (default = source_embedding_dim)
         You can specify an embedding dimensionality for the target side. If not, we'll use the same
         value as the source embedder's.
+    regularizer : ``RegularizerApplicator``, optional (default=``None``)
+        If provided, will be used to calculate the regularization penalty during training.
     """
     def __init__(self,
                  vocab: Vocabulary,
@@ -65,8 +68,9 @@ def __init__(self,
                  beam_size: int = 10,
                  target_names: List[str] = None,
                  target_namespace: str = "tokens",
-                 target_embedding_dim: int = None) -> None:
-        super().__init__(vocab)
+                 target_embedding_dim: int = None,
+                 regularizer: Optional[RegularizerApplicator] = None) -> None:
+        super().__init__(vocab, regularizer)
         target_names = target_names or ["xintent", "xreact", "oreact"]
 
         # Note: The original tweaks the embeddings for "personx" to be the mean

allennlp/models/language_model.py

@@ -1,4 +1,4 @@
-from typing import Dict, List, Tuple, Union
+from typing import Dict, List, Tuple, Union, Optional
 
 import torch
 import numpy as np
@@ -10,7 +10,7 @@
 from allennlp.modules.sampled_softmax_loss import SampledSoftmaxLoss
 from allennlp.modules.seq2seq_encoders import Seq2SeqEncoder
 from allennlp.nn.util import get_text_field_mask
-from allennlp.nn import InitializerApplicator
+from allennlp.nn import InitializerApplicator, RegularizerApplicator
 from allennlp.training.metrics import Perplexity
 
 
@@ -88,6 +88,8 @@ class LanguageModel(Model):
         Train a bidirectional language model, where the contextualizer
         is used to predict the next and previous token for each input token.
         This must match the bidirectionality of the contextualizer.
+    regularizer : ``RegularizerApplicator``, optional (default=``None``)
+        If provided, will be used to calculate the regularization penalty during training.
     """
     def __init__(self,
                  vocab: Vocabulary,
@@ -97,8 +99,9 @@ def __init__(self,
                  num_samples: int = None,
                  sparse_embeddings: bool = False,
                  bidirectional: bool = False,
-                 initializer: InitializerApplicator = None) -> None:
-        super().__init__(vocab)
+                 initializer: InitializerApplicator = None,
+                 regularizer: Optional[RegularizerApplicator] = None) -> None:
+        super().__init__(vocab, regularizer)
         self._text_field_embedder = text_field_embedder
 
         if contextualizer.is_bidirectional() is not bidirectional: