storage policy abstract
+from abc import ABC, abstractmethod
+from typing import TypeVar, Optional, Dict, TYPE_CHECKING
+
from torch.utils.data import random_split, ConcatDataset
from avalanche.benchmarks.utils import AvalancheConcatDataset
@@ -30,6 +33,8 @@ class ReplayPlugin(StrategyPlugin):
self.mem_size = mem_size
self.ext_mem = {} # a Dict<task_id, Dataset>
+ self.storage_policy = RandomStoragePolicy(replay_mem=self.ext_mem,
+ experience_balanced=True)
def before_training_exp(self, strategy, num_workers=0, shuffle=True,
**kwargs):
@@ -52,7 +57,65 @@ class ReplayPlugin(StrategyPlugin):
the current training batch/task, adding the patterns from the new
experience and removing those from past experiences to comply the limit
of the total number of patterns in memory """
+ self.storage_policy(strategy)
+
+
+class StoragePolicy(ABC):
+ """ Any policy to store exemplars in a replay memory."""
+
+ @abstractmethod
+ def __call__(self, replay_mem: Dict, **kwargs):
+ """Store exemplars in the replay memory"""
+ pass
+
+
+class ExperienceBalancedStoragePolicy(StoragePolicy):
+
+ def __init__(self,
+ replay_mem: Dict, mem_size: int,
+ class_balanced=False,
+ ):
+ self.replay_mem = replay_mem
+ """
+ The replay memory as a Dict. Keys are the conditionals (experience/class), values are the samples stored.
+ """
+
+ assert not (experience_balanced and class_balanced), \
+ "Can't be both experience and class balanced."
+ self.experience_balanced = experience_balanced
+ """
+ Each experience gets the same amount of space in the replay memory.
+ """
+
+ self.class_balanced = class_balanced
+ """
+ Each class label gets the same amount of space in the replay memory.
+ """
+
+ self.allocate_size = allocate_size
+ """
+ None = adaptive: Memory capacity is divided over all keys in the replay memory.
+ This exploits all the memory capacity over the experiences/classes (keys).
+ int = fixed: All keys in the replay memory get a fixed memory capacity.
+ Useful when experiences/classes known beforehand.
+ """
+ self.adaptive_size = self.allocate_size < 0
+ """
+ Memory capacity is divided over all keys in the replay memory.
+ """
+ self.fixed_size = self.allocate_size > 0
+ """
+ All keys in the replay memory get a fixed memory capacity.
+ Useful when experiences/classes known beforehand.
+ """
+
+ def __call__(self, strategy, **kwargs):
+ """
+ :param strategy:
+ :param kwargs:
+ :return:
+ """
curr_task_id = strategy.experience.task_label
curr_data = strategy.experience.dataset
@@ -98,3 +161,6 @@ class ReplayPlugin(StrategyPlugin):
[rem_len, current_mem_size])
ext_mem[task_id] = saved_part
self.ext_mem = ext_mem
+
+class ClassBalancedStoragePolicy(StoragePolicy):
+ pass
\ No newline at end of file
