Skip to content

Add LayerNorm in 3C1F for Fashion-MNIST. #1

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
b-mu wants to merge 9 commits into
base: master
Choose a base branch
from
Open

Add LayerNorm in 3C1F for Fashion-MNIST. #1

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
9 commits
Select commit Hold shift + click to select a range
53e2431
Add LayerNorm in 3C1F for Fashion-MNIST.
b-mu Jun 14, 2021
26738a2
Try take mean and var over surface only (not volume) for 4D tensor.
b-mu Jun 14, 2021
781daaf
Add LayerNorm to main loop.
b-mu Jun 14, 2021
8825596
Add bias.
b-mu Jun 16, 2021
950a085
Update model.
b-mu Jun 17, 2021
f94a1b2
Update model.
b-mu Jun 17, 2021
9178beb
Save Fisher inv block of LayerNorm.
b-mu Jun 17, 2021
b6a1ab8
Add save NGD_kernel option.
b-mu Jun 19, 2021
8103e98
Try no learnable parameter for LayerNorm.
b-mu Jun 19, 2021
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
87 changes: 83 additions & 4 deletions main.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -50,6 +50,7 @@
parser.add_argument('--load_path', default='', type=str)
parser.add_argument('--log_dir', default='runs/pretrain', type=str)
parser.add_argument('--save_inv', default='false', type=str)
parser.add_argument('--save_kernel', default='false', type=str)


parser.add_argument('--optimizer', default='kfac', type=str)
Expand Down Expand Up @@ -182,6 +183,8 @@
buf[name] = torch.zeros_like(param.data).to(args.device)
if args.save_inv == 'true':
os.mkdir('ngd')
if args.save_kernel == 'true':
os.mkdir('ngd_kernel')

elif optim_name == 'exact_ngd':
print('Exact NGD optimizer selected.')
Expand Down Expand Up @@ -457,7 +460,7 @@ def closure():
sampled_y = torch.multinomial(torch.nn.functional.softmax(outputs, dim=1),1).squeeze().to(args.device)

if args.trial == 'true':
update_list, loss = optimal_JJT_v2(outputs, sampled_y, args.batch_size, damping=damp, alpha=0.95, low_rank=args.low_rank, gamma=args.gamma, memory_efficient=args.memory_efficient, super_opt=args.super_opt)
update_list, loss = optimal_JJT_v2(outputs, sampled_y, args.batch_size, damping=damp, alpha=0.95, low_rank=args.low_rank, gamma=args.gamma, memory_efficient=args.memory_efficient, super_opt=args.super_opt, save_kernel=args.save_kernel)
else:
update_list, loss = optimal_JJT(outputs, sampled_y, args.batch_size, damping=damp, alpha=0.95, low_rank=args.low_rank, gamma=args.gamma, memory_efficient=args.memory_efficient)

Expand Down Expand Up @@ -587,7 +590,49 @@ def closure():
gv = gv / n
update = (grad - gv)/damp
m.weight.grad.copy_(update)

elif isinstance(m, nn.LayerNorm):
I, G = m.I, m.G
if len(I.shape) == 2:
mean = I.mean(dim=-1).unsqueeze(-1)
var = I.var(dim=-1, unbiased=False).unsqueeze(-1)
else:
mean = I.mean((-2, -1), keepdims=True)
var = I.var((-2, -1), unbiased=False, keepdims=True)
x_hat = (I - mean) / (var + m.eps).sqrt()

J = G * x_hat
J = J.reshape(J.shape[0], -1)
JJT = torch.matmul(J, J.t())

grad_prod = torch.matmul(J, grad.reshape(-1))

NGD_kernel = JJT / n
NGD_inv = torch.linalg.inv(NGD_kernel + damp * torch.eye(n).to(grad.device))
v = torch.matmul(NGD_inv, grad_prod)

gv = torch.matmul(J.t(), v) / n

update = (grad.reshape(-1) - gv) / damp
update = update.reshape(m.weight.grad.shape)
m.weight.grad.copy_(update)

grad = m.bias.grad.reshape(-1)

J = G
J = J.reshape(J.shape[0], -1)
JJT = torch.matmul(J, J.t())

grad_prod = torch.matmul(J, grad)

NGD_kernel = JJT / n
NGD_inv = torch.linalg.inv(NGD_kernel + damp * torch.eye(n).to(grad.device))
v = torch.matmul(NGD_inv, grad_prod)

gv = torch.matmul(J.t(), v) / n

update = (grad - gv) / damp
update = update.reshape(m.bias.grad.shape)
m.bias.grad.copy_(update)


# last part of SMW formula
Expand Down Expand Up @@ -654,6 +699,22 @@ def closure():
train_loss = train_loss/(batch_idx + 1)
if args.step_info == 'true':
TRAIN_INFO['epoch_time'].append(float("{:.4f}".format(epoch_time)))

# save NGD kernels
if args.save_kernel == 'true' and optim_name == 'ngd':
if module_names == 'children':
all_modules = net.children()
elif module_names == 'features':
all_modules = net.features.children()

