diff --git a/poetry.lock b/poetry.lock index 9ab99b42d..328d5927b 100644 --- a/poetry.lock +++ b/poetry.lock @@ -162,6 +162,24 @@ docs = ["furo", "myst-parser", "sphinx", "sphinx-notfound-page", "sphinxcontrib- tests = ["attrs[tests-no-zope]", "zope-interface"] tests-no-zope = ["cloudpickle", "hypothesis", "mypy (>=1.1.1)", "pympler", "pytest (>=4.3.0)", "pytest-mypy-plugins", "pytest-xdist[psutil]"] +[[package]] +name = "beartype" +version = "0.17.2" +description = "Unbearably fast runtime type checking in pure Python." +optional = false +python-versions = ">=3.8.0" +files = [ + {file = "beartype-0.17.2-py3-none-any.whl", hash = "sha256:c22b21e1f785cfcf5c4d3d13070f532b6243a3ad67e68d2298ff08d539847dce"}, + {file = "beartype-0.17.2.tar.gz", hash = "sha256:e911e1ae7de4bccd15745f7643609d8732f64de5c2fb844e89cbbed1c5a8d495"}, +] + +[package.extras] +all = ["typing-extensions (>=3.10.0.0)"] +dev = ["autoapi (>=0.9.0)", "coverage (>=5.5)", "equinox", "mypy (>=0.800)", "numpy", "pandera", "pydata-sphinx-theme (<=0.7.2)", "pytest (>=4.0.0)", "sphinx", "sphinx (>=4.2.0,<6.0.0)", "sphinxext-opengraph (>=0.7.5)", "tox (>=3.20.1)", "typing-extensions (>=3.10.0.0)"] +doc-rtd = ["autoapi (>=0.9.0)", "pydata-sphinx-theme (<=0.7.2)", "sphinx (>=4.2.0,<6.0.0)", "sphinxext-opengraph (>=0.7.5)"] +test-tox = ["equinox", "mypy (>=0.800)", "numpy", "pandera", "pytest (>=4.0.0)", "sphinx", "typing-extensions (>=3.10.0.0)"] +test-tox-coverage = ["coverage (>=5.5)"] + [[package]] name = "certifi" version = "2023.7.22" @@ -386,6 +404,21 @@ files = [ {file = "distlib-0.3.7.tar.gz", hash = "sha256:9dafe54b34a028eafd95039d5e5d4851a13734540f1331060d31c9916e7147a8"}, ] +[[package]] +name = "ema-pytorch" +version = "0.4.2" +description = "Easy way to keep track of exponential moving average version of your pytorch module" +optional = false +python-versions = "*" +files = [ + {file = "ema-pytorch-0.4.2.tar.gz", hash = "sha256:a1017c04994151e6d1b93f3e82fd02d828390e01c64a73cf6a302544c0cde4a2"}, + {file = "ema_pytorch-0.4.2-py3-none-any.whl", hash = "sha256:1734a5be32cd3c56b34ebc8ef54c726e47c14748493bfb71e7bc978ec8a3d9e7"}, +] + +[package.dependencies] +beartype = "*" +torch = ">=1.6" + [[package]] name = "exceptiongroup" version = "1.1.3" @@ -1944,4 +1977,4 @@ multidict = ">=4.0" [metadata] lock-version = "2.0" python-versions = "^3.9" -content-hash = "cb5b3ad5422c31ab6ce27a1f34a7ea982aba476d9c53e9a97018ddb3d2ffcb00" +content-hash = "7146302f48d40292143d44ce85f75c797ef49988f63354256150a0713c7891a4" diff --git a/pyproject.toml b/pyproject.toml index c0d108748..71f5df3b5 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -32,6 +32,7 @@ pytorch-crf = ">=0.7.2" networkx = "^3.0.0" # because of BartModelWithDecoderPositionIds transformers = "^4.35.0" +ema-pytorch = "^0.4.2" [tool.poetry.group.dev.dependencies] torch = {version = "^2.1.0+cpu", source = "pytorch"} diff --git a/src/pie_modules/models/simple_generative.py b/src/pie_modules/models/simple_generative.py index a4bdae0b8..a0943e7ff 100644 --- a/src/pie_modules/models/simple_generative.py +++ b/src/pie_modules/models/simple_generative.py @@ -3,6 +3,7 @@ from typing import Any, Dict, Optional, Tuple, Type, Union import torch +from ema_pytorch import EMA from pytorch_ie.core import PyTorchIEModel from pytorch_lightning.utilities.types import OptimizerLRScheduler from torch import FloatTensor, LongTensor @@ -11,7 +12,7 @@ from transformers.modeling_outputs import Seq2SeqLMOutput from typing_extensions import TypeAlias -from pie_modules.models.common import ModelWithBoilerplate +from pie_modules.models.common import TRAINING, ModelWithBoilerplate from pie_modules.utils import resolve_type logger = logging.getLogger(__name__) @@ -65,6 +66,8 @@ def __init__( override_generation_kwargs: Optional[Dict[str, Any]] = None, # scheduler / optimizer warmup_proportion: float = 0.0, + # Exponential Moving Average (EMA) + use_ema: bool = False, # important: the following entries are only used if the base model does not have a configure_optimizer method! learning_rate: Optional[float] = None, optimizer_type: Optional[Union[str, Type[Optimizer]]] = None, @@ -83,6 +86,14 @@ def __init__( resolved_base_model_type: Type[PreTrainedModel] = resolve_type(base_model_type) self.model = resolved_base_model_type.from_pretrained(**base_model_config) self.generation_config = self.configure_generation(**(override_generation_kwargs or {})) + self.use_ema = use_ema + if use_ema: + self.ema = EMA( + self.model, + beta=0.9999, # exponential moving average factor + update_after_step=100, # only after this number of .update() calls will it start updating + update_every=10, # how often to actually update, to save on compute (updates every 10th .update() call) + ) def configure_generation(self, **kwargs) -> Dict[str, Any]: if self.taskmodule is not None: @@ -97,6 +108,12 @@ def configure_generation(self, **kwargs) -> Dict[str, Any]: generation_config.update(kwargs) return generation_config + def training_step(self, batch: StepInputType, batch_idx: int) -> StepOutputType: + result = self._step(stage=TRAINING, batch=batch) + if self.use_ema: + self.ema.update() + return result + def predict(self, inputs, **kwargs) -> TargetType: is_training = self.training self.eval() @@ -149,7 +166,7 @@ def configure_optimizers(self) -> OptimizerLRScheduler: resolved_optimizer_type = resolve_type( self.optimizer_type, expected_super_type=Optimizer ) - optimizer = resolved_optimizer_type(self.parameters(), lr=self.learning_rate) + optimizer = resolved_optimizer_type(self.model.parameters(), lr=self.learning_rate) if self.warmup_proportion > 0.0: stepping_batches = self.trainer.estimated_stepping_batches diff --git a/tests/models/test_simple_generative.py b/tests/models/test_simple_generative.py index fa6a2ceb5..ea7419ed3 100644 --- a/tests/models/test_simple_generative.py +++ b/tests/models/test_simple_generative.py @@ -25,8 +25,8 @@ def taskmodule(): ) -@pytest.fixture(scope="module") -def model(taskmodule) -> SimpleGenerativeModel: +@pytest.fixture(scope="module", params=[False, True]) +def model(taskmodule, request) -> SimpleGenerativeModel: return SimpleGenerativeModel( base_model_type="transformers.AutoModelForSeq2SeqLM", base_model_config=dict(pretrained_model_name_or_path=MODEL_ID), @@ -36,6 +36,7 @@ def model(taskmodule) -> SimpleGenerativeModel: # use a strange learning rate to make sure it is passed through learning_rate=13e-3, optimizer_type="torch.optim.Adam", + use_ema=request.param, )