Add support for selective finetune (freeze parameters by regex from config file)

allennlp/commands/fine_tune.py

@@ -10,7 +10,7 @@
 import logging
 import os
 from copy import deepcopy
-
+import re
 from allennlp.commands.evaluate import evaluate
 from allennlp.commands.subcommand import Subcommand
 from allennlp.commands.train import datasets_from_params
@@ -166,6 +166,12 @@ def fine_tune_model(model: Model,
     test_data = all_datasets.get('test')
 
     trainer_params = params.pop("trainer")
+    nograd_regex_list = trainer_params.pop("no_grad", ())
+    if nograd_regex_list:
+        nograd_regex = "(" + ")|(".join(nograd_regex_list) + ")"
+        for name, parameter in model.named_parameters():
+            if re.search(nograd_regex, name):
+                parameter.requires_grad_(False)
     trainer = Trainer.from_params(model,
                                   serialization_dir,
                                   iterator,

allennlp/commands/train.py

@@ -37,7 +37,7 @@
 import logging
 import os
 from copy import deepcopy
-
+import re
 import torch
 
 from allennlp.commands.evaluate import evaluate
@@ -282,6 +282,12 @@ def train_model(params: Params,
     test_data = all_datasets.get('test')
 
     trainer_params = params.pop("trainer")
+    nograd_regex_list = trainer_params.pop("no_grad", ())
+    if nograd_regex_list:
+        nograd_regex = "(" + ")|(".join(nograd_regex_list) + ")"
+        for name, parameter in model.named_parameters():
+            if re.search(nograd_regex, name):
+                parameter.requires_grad_(False)
     trainer = Trainer.from_params(model,
                                   serialization_dir,
                                   iterator,

allennlp/tests/commands/fine_tune_test.py

@@ -1,8 +1,13 @@
 # pylint: disable=invalid-name,no-self-use
 import argparse
-
+import re
+import shutil
+import pytest
 from allennlp.common.testing import AllenNlpTestCase
-from allennlp.commands.fine_tune import FineTune, fine_tune_model_from_file_paths, fine_tune_model_from_args
+from allennlp.commands.fine_tune import FineTune, fine_tune_model_from_file_paths, \
+                               fine_tune_model_from_args, fine_tune_model
+from allennlp.common.params import Params
+from allennlp.models import load_archive
 
 class TestFineTune(AllenNlpTestCase):
     def setUp(self):
@@ -50,3 +55,36 @@ def test_fine_tune_fails_without_required_args(self):
         with self.assertRaises(SystemExit) as context:
             self.parser.parse_args(["fine-tune", "-s", "serialization_dir", "-c", "path/to/config"])
             assert context.exception.code == 2  # argparse code for incorrect usage
+
+    def test_fine_tune_nograd_regex(self):
+        original_model = load_archive(self.model_archive).model
+        name_parameters_original = dict(original_model.named_parameters())
+        regex_lists = [[],
+                       [".*attend_feedforward.*", ".*token_embedder.*"],
+                       [".*compare_feedforward.*"]]
+        for regex_list in regex_lists:
+            params = Params.from_file(self.config_file)
+            params["trainer"]["no_grad"] = regex_list
+            shutil.rmtree(self.serialization_dir, ignore_errors=True)
+            tuned_model = fine_tune_model(model=original_model,
+                                          params=params,
+                                          serialization_dir=self.serialization_dir)
+            # If regex is matched, parameter name should have requires_grad False
+            # If regex is matched, parameter name should have same requires_grad
+            # as the originally loaded model
+            if regex_list:
+                nograd_regex = "(" + ")|(".join(regex_list) + ")"
+                for name, parameter in tuned_model.named_parameters():
+                    if re.search(nograd_regex, name):
+                        assert not parameter.requires_grad
+                    else:
+                        assert parameter.requires_grad \
+                        == name_parameters_original[name].requires_grad
+        # If all parameters have requires_grad=False, then error.
+        with pytest.raises(Exception) as _:
+            params = Params.from_file(self.config_file)
+            params["trainer"]["no_grad"] = ["*"]
+            shutil.rmtree(self.serialization_dir, ignore_errors=True)
+            tuned_model = fine_tune_model(model=original_model,
+                                          params=params,
+                                          serialization_dir=self.serialization_dir)

allennlp/tests/commands/train_test.py

@@ -2,7 +2,8 @@
 import argparse
 from typing import Iterable
 import os
-
+import shutil
+import re
 import pytest
 import torch
 
@@ -232,3 +233,54 @@ def test_train_with_test_set(self):
         })
 
         train_model(params, serialization_dir=os.path.join(self.TEST_DIR, 'lazy_test_set'))
+
+    def test_train_nograd_regex(self):
+        params_get = lambda: Params({
+                "model": {
+                        "type": "simple_tagger",
+                        "text_field_embedder": {
+                                "tokens": {
+                                        "type": "embedding",
+                                        "embedding_dim": 5
+                                }
+                        },
+                        "encoder": {
+                                "type": "lstm",
+                                "input_size": 5,
+                                "hidden_size": 7,
+                                "num_layers": 2
+                        }
+                },
+                "dataset_reader": {"type": "sequence_tagging"},
+                "train_data_path": SEQUENCE_TAGGING_DATA_PATH,
+                "validation_data_path": SEQUENCE_TAGGING_DATA_PATH,
+                "iterator": {"type": "basic", "batch_size": 2},
+                "trainer": {
+                        "num_epochs": 2,
+                        "optimizer": "adam"
+                }
+        })
+        serialization_dir = os.path.join(self.TEST_DIR, 'test_train_nograd')
+        regex_lists = [[],
+                       [".*text_field_embedder.*"],
+                       [".*text_field_embedder.*", ".*encoder.*"]]
+        for regex_list in regex_lists:
+            params = params_get()
+            params["trainer"]["no_grad"] = regex_list
+            shutil.rmtree(serialization_dir, ignore_errors=True)
+            model = train_model(params, serialization_dir=serialization_dir)
+            # If regex is matched, parameter name should have requires_grad False
+            # Or else True
+            if regex_list:
+                nograd_regex = "(" + ")|(".join(regex_list) + ")"
+                for name, parameter in model.named_parameters():
+                    if re.search(nograd_regex, name):
+                        assert not parameter.requires_grad
+                    else:
+                        assert parameter.requires_grad
+        # If all parameters have requires_grad=False, then error.
+        params = params_get()
+        params["trainer"]["no_grad"] = ["*"]
+        shutil.rmtree(serialization_dir, ignore_errors=True)
+        with pytest.raises(Exception) as _:
+            model = train_model(params, serialization_dir=serialization_dir)