Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

Make offline ER us total batch size in first update #381

Merged
merged 6 commits into from
Aug 18, 2023
Merged

Make offline ER us total batch size in first update #381

Changes from 1 commit
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
75125ae
Make offline ER use full batch size in first update
lballes Aug 17, 2023
d345fe4
Reuse dataloader code and adapt integ test
lballes Aug 17, 2023
f35b22f
Foo
lballes Aug 17, 2023
c95c465
Foo
lballes Aug 17, 2023
636710c
Remove unused import
lballes Aug 17, 2023
9beec20
Adapt integ test
lballes Aug 17, 2023
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
20 changes: 18 additions & 2 deletions src/renate/updaters/experimental/offline_er.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -13,6 +13,7 @@
from torch.utils.data import DataLoader, Dataset

from renate import defaults
from renate.data.datasets import _TransformedDataset
from renate.models import RenateModule
from renate.types import NestedTensors
from renate.updaters.learner import ReplayLearner
Expand Down Expand Up @@ -71,9 +72,9 @@ def on_model_update_start(
self._num_points_current_task = len(train_dataset)

def train_dataloader(self) -> DataLoader:
train_loader = super().train_dataloader()
loaders = {"current_task": train_loader}
loaders = {}
if len(self._memory_buffer) > self._memory_batch_size:
loaders["current_task"] = super().train_dataloader()
loaders["memory"] = DataLoader(
dataset=self._memory_buffer,
batch_size=self._memory_batch_size,
Expand All @@ -83,6 +84,21 @@ def train_dataloader(self) -> DataLoader:
pin_memory=True,
collate_fn=self._train_collate_fn,
)
else:
# Manually create a dataloader for the current task with total batch size.
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Can we reuse the code instead of overriding it?

self._batch_size += self._memory_batch_size
loaders["current_task"] = super().train_dataloader()

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I can do something like that if you prefer it. We would have to undo the change after creating the data loader though.

train_dataset = _TransformedDataset(
self._train_dataset,
transform=self._train_transform,
target_transform=self._train_target_transform,
)
loaders["current_task"] = DataLoader(
train_dataset,
batch_size=self._batch_size + self._memory_batch_size,
shuffle=True,
generator=self._rng,
pin_memory=True,
collate_fn=self._train_collate_fn,
)
return CombinedLoader(loaders, mode="max_size_cycle")

def on_model_update_end(self) -> None:
Expand Down
Loading