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
|
"""Standard (layer-wise) feedback alignment: Linear / Conv2d layers whose backward uses a fixed random
weight B in place of W^T (Lillicrap et al. 2016). Nonlinearity derivatives are the true ones, taken at
the forward activations, as in the original algorithm. Use `apply_fa(module)` to convert all Linear/Conv2d
layers of a block in place."""
import math
import torch
import torch.nn as nn
import torch.nn.functional as F
class _FALinearFn(torch.autograd.Function):
@staticmethod
def forward(ctx, x, W, b, B):
ctx.save_for_backward(x, B)
ctx.has_bias = b is not None
y = x.matmul(W.t())
return y + b if b is not None else y
@staticmethod
def backward(ctx, gy):
x, B = ctx.saved_tensors
gx = gy.matmul(B) # random feedback instead of W
gW = gy.reshape(-1, gy.shape[-1]).t().matmul(x.reshape(-1, x.shape[-1]))
gb = gy.reshape(-1, gy.shape[-1]).sum(0) if ctx.has_bias else None
return gx, gW, gb, None
class _FAConvFn(torch.autograd.Function):
@staticmethod
def forward(ctx, x, W, b, B, stride, padding):
ctx.save_for_backward(x, W, B)
ctx.stride, ctx.padding = stride, padding
ctx.has_bias = b is not None
return F.conv2d(x, W, b, stride, padding)
@staticmethod
def backward(ctx, gy):
x, W, B = ctx.saved_tensors
gx = torch.nn.grad.conv2d_input(x.shape, B, gy, ctx.stride, ctx.padding) # random feedback kernels
gW = torch.nn.grad.conv2d_weight(x, W.shape, gy, ctx.stride, ctx.padding)
gb = gy.sum((0, 2, 3)) if ctx.has_bias else None
return gx, gW, gb, None, None, None
class FALinear(nn.Linear):
def __init__(self, *a, **k):
super().__init__(*a, **k)
self.register_buffer("B", torch.randn_like(self.weight) / math.sqrt(self.in_features))
def forward(self, x):
return _FALinearFn.apply(x, self.weight, self.bias, self.B)
class FAConv2d(nn.Conv2d):
def __init__(self, *a, **k):
super().__init__(*a, **k)
fan_in = self.in_channels * self.kernel_size[0] * self.kernel_size[1]
self.register_buffer("B", torch.randn_like(self.weight) / math.sqrt(fan_in))
def forward(self, x):
return _FAConvFn.apply(x, self.weight, self.bias, self.B, self.stride, self.padding)
def apply_fa(module):
"""Replace every nn.Linear / nn.Conv2d inside `module` (recursively) by its FA variant, keeping weights."""
for name, child in list(module.named_children()):
if type(child) is nn.Linear:
new = FALinear(child.in_features, child.out_features, bias=child.bias is not None)
new.weight = child.weight
if child.bias is not None:
new.bias = child.bias
new.B = new.B.to(child.weight.device)
setattr(module, name, new)
elif type(child) is nn.Conv2d:
new = FAConv2d(child.in_channels, child.out_channels, child.kernel_size, child.stride, child.padding,
bias=child.bias is not None)
new.weight = child.weight
if child.bias is not None:
new.bias = child.bias
new.B = new.B.to(child.weight.device)
setattr(module, name, new)
else:
apply_fa(child)
return module
|
