| import os |
|
|
| import numpy as np |
| import torch |
| from PIL import Image |
| from torchvision import transforms |
|
|
| from .finetune import Finetune |
|
|
|
|
| class PaperER(Finetune): |
| """Online ER baseline used by the original IDER code. |
| |
| Each training step optimizes CE on the current mini-batch concatenated with |
| a random replay mini-batch. The current mini-batch is inserted into the |
| replay memory by reservoir sampling after the optimizer step. |
| """ |
|
|
| def __init__(self, backbone, feat_dim, num_class, **kwargs): |
| super().__init__(backbone, feat_dim, num_class, **kwargs) |
| self.buffer = None |
| self.buffer_dataset = None |
| self.minibatch_size = kwargs.get("minibatch_size", kwargs.get("buffer_batch_size", 32)) |
| self.buffer_storage = kwargs.get("buffer_storage", "tensor") |
| self._pending_buffer_data = None |
| self._pending_buffer_labels = None |
| self._to_tensor = transforms.ToTensor() |
| self._to_pil = transforms.ToPILImage() |
|
|
| def before_task(self, task_idx, buffer, train_loader, test_loaders): |
| 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 |
|
|
| def _logits(self, x): |
| return self.classifier(self.backbone(x)["features"]) |
|
|
| def _load_buffer_image(self, image_ref): |
| dataset = self.buffer_dataset |
| if dataset is None: |
| raise RuntimeError("PaperER 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 |
|
|
| dataset = self.buffer_dataset |
| if torch.is_tensor(image_ref): |
| return image_ref.detach().cpu() |
| 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 |
|
|
| batch_size = min(self.minibatch_size, len(self.buffer.labels)) |
| indices = np.random.choice(len(self.buffer.labels), 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 ( |
| 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: |
| self.buffer.images[replace_idx] = image_ref |
| self.buffer.labels[replace_idx] = int(label) |
| self.buffer.num_seen_examples += 1 |
|
|
| def observe(self, data): |
| x, y = data["image"].to(self.device), data["label"].to(self.device) |
| real_batch_size = x.size(0) |
|
|
| replay_batch = self._sample_buffer_batch() |
| if replay_batch is not None: |
| buf_x, buf_y = replay_batch |
| train_x = torch.cat((x, buf_x), dim=0) |
| train_y = torch.cat((y, buf_y), dim=0) |
| else: |
| train_x, train_y = x, y |
|
|
| logits = self._logits(train_x) |
| loss = self.loss_fn(logits, train_y) |
|
|
| current_logits = logits[:real_batch_size] |
| pred = torch.argmax(current_logits, dim=1) |
| acc = torch.sum(pred == y).item() |
|
|
| self._pending_buffer_data = data |
| self._pending_buffer_labels = y.detach() |
| return pred, acc / real_batch_size, 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 inference(self, data, task_id=-1): |
| x, y = data["image"].to(self.device), data["label"].to(self.device) |
| logits = self._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 _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 predict_logits(self, data): |
| x = data["image"].to(self.device) |
| return self._logits(x) |
|
|
