Spaces:
Runtime error
Runtime error
File size: 4,088 Bytes
e56055d |
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 |
import torch
import torch.nn as nn
import higher
from editable_model import EditableModel
from utils import _logits
def fomaml_callback(all_grads):
return [g.detach() if g is not None else None for g in all_grads]
class ENN(EditableModel):
def __init__(self, model, config, model_constructor, edit_lrs=None, edit_loss_fn=None):
super().__init__(model, config, model_constructor)
if edit_lrs is None:
edit_lrs = nn.Parameter(torch.tensor([config.edit_lr] * len(self.config.model.inner_params)))
self.edit_lrs = edit_lrs
if edit_loss_fn is not None:
self.edit_loss_fn = edit_loss_fn
self.grad_callback = fomaml_callback if config.enn.first_order else lambda x: x
def outer_parameters(self, grouped=False):
extra_params = [self.edit_lrs]
if self.config.no_grad_layers is None:
model_params = self.model.parameters() if type(self.model.parameters()) == list else list(self.model.parameters())
else:
model_params = []
for m in self.model.modules():
if isinstance(m, nn.ModuleList):
model_params.extend(list(m[self.config.no_grad_layers:].parameters()))
if grouped:
return [
dict(params=model_params, lr=self.config.lr),
dict(params=extra_params, lr=self.config.lr_lr)
]
else:
return model_params + extra_params
def get_state_dict(self):
return self.state_dict()
def edit(self, batch, condition=None, detach_history=False):
opt = torch.optim.SGD([{"params": p, "lr": None}
for (n, p) in self.model.named_parameters() if n in self.config.model.inner_params])
with torch.enable_grad(), higher.innerloop_ctx(
self.model,
opt,
override={'lr': list(self.edit_lrs)},
copy_initial_weights=False,
track_higher_grads=self.training,
in_place=True
) as (fmodel, diffopt):
fmodel.eval()
for edit_step in range(self.config.enn.n_edit_steps):
output = _logits(fmodel(**batch))
loss = self.edit_loss_fn(output, batch["labels"])["nll"]
diffopt.step(loss, grad_callback=self.grad_callback)
if not detach_history:
model_edited = fmodel
else:
model_edited = self.model_constructor()
model_edited.load_state_dict(fmodel.state_dict())
model_edited.train(self.training)
return ENN(model_edited, self.config, self.model_constructor, edit_lrs=self.edit_lrs, edit_loss_fn=self.edit_loss_fn), {}
def test():
import transformers
import types
import copy
model = transformers.GPT2LMHeadModel.from_pretrained("gpt2")
config = types.SimpleNamespace()
config.edit_lr = 0.1
config.model.inner_params = [
"transformer.h.9.mlp.c_fc.weight",
"transformer.h.9.mlp.c_proj.weight",
"transformer.h.10.mlp.c_fc.weight",
"transformer.h.10.mlp.c_proj.weight",
"transformer.h.11.mlp.c_fc.weight",
"transformer.h.11.mlp.c_proj.weight",
]
config.enn = {
"n_edit_steps": 2,
"first_order": False
}
enn = ENN(model, config, lambda: copy.deepcopy(model)).cuda()
x = torch.arange(100).view(5, 20).cuda() + 1000
edited = enn.edit(x, masks=torch.ones_like(x), labels=x)
orig_param = [p for (n, p) in enn.model.named_parameters() if n == config.model.inner_params[-1]][0]
edited_param = [p for (n, p) in edited.model.named_parameters() if n == config.model.inner_params[-1]][0]
print((orig_param - edited_param).abs().max())
edited.eval()
print(enn(x, labels=x).loss, edited(x, labels=x).loss, edited.edit_loss_fn(edited(x).logits, x)["nll"])
edited.edit_loss_fn(edited(x).logits, x).backward()
import pdb; pdb.set_trace()
if __name__ == '__main__':
with torch.autograd.set_detect_anomaly(True):
test()
|