-
Notifications
You must be signed in to change notification settings - Fork 4
/
continual_learner.py
145 lines (114 loc) · 6.42 KB
/
continual_learner.py
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
import abc
import numpy as np
import torch
from torch import nn
from torch.nn import functional as F
import utils
class ContinualLearner(nn.Module, metaclass=abc.ABCMeta):
'''Abstract module to add continual learning capabilities to a classifier.'''
def __init__(self):
super().__init__()
#----------------- EWC-specifc parameters -----------------#
self.ewc = False
self.ewc_lambda = 5000 #-> hyperparam: how strong to weigh EWC-loss ("regularisation strength")
self.gamma = 1. #-> hyperparam (online EWC): decay-term for old tasks' contribution to quadratic term
self.online = True #-> "online" (=single quadratic term) or "offline" (=quadratic term per task) EWC
self.fisher_n = None #-> sample size for estimating FI-matrix (if "None", full pass over dataset)
self.emp_FI = False #-> if True, use provided labels to calculate FI ("empirical FI"); else predicted labels
self.EWC_task_count = 0 #-> keeps track of number of quadratic loss terms (for "offline EWC")
#----------------- Distillation-specifc parameters -----------------#
self.distill = False
self.KD_temp = 2.0
#----------------- EWC-specifc functions -----------------#s
def estimate_fisher(self, dataset, allowed_classes=None, collate_fn=None):
'''After completing training on a task, estimate diagonal of Fisher Information matrix.
[dataset]: <DataSet> to be used to estimate FI-matrix
[allowed_classes]: <list> with class-indeces of 'allowed' or 'active' classes'''
# Prepare <dict> to store estimated Fisher Information matrix
est_fisher_info = {}
for n, p in self.named_parameters():
if p.requires_grad:
n = n.replace('.', '__')
est_fisher_info[n] = p.detach().clone().zero_()
# Set model to evaluation mode
mode = self.training
self.eval()
# Create data-loader to give batches of size 1
data_loader = utils.get_data_loader(dataset, batch_size=1, cuda=self._is_on_cuda(), collate_fn=collate_fn)
# Estimate the FI-matrix for [self.fisher_n] batches of size 1
for index,(x,y) in enumerate(data_loader):
# break from for-loop if max number of samples has been reached
if self.fisher_n is not None:
if index >= self.fisher_n:
break
# run forward pass of model
x = x.to(self._device())
output = self(x) if allowed_classes is None else self(x)[:, allowed_classes]
if self.emp_FI:
# -use provided label to calculate loglikelihood --> "empirical Fisher":
label = torch.LongTensor([y]) if type(y)==int else y
if allowed_classes is not None:
label = [int(np.where(i == allowed_classes)[0][0]) for i in label.numpy()]
label = torch.LongTensor(label)
label = label.to(self._device())
else:
# -use predicted label to calculate loglikelihood:
label = output.max(1)[1]
# calculate negative log-likelihood
negloglikelihood = F.nll_loss(F.log_softmax(output, dim=1), label)
# Calculate gradient of negative loglikelihood
self.zero_grad()
negloglikelihood.backward()
# Square gradients and keep running sum
for n, p in self.named_parameters():
if p.requires_grad:
n = n.replace('.', '__')
if p.grad is not None:
est_fisher_info[n] += p.grad.detach() ** 2
# Normalize by sample size used for estimation
est_fisher_info = {n: p/index for n, p in est_fisher_info.items()}
# Store new values in the network
for n, p in self.named_parameters():
if p.requires_grad:
n = n.replace('.', '__')
# -mode (=MAP parameter estimate)
self.register_buffer('{}_EWC_prev_task{}'.format(n, "" if self.online else self.EWC_task_count+1),
p.detach().clone())
# -accuracy (approximated by diagonal Fisher Information matrix)
if self.online and self.EWC_task_count==1:
existing_values = getattr(self, '{}_EWC_estimated_fisher'.format(n))
est_fisher_info[n] += self.gamma * existing_values
self.register_buffer('{}_EWC_estimated_fisher{}'.format(n, "" if self.online else self.EWC_task_count+1),
est_fisher_info[n])
# If "offline EWC", increase task-count (for "online EWC", set it to 1 to indicate EWC-loss can be calculated)
self.EWC_task_count = 1 if self.online else self.EWC_task_count + 1
# Set model back to its initial mode
self.train(mode=mode)
def ewc_loss(self):
'''Calculate EWC-loss.'''
if self.EWC_task_count>0:
losses = []
# If "offline EWC", loop over all previous tasks (if "online EWC", [EWC_task_count]=1 so only 1 iteration)
for task in range(1, self.EWC_task_count+1):
for n, p in self.named_parameters():
if p.requires_grad:
# Retrieve stored mode (MAP estimate) and accuracy (Fisher Information matrix)
n = n.replace('.', '__')
mean = getattr(self, '{}_EWC_prev_task{}'.format(n, "" if self.online else task))
fisher = getattr(self, '{}_EWC_estimated_fisher{}'.format(n, "" if self.online else task))
# If "online EWC", apply decay-term to the running sum of the Fisher Information matrices
fisher = self.gamma*fisher if self.online else fisher
# Calculate EWC-loss
losses.append((fisher * (p-mean)**2).sum())
# Sum EWC-loss from all parameters (and from all tasks, if "offline EWC")
return (1./2)*sum(losses)
else:
# EWC-loss is 0 if there are no stored mode and accuracy yet
return torch.tensor(0., device=self._device())
def _device(self):
return next(self.parameters()).device
def _is_on_cuda(self):
return next(self.parameters()).is_cuda
@abc.abstractmethod
def forward(self, x):
pass