具有GPU支持的Levenberg-Marquardt算法

Question

对于浅层神经网络，LM算法做得非常好。

然而，只有MatLab和pyrenn(Python包(似乎有一个健壮的实现。这两种实现的一个问题是它们没有 GPU 支持。我也尝试了neupy(一个python包(，但它并不健壮，当你尝试训练更长的epoch或大型数据集时会失败。你们知道一个好的 NN LM python 包可以使用 GPU 训练吗？

Answer 1

我不确定这样的实现是否对除了最琐碎的网络之外的任何东西都有用，计算 Hessian 是非常不切实际的。然而，我认为在pytorch中为浅层网络实现Levenberg-Marquardt并不是那么难。至少，即使您的实现不是最佳的，与 CPU 版本相比，您仍然可以获得一些加速。这是一个在 GPU 上运行的非常快速且因此非常不完美/次优的示例。对于1080ti上大小为3x20x20的输入，加速比约为 3 倍(我目前没有免费的 GPU 来测试更大的输入(。但是，如果你对损失函数有一个先验(例如，如果它是最小二乘的，并且你知道黑森可以近似为2*Jacobian^t*Jacobian(，那么下面代码的一些变体可能会有用：

import torch
import numpy as np
import functools
import matplotlib.pyplot as plt
def LM(model,loss,n_iter=30):
alpha=1e-3
loss_hist=[]
for i in range(n_iter):
model.train()
out=model(toy_input).unsqueeze(1)
loss_out=loss(out)
prev_loss=loss_out.item()
gradients=torch.autograd.grad(loss_out, model.parameters(), create_graph=True)
model.eval()
Hessian, g_vector=eval_hessian(gradients, model)
dx=-1(alpha*torch.eye(Hessian.shape[-1]).cuda()+Hessian).inverse().mm(g_vector).detach()
cnt=0
model.zero_grad()
for p in model.parameters():
mm=torch.Tensor([p.shape]).tolist()[0]
num=int(functools.reduce(lambda x,y:x*y,mm,1))
p.requires_grad=False
p+=dx[cnt:cnt+num,:].reshape(p.shape)
cnt+=num
p.requires_grad=True

out=model(toy_input).unsqueeze(1)
loss_out=loss(out)
if loss_out<prev_loss:
print("Successful iteration")
loss_hist.append(loss_out)
alpha/=10
else:
print("Augmenting step size")
alpha*=10
cnt=0
for p in model.parameters():
mm=torch.Tensor([p.shape]).tolist()[0]
num=int(functools.reduce(lambda x,y:x*y,mm,1))
p.requires_grad=False
p-=dx[cnt:cnt+num,:].reshape(p.shape)
cnt+=num
p.requires_grad=True
return loss_hist 

def eval_hessian(loss_grad, model):
cnt = 0
for g in loss_grad:
g_vector = g.contiguous().view(-1) if cnt == 0 else torch.cat([g_vector,     g.contiguous().view(-1)])
cnt = 1
l = g_vector.size(0)
hessian = torch.zeros(l, l).cuda()
for idx in range(l):
grad2rd = torch.autograd.grad(g_vector[idx], model.parameters(), create_graph=True)
cnt = 0
for g in grad2rd:
g2 = g.contiguous().view(-1) if cnt == 0 else torch.cat([g2, g.contiguous().view(-1)])
cnt = 1
hessian[idx] = g2
return hessian, g_vector.unsqueeze(1)
def toy_loss(vec):
return vec.transpose(0,1).mm(vec)
class toy_model(torch.nn.Module):
def __init__(self,in_c,width,height):
super().__init__()
self.cnv=torch.nn.Conv2d(in_c,1,3,1,padding=1)
self.lin=torch.nn.Linear(1*width*height,16)
def forward(self,tns):
out=self.cnv(tns)
out=self.lin(out.view(-1))
return out
if __name__=="__main__":
H=20
W=20
toy_input=torch.rand(1,3,H,W).cuda()
toy_mdl=toy_model(3,W,H)
toy_mdl.cuda()
loss_hist=LM(toy_mdl,lambda x:toy_loss(x))

请注意，我从这里获取了eval_hessian的代码。