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

* Add support in fine_tune to selectively tune (freeze some parameters set through config file)

* Add tests for selective fine tuning.

* Allow for turning off gradients in train command (since in fine-tune as well this is happening with "trainer" configs).

* Add missing imports in fine_tune_test.py

* add tests for using 'no_grad' config with train command

* Code cleanup: 1. for regex matches 2. follow import convention

* Add logging statements for knowing tunable and frozen parameters.

Author: HarshTrivedi

allennlp/commands/fine_tune.py

@@ -10,6 +10,7 @@
 import logging
 import os
 from copy import deepcopy
+import re
 
 from allennlp.commands.evaluate import evaluate
 from allennlp.commands.subcommand import Subcommand
@@ -180,6 +181,24 @@ def fine_tune_model(model: Model,
     test_data = all_datasets.get('test')
 
     trainer_params = params.pop("trainer")
+    no_grad_regexes = trainer_params.pop("no_grad", ())
+
+    nograd_parameter_names = []
+    grad_parameter_names = []
+    for name, parameter in model.named_parameters():
+        if any(re.search(regex, name) for regex in no_grad_regexes):
+            parameter.requires_grad_(False)
+            nograd_parameter_names.append(name)
+        else:
+            grad_parameter_names.append(name)
+
+    logger.info("Following parameters are Frozen  (without gradient):")
+    for name in nograd_parameter_names:
+        logger.info(name)
+    logger.info("Following parameters are Tunable (with gradient):")
+    for name in grad_parameter_names:
+        logger.info(name)
+
     trainer = Trainer.from_params(model,
                                   serialization_dir,
                                   iterator,

allennlp/commands/train.py

@@ -37,6 +37,7 @@
 import logging
 import os
 from copy import deepcopy
+import re
 
 import torch
 
@@ -282,6 +283,24 @@ def train_model(params: Params,
     test_data = all_datasets.get('test')
 
     trainer_params = params.pop("trainer")
+    no_grad_regexes = trainer_params.pop("no_grad", ())
+
+    nograd_parameter_names = []
+    grad_parameter_names = []
+    for name, parameter in model.named_parameters():
+        if any(re.search(regex, name) for regex in no_grad_regexes):
+            parameter.requires_grad_(False)
+            nograd_parameter_names.append(name)
+        else:
+            grad_parameter_names.append(name)
+
+    logger.info("Following parameters are Frozen  (without gradient):")
+    for name in nograd_parameter_names:
+        logger.info(name)
+    logger.info("Following parameters are Tunable (with gradient):")
+    for name in grad_parameter_names:
+        logger.info(name)
+
     trainer = Trainer.from_params(model,
                                   serialization_dir,
                                   iterator,

allennlp/tests/commands/fine_tune_test.py

@@ -1,8 +1,15 @@
 # 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 +57,34 @@ 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
+            for name, parameter in tuned_model.named_parameters():
+                if any(re.search(regex, name) for regex in regex_list):
+                    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,6 +2,8 @@
 import argparse
 from typing import Iterable
 import os
+import shutil
+import re
 
 import pytest
 import torch
@@ -232,3 +234,52 @@ 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
+            for name, parameter in model.named_parameters():
+                if any(re.search(regex, name) for regex in regex_list):
+                    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)