支持GPU的Levenberg-Marquardt算法

Question

对于浅层神经网络，LM算法取得了惊人的效果。

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

Answer 1

我不确定这样的实现对于任何东西都有用，除了最简单的网络之外，计算Hessian是非常不切实际的。然而，我认为在浅层网络中实现Levenberg-Marquardt并不是那么困难。至少，即使您的实现不是最优的，与CPU版本相比，您仍然会得到一些速度上的提高。下面是一个运行在GPU上的非常快速的、非常不完美的/次优的例子。对于大小为3x20x20的1080ti输入，加速比约为3倍(我目前没有免费的GPU来测试更大的输入)。但是，如果您对损失函数有一个先验(例如，如果它是最小二乘的，并且您知道Hessian可以近似为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))

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