count = 0
for m in all_modules:
if m.__class__.__name__ in ['Linear', 'Conv2d', 'LayerNorm']:
if hasattr(m, "NGD_kernel"):
with open('ngd_kernel/' + str(epoch) + '_m_' + str(count) + '_kernel.npy', 'wb') as f:
np.save(f, m.NGD_kernel.cpu().numpy())
count += 1

# save diagonal blocks of exact Fisher inverse or its approximations
if args.save_inv == 'true':
if module_names == 'children':
Expand Down Expand Up @@ -686,6 +747,24 @@ def closure():
np.save(f, ((torch.eye(JTDJ.size(0)).to(JTDJ.device) - JTDJ) / damping).cpu().numpy())
count += 1

elif m.__class__.__name__ == 'LayerNorm':
with torch.no_grad():
I, G = m.I, m.G
if len(I.shape) == 2:
mean = I.mean(dim=-1).unsqueeze(-1)
var = I.var(dim=-1, unbiased=False).unsqueeze(-1)
else:
mean = I.mean((-2, -1), keepdims=True)
var = I.var((-2, -1), unbiased=False, keepdims=True)
x_hat = (I - mean) / (var + m.eps).sqrt()

J = G * x_hat
J = J.reshape(J.shape[0], -1)
JTDJ = torch.matmul(J.t(), torch.matmul(m.NGD_inv, J)) / args.batch_size

with open('ngd/' + str(epoch) + '_m_' + str(count) + '_inv.npy', 'wb') as f:
np.save(f, ((torch.eye(JTDJ.size(0)).to(JTDJ.device) - JTDJ) / damping).cpu().numpy())
count += 1
elif optim_name == 'exact_ngd':
for m in all_modules:
if m.__class__.__name__ in ['Conv2d', 'Linear']:
Expand Down Expand Up @@ -775,11 +854,11 @@ def optimal_JJT(outputs, targets, batch_size, damping=1.0, alpha=0.95, low_rank=
update_list[name] = fisher_vals[2]
return update_list, loss

def optimal_JJT_v2(outputs, targets, batch_size, damping=1.0, alpha=0.95, low_rank='false', gamma=0.95, memory_efficient='false', super_opt='false'):
def optimal_JJT_v2(outputs, targets, batch_size, damping=1.0, alpha=0.95, low_rank='false', gamma=0.95, memory_efficient='false', super_opt='false', save_kernel='false'):
jac_list = 0
vjp = 0
update_list = {}
with backpack(FisherBlockEff(damping, alpha, low_rank, gamma, memory_efficient, super_opt)):
with backpack(FisherBlockEff(damping, alpha, low_rank, gamma, memory_efficient, super_opt, save_kernel)):
loss = criterion(outputs, targets)
loss.backward()
for name, param in net.named_parameters():
Expand Down
7 changes: 7 additions & 0 deletions models/mnist/convnet.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -12,14 +12,21 @@ def __init__(self, num_classes=10, **kwargs):
super(ConvNet, self).__init__()
self.features = nn.Sequential(
nn.Conv2d(1, 128, kernel_size=3, stride=1, padding=1),
# nn.LayerNorm([128, 28, 28], elementwise_affine=False),
nn.LayerNorm([28, 28], elementwise_affine=False),
nn.ReLU(),
nn.Conv2d(128, 128, kernel_size=3, stride=1, padding=1),
# nn.LayerNorm([128, 28, 28], elementwise_affine=False),
nn.LayerNorm([28, 28], elementwise_affine=False),
nn.ReLU(),
nn.Conv2d(128, 128, kernel_size=3, stride=1, padding=1),
# nn.LayerNorm([128, 28, 28], elementwise_affine=False),
nn.LayerNorm([28, 28], elementwise_affine=False),
nn.ReLU(),
nn.MaxPool2d(kernel_size=3),
nn.Flatten(),
nn.Linear(9*9*128, 500),
nn.LayerNorm([500], elementwise_affine=False),
nn.ReLU(),
nn.Linear(500, 10),
)
Expand Down
3 changes: 3 additions & 0 deletions models/mnist/toy.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -11,12 +11,15 @@ def __init__(self, num_classes=10, **kwargs):
super(ToyNet, self).__init__()
self.features = nn.Sequential(
nn.Conv2d(1, 16, kernel_size=3, stride=1, padding=1),
nn.LayerNorm([16, 28, 28]),
nn.ReLU(),
nn.MaxPool2d(kernel_size=3),
nn.Conv2d(16, 16, kernel_size=3, stride=1, padding=1),
nn.LayerNorm([16, 9, 9]),
nn.ReLU(),
nn.MaxPool2d(kernel_size=3),
nn.Conv2d(16, 16, kernel_size=3, stride=1, padding=1),
nn.LayerNorm([16, 3, 3]),
nn.ReLU(),
nn.Flatten(),
nn.Linear(3*3*16, 10)
Expand Down