From 192625cc7f661cfc91bd7906f74ab3f1ffdb290e Mon Sep 17 00:00:00 2001 From: root <23239305+b-chu@users.noreply.github.com> Date: Mon, 17 Jun 2024 20:45:39 +0000 Subject: [PATCH] Add setter for epoch in iteration --- composer/core/time.py | 25 +++++++++++++++++-------- tests/test_time.py | 7 +++++++ 2 files changed, 24 insertions(+), 8 deletions(-) diff --git a/composer/core/time.py b/composer/core/time.py index 3916dd7659..00af1fd456 100644 --- a/composer/core/time.py +++ b/composer/core/time.py @@ -525,13 +525,8 @@ def __init__( raise ValueError(f'The `token` argument has units of {token.unit}; not {TimeUnit.TOKEN}.') self._token = token - epoch_in_iteration = Time.from_input(epoch_in_iteration, TimeUnit.EPOCH) - if epoch_in_iteration.unit != TimeUnit.EPOCH: - raise ValueError(( - f'The `epoch_in_iteration` argument has units of {epoch_in_iteration.unit}; ' - f'not {TimeUnit.EPOCH}.' - )) - self._epoch_in_iteration = epoch_in_iteration + self._epoch_in_iteration = Time(0, TimeUnit.EPOCH) + self.epoch_in_iteration = epoch_in_iteration token_in_iteration = Time.from_input(token_in_iteration, TimeUnit.TOKEN) if token_in_iteration.unit != TimeUnit.TOKEN: @@ -619,7 +614,7 @@ def load_state_dict(self, state: dict[str, Any]) -> None: if 'iteration' in state: self._iteration = Time(state['iteration'], TimeUnit.ITERATION) if 'epoch_in_iteration' in state: - self._epoch_in_iteration = Time(state['epoch_in_iteration'], TimeUnit.EPOCH) + self.epoch_in_iteration = Time(state['epoch_in_iteration'], TimeUnit.EPOCH) if 'token_in_iteration' in state: self._token_in_iteration = Time(state['token_in_iteration'], TimeUnit.TOKEN) if 'iteration_wct' in state: @@ -655,6 +650,20 @@ def epoch_in_iteration(self) -> Time[int]: """The epoch count in the current iteration (resets at 0 at the beginning of every iteration).""" return self._epoch_in_iteration + @epoch_in_iteration.setter + def epoch_in_iteration( + self, + epoch_in_iteration: Union[int, Time[int]], # pyright: ignore[reportPropertyTypeMismatch] + ): + """Sets epoch count in the current iteration.""" + epoch_in_iteration = Time.from_input(epoch_in_iteration, TimeUnit.EPOCH) + if epoch_in_iteration.unit != TimeUnit.EPOCH: + raise ValueError(( + f'The `epoch_in_iteration` argument has units of {epoch_in_iteration.unit}; ' + f'not {TimeUnit.EPOCH}.' + )) + self._epoch_in_iteration = epoch_in_iteration + @property def token_in_iteration(self) -> Time[int]: """The token count in the current iteration (resets at 0 at the beginning of every iteration).""" diff --git a/tests/test_time.py b/tests/test_time.py index 1545eaa3b1..d585d9af36 100644 --- a/tests/test_time.py +++ b/tests/test_time.py @@ -146,6 +146,13 @@ def test_timestamp_update(): assert timestamp is not timestamp_2 +def test_set_timestamp(): + timestamp = Timestamp(epoch_in_iteration=1) + assert timestamp.epoch_in_iteration == 1 + timestamp.epoch_in_iteration = 2 + assert timestamp.epoch_in_iteration == 2 + + def test_timestamp_to_next_batch_epoch_iteration(): timestamp = Timestamp() # Step batch 0 in epoch 0