core/model/ider.py

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()}]