File size: 6,995 Bytes
3cdd985 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 | 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)
|