File size: 5,595 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 | import torch
import torch.nn as nn
import torch.nn.functional as F
from .finetune import Finetune
class IncrementalLinearNet(nn.Module):
def __init__(self, backbone, feat_dim, num_class):
super().__init__()
self.backbone = backbone
self.classifier = nn.Linear(feat_dim, num_class)
def _features(self, x):
output = self.backbone(x)
if isinstance(output, dict):
return output["features"]
return output
def forward(self, x):
return self.classifier(self._features(x))
def update_classifier(self, num_class):
old_classifier = self.classifier
if old_classifier.out_features == num_class:
return
new_classifier = nn.Linear(old_classifier.in_features, num_class)
new_classifier = new_classifier.to(old_classifier.weight.device)
with torch.no_grad():
old_out = old_classifier.out_features
new_classifier.weight[:old_out].copy_(old_classifier.weight)
new_classifier.bias[:old_out].copy_(old_classifier.bias)
self.classifier = new_classifier
class EWCDR(Finetune):
"""EWC Done Right with logits reversal during importance estimation."""
def __init__(self, backbone, feat_dim, num_class, **kwargs):
super().__init__(backbone, feat_dim, num_class, **kwargs)
self.init_cls_num = kwargs["init_cls_num"]
self.inc_cls_num = kwargs["inc_cls_num"]
self.lamda = kwargs["lamda"]
self.omega_max = kwargs.get("omega_max", 1e-4)
self.task_idx = 0
self.known_cls_num = 0
self.total_cls_num = self.init_cls_num
self.network = IncrementalLinearNet(backbone, feat_dim, self.init_cls_num)
self.ref_param = {}
self.omega = None
def before_task(self, task_idx, buffer, train_loader, test_loaders):
self.task_idx = task_idx
self.known_cls_num = (
0 if task_idx == 0 else self.init_cls_num + (task_idx - 1) * self.inc_cls_num
)
self.total_cls_num = self.init_cls_num + task_idx * self.inc_cls_num
self.network.update_classifier(self.total_cls_num)
self.network.to(self.device)
def observe(self, data):
x = data["image"].to(self.device)
y = data["label"].to(self.device)
logits = self.network(x)
if self.task_idx == 0:
loss = F.cross_entropy(logits, y)
else:
loss = F.cross_entropy(
logits[:, self.known_cls_num:], y - self.known_cls_num
)
loss = loss + self.lamda * self.compute_ewc()
pred = torch.argmax(logits, dim=1)
acc = torch.sum(pred == y).item()
return pred, acc / x.size(0), loss
def inference(self, data):
x = data["image"].to(self.device)
y = data["label"].to(self.device)
logits = self.network(x)
pred = torch.argmax(logits, dim=1)
acc = torch.sum(pred == y).item()
return pred, acc / x.size(0)
def after_task(self, task_idx, buffer, train_loader, test_loaders):
new_omega = self.get_importance(train_loader)
if self.omega is not None:
alpha = self.known_cls_num / self.total_cls_num
for name, old_omega in self.omega.items():
if name not in new_omega:
continue
old_slice = self._prefix_slice(old_omega)
new_omega[name][old_slice] = (
alpha * old_omega + (1 - alpha) * new_omega[name][old_slice]
)
self.omega = new_omega
self.ref_param = {
name: param.clone().detach()
for name, param in self.network.named_parameters()
if param.requires_grad
}
def get_importance(self, train_loader):
omega = {
name: torch.zeros_like(param, device=self.device)
for name, param in self.network.named_parameters()
if param.requires_grad
}
was_training = self.network.training
self.network.train()
for data in train_loader:
x = data["image"].to(self.device)
y = data["label"].to(self.device)
logits = -self.network(x)
loss = F.cross_entropy(logits, y)
self.network.zero_grad(set_to_none=True)
loss.backward()
for name, param in self.network.named_parameters():
if param.grad is not None and name in omega:
omega[name] += param.grad.pow(2).detach()
for name, value in omega.items():
value = value / len(train_loader)
omega[name] = torch.clamp(value, max=self.omega_max)
self.network.zero_grad(set_to_none=True)
self.network.train(was_training)
return omega
def compute_ewc(self):
if self.omega is None:
return torch.zeros((), device=self.device)
loss = torch.zeros((), device=self.device)
for name, param in self.network.named_parameters():
if name not in self.omega:
continue
ref = self.ref_param[name]
old_slice = self._prefix_slice(ref)
loss += torch.sum(self.omega[name] * (param[old_slice] - ref).pow(2)) / 2
return loss
@staticmethod
def _prefix_slice(tensor):
return (slice(0, tensor.shape[0]),) + (slice(None),) * (tensor.dim() - 1)
def forward(self, x):
return self.network(x)
def get_parameters(self, config):
return [{"params": self.network.parameters()}]
|