forked from SamsungSAILMontreal/TinyRecursiveModels
-
Notifications
You must be signed in to change notification settings - Fork 0
Expand file tree
/
Copy pathcustom_train.py
More file actions
40 lines (32 loc) · 1.54 KB
/
Copy pathcustom_train.py
File metadata and controls
40 lines (32 loc) · 1.54 KB
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
for i, (input, target) in enumerate(zip(input_sequence, target_sequence, strict=True)):
memories_wrt.append(memory.detach().requires_grad_(True))
memory = model(input.unsqueeze(0).to(dtype=torch.float32), memories_wrt[-1])
output = head(memory)
loss = criterion(output, target.unsqueeze(0).to(dtype=torch.float32))
memories.append(memory)
transform = Accumulate() << Aggregate(aggregator, OrderedSet(list(model.parameters())))
print(f"{loss.item():.1e}")
if (i + 1) % update_every == 0:
grad_output = torch.autograd.grad(loss, [memories[-1]], retain_graph=True)
for j in range(update_every):
grads = torch.autograd.grad(
memories[-j - 1],
list(model.parameters()) + [memories_wrt[-j - 1]],
grad_outputs=grad_output,
)
grads_wrt_params = grads[:-1]
grad_output = grads[-1]
for param, grad in zip(model.parameters(), grads_wrt_params, strict=True):
param_to_gradients[param].append(grad)
param_to_jacobian = {
param: torch.stack(gradients, dim=0) for param, gradients in param_to_gradients.items()
}
optimizer.zero_grad()
transform(param_to_jacobian) # This stores the aggregated Jacobian in the .grad fields
optimizer.step()
memories = []
memories_wrt = []
param_to_gradients = defaultdict(list)
head_optimizer.zero_grad()
torch.autograd.backward(loss, inputs=list(head.parameters()))
head_optimizer.step()