| import copy |
| import os |
| from collections import Counter |
|
|
| import numpy as np |
| import torch |
| import torch.nn as nn |
| import torch.nn.functional as F |
| from PIL import Image |
| from torchvision import transforms |
|
|
| from .finetune import Finetune |
|
|
|
|
| class IDER(Finetune): |
| """Idempotent Experience Replay for LibContinual. |
| |
| This follows the official implementation: the current batch receives the |
| standard idempotent CE loss, while replay samples are drawn inside |
| ``observe`` for supervised ER and old-checkpoint idempotent distillation. |
| """ |
|
|
| def __init__(self, backbone, feat_dim, num_class, **kwargs): |
| super().__init__(backbone, feat_dim, num_class, **kwargs) |
| self.classifier = nn.Identity() |
| self.weighta = kwargs.get("weighta", 0.5) |
| self.weightb = kwargs.get("weightb", 0.4) |
| self.weightc = kwargs.get("weightc", 0.5) |
| self.empty_prob = kwargs.get("weightmask", kwargs.get("empty_prob", 1.0)) |
| self.refine_inference = kwargs.get("refine_inference", False) |
| self.mask_unseen = kwargs.get("mask_unseen", False) |
| self.old_backbone = None |
| self.old_seen_cls_num = 0 |
| self.seen_cls_num = num_class |
| self.task_idx = 0 |
| self.buffer = None |
| self.buffer_dataset = None |
| self.minibatch_size = kwargs.get("minibatch_size", kwargs.get("buffer_batch_size", 32)) |
| self.online_buffer_update = kwargs.get("online_buffer_update", True) |
| self.class_balance = kwargs.get("class_balance", True) |
| self.print_buffer_stats = kwargs.get("print_buffer_stats", False) |
| self.buffer_update_after_step = kwargs.get("buffer_update_after_step", False) |
| self.buffer_storage = kwargs.get("buffer_storage", "raw") |
| self.ce_detach_signal = kwargs.get("ce_detach_signal", True) |
| self.distill_detach_signal = kwargs.get("distill_detach_signal", True) |
| self.old_model_eval = kwargs.get("old_model_eval", True) |
| self.freeze_old_model = kwargs.get("freeze_old_model", True) |
| self._pending_buffer_data = None |
| self._pending_buffer_labels = None |
| self._to_tensor = transforms.ToTensor() |
| self._to_pil = transforms.ToPILImage() |
|
|
| def _empty_signal(self, batch_size, upto=None): |
| upto = self.num_class if upto is None else upto |
| if hasattr(self.backbone, "empty_signal"): |
| return self.backbone.empty_signal(batch_size, self.device, upto=upto) |
| signal = torch.zeros(batch_size, self.num_class, device=self.device) |
| signal[:, :upto] = 1.0 / max(upto, 1) |
| return signal |
|
|
| def _one_hot(self, labels): |
| return F.one_hot(labels, num_classes=self.num_class).float() |
|
|
| def _sample_second_input(self, labels): |
| empty = self._empty_signal(labels.size(0)) |
| one_hot = self._one_hot(labels) |
| if self.empty_prob >= 1: |
| return empty |
| if self.empty_prob <= 0: |
| return one_hot |
| use_empty = (torch.rand(1, device=self.device) <= self.empty_prob).item() |
| return empty if use_empty else one_hot |
|
|
| def _mask_logits(self, logits, upto=None): |
| if not self.mask_unseen: |
| return logits |
| upto = self.seen_cls_num if upto is None else upto |
| if upto >= logits.size(1): |
| return logits |
| masked = logits.clone() |
| masked[:, upto:] = -1e9 |
| return masked |
|
|
| def _backbone_logits(self, x, y_signal=None, upto=None, backbone=None): |
| backbone = self.backbone if backbone is None else backbone |
| try: |
| output = backbone(x, y_signal, upto=upto) |
| except TypeError: |
| output = backbone(x, y_signal) |
|
|
| if isinstance(output, dict): |
| return output["logits"] |
| if isinstance(output, tuple): |
| return output[0] |
| return output |
|
|
| def _refined_logits(self, x): |
| empty = self._empty_signal(x.size(0)) |
| logits0 = self._backbone_logits(x, empty) |
| logits0 = self._mask_logits(logits0) |
| if not self.refine_inference: |
| return logits0 |
|
|
| probs0 = F.softmax(logits0, dim=1).detach() |
| logits1 = self._backbone_logits(x, probs0) |
| return self._mask_logits(logits1) |
|
|
| def before_task(self, task_idx, buffer, train_loader, test_loaders): |
| self.task_idx = task_idx |
| self.seen_cls_num = self.kwargs["init_cls_num"] + task_idx * self.kwargs["inc_cls_num"] |
| self.seen_cls_num = min(self.seen_cls_num, self.num_class) |
| self.buffer = buffer |
| self.buffer_dataset = train_loader.dataset if train_loader is not None else None |
| if buffer is not None and hasattr(buffer, "batch_size"): |
| self.minibatch_size = buffer.batch_size |
| if self.old_backbone is not None: |
| self.old_backbone = self.old_backbone.to(self.device) |
| if self.old_model_eval: |
| self.old_backbone.eval() |
| else: |
| self.old_backbone.train() |
|
|
| def _load_buffer_image(self, image_ref): |
| dataset = self.buffer_dataset |
| if dataset is None: |
| raise RuntimeError("IDER needs the train dataset to sample replay images.") |
|
|
| if torch.is_tensor(image_ref): |
| image = self._to_pil(image_ref.cpu()) |
| elif dataset.dataset in ["binary_cifar10", "binary_cifar100"]: |
| image = Image.fromarray(np.uint8(image_ref)) |
| elif dataset.dataset == "tiny-imagenet": |
| image = Image.open(image_ref).convert("RGB") |
| elif dataset.dataset == "processed_tinyimg": |
| if np.max(image_ref) <= 1.0: |
| image_ref = np.uint8(255 * image_ref) |
| image = Image.fromarray(np.uint8(image_ref)).convert("RGB") |
| else: |
| image = Image.open(os.path.join(dataset.data_root, dataset.mode, image_ref)).convert("RGB") |
| return dataset.trfms(image) |
|
|
| def _make_buffer_image_ref(self, image_ref): |
| if self.buffer_storage != "tensor": |
| return image_ref |
|
|
| if torch.is_tensor(image_ref): |
| return image_ref.detach().cpu() |
|
|
| dataset = self.buffer_dataset |
| if dataset.dataset in ["binary_cifar10", "binary_cifar100"]: |
| image = Image.fromarray(np.uint8(image_ref)) |
| elif dataset.dataset == "tiny-imagenet": |
| image = Image.open(image_ref).convert("RGB") |
| elif dataset.dataset == "processed_tinyimg": |
| if np.max(image_ref) <= 1.0: |
| image_ref = np.uint8(255 * image_ref) |
| image = Image.fromarray(np.uint8(image_ref)).convert("RGB") |
| else: |
| image = Image.open(os.path.join(dataset.data_root, dataset.mode, image_ref)).convert("RGB") |
| return self._to_tensor(image).cpu() |
|
|
| def _sample_buffer_batch(self): |
| if self.buffer is None or self.buffer.is_empty(): |
| return None |
|
|
| n_items = len(self.buffer.labels) |
| batch_size = min(self.minibatch_size, n_items) |
| indices = np.random.choice(n_items, size=batch_size, replace=False) |
| images = [self._load_buffer_image(self.buffer.images[idx]) for idx in indices] |
| labels = [self.buffer.labels[idx] for idx in indices] |
| return torch.stack(images).to(self.device), torch.tensor(labels, dtype=torch.long, device=self.device) |
|
|
| def _online_update_buffer(self, data, labels): |
| if ( |
| not self.online_buffer_update |
| or self.buffer is None |
| or self.buffer.buffer_size <= 0 |
| or self.buffer_dataset is None |
| or "index" not in data |
| ): |
| return |
|
|
| if not hasattr(self.buffer, "num_seen_examples"): |
| self.buffer.num_seen_examples = len(self.buffer.labels) |
|
|
| indices = data["index"].detach().cpu().numpy().tolist() |
| labels = labels.detach().cpu().numpy().tolist() |
|
|
| for dataset_idx, label in zip(indices, labels): |
| image_ref = self._make_buffer_image_ref(self.buffer_dataset.images[int(dataset_idx)]) |
| seen = self.buffer.num_seen_examples |
| if len(self.buffer.labels) < self.buffer.buffer_size: |
| self.buffer.images.append(image_ref) |
| self.buffer.labels.append(int(label)) |
| else: |
| replace_idx = np.random.randint(0, seen + 1) |
| if replace_idx < self.buffer.buffer_size: |
| if self.class_balance and len(self.buffer.labels) > 0: |
| counts = Counter(self.buffer.labels) |
| max_count = max(counts.values()) |
| majority_classes = {cls for cls, cnt in counts.items() if cnt == max_count} |
| majority_indices = [ |
| idx for idx, old_label in enumerate(self.buffer.labels) |
| if old_label in majority_classes |
| ] |
| replace_idx = int(np.random.choice(majority_indices)) |
| self.buffer.images[replace_idx] = image_ref |
| self.buffer.labels[replace_idx] = int(label) |
| self.buffer.num_seen_examples += 1 |
|
|
| def _idempotent_ce_loss(self, x, y): |
| y_signal = self._sample_second_input(y) |
| logits1 = self._backbone_logits(x, y_signal) |
| signal_logits = logits1.detach() if self.ce_detach_signal else logits1 |
| logits2 = self._backbone_logits(x, F.softmax(signal_logits, dim=1)) |
| loss1 = self.loss_fn(self._mask_logits(logits1), y) |
| loss2 = self.loss_fn(self._mask_logits(logits2), y) |
| return 0.5 * (loss1 + loss2), logits1, logits2, y_signal |
|
|
| def _distill_loss(self, x): |
| uniform = self._empty_signal(x.size(0)) |
| logits1 = self._backbone_logits(x, uniform) |
| signal_logits = logits1.detach() if self.distill_detach_signal else logits1 |
| if self.old_model_eval: |
| self.old_backbone.eval() |
| else: |
| self.old_backbone.train() |
| logits2 = self._backbone_logits( |
| x, |
| F.softmax(signal_logits, dim=1), |
| backbone=self.old_backbone, |
| ) |
| return F.mse_loss(logits1, logits2) |
|
|
| def observe(self, data): |
| x, y = data["image"], data["label"] |
| x = x.to(self.device) |
| y = y.to(self.device) |
|
|
| loss, logits1, logits2, _ = self._idempotent_ce_loss(x, y) |
|
|
| if self.old_backbone is not None and self.old_seen_cls_num > 0 and self.weightb > 0: |
| loss = loss + self.weightb * self._distill_loss(x) |
|
|
| if self.weightc != 0: |
| buffer_batch = self._sample_buffer_batch() |
| if buffer_batch is not None: |
| buffer_loss, _, _, _ = self._idempotent_ce_loss(*buffer_batch) |
| loss = loss + self.weightc * 2.0 * buffer_loss |
|
|
| if self.old_backbone is not None and self.old_seen_cls_num > 0 and self.weighta != 0: |
| buffer_batch = self._sample_buffer_batch() |
| if buffer_batch is not None: |
| loss = loss + self.weighta * self._distill_loss(buffer_batch[0]) |
|
|
| pred = torch.argmax(self._mask_logits(logits2), dim=1) |
| acc = torch.sum(pred == y).item() |
| if self.buffer_update_after_step: |
| self._pending_buffer_data = data |
| self._pending_buffer_labels = y.detach() |
| else: |
| self._online_update_buffer(data, y) |
| return pred, acc / x.size(0), loss |
|
|
| def after_observe(self): |
| if self._pending_buffer_data is None or self._pending_buffer_labels is None: |
| return |
| self._online_update_buffer(self._pending_buffer_data, self._pending_buffer_labels) |
| self._pending_buffer_data = None |
| self._pending_buffer_labels = None |
|
|
| def _task_class_range(self, task_id): |
| if task_id < 0: |
| return 0, self.num_class |
| if task_id == 0: |
| start = 0 |
| end = self.kwargs["init_cls_num"] |
| else: |
| start = self.kwargs["init_cls_num"] + (task_id - 1) * self.kwargs["inc_cls_num"] |
| end = start + self.kwargs["inc_cls_num"] |
| return start, min(end, self.num_class) |
|
|
| def inference(self, data, task_id=-1): |
| x, y = data["image"], data["label"] |
| x = x.to(self.device) |
| y = y.to(self.device) |
|
|
| logits = self._refined_logits(x) |
| if task_id > -1: |
| start, end = self._task_class_range(task_id) |
| task_logits = torch.full_like(logits, -1e9) |
| task_logits[:, start:end] = logits[:, start:end] |
| logits = task_logits |
| pred = torch.argmax(logits, dim=1) |
| acc = torch.sum(pred == y).item() |
| return pred, acc / x.size(0) |
|
|
| def predict_logits(self, data): |
| x = data["image"].to(self.device) |
| return self._refined_logits(x) |
|
|
| def forward(self, x): |
| return self._refined_logits(x) |
|
|
| def after_task(self, task_idx, buffer, train_loader, test_loaders): |
| if self.print_buffer_stats and buffer is not None and hasattr(buffer, "labels") and len(buffer.labels) > 0: |
| counts = Counter(buffer.labels) |
| seen_classes = self.kwargs["init_cls_num"] + task_idx * self.kwargs["inc_cls_num"] |
| seen_classes = min(seen_classes, self.num_class) |
| values = [counts.get(cls, 0) for cls in range(seen_classes)] |
| zero_classes = sum(1 for value in values if value == 0) |
| print( |
| "[IDER] Buffer class counts after task " |
| f"{task_idx}: size={len(buffer.labels)}, classes={len(counts)}, " |
| f"min={min(values)}, max={max(values)}, zero={zero_classes}, " |
| f"head={values[:min(20, len(values))]}" |
| ) |
| self.old_backbone = copy.deepcopy(self.backbone).to(self.device) |
| if self.old_model_eval: |
| self.old_backbone.eval() |
| else: |
| self.old_backbone.train() |
| for param in self.old_backbone.parameters(): |
| param.requires_grad = not self.freeze_old_model |
| self.old_seen_cls_num = self.seen_cls_num |
|
|
| def get_parameters(self, config): |
| return [{"params": self.backbone.parameters()}] |
|
|